Merge pull request #1077 from suketa/add-date-support-to-scalar-function

Add DATE support to scalar functions

Author: suketa

ext/duckdb/scalar_function.c

@@ -278,6 +278,22 @@ static void vector_set_value_at(duckdb_vector vector, duckdb_logical_type elemen
             ((duckdb_timestamp *)vector_data)[index] = ts;
             break;
         }
+        case DUCKDB_TYPE_DATE: {
+            // Convert Ruby Date to DuckDB date
+            // Ruby Date is defined in date library, check with Class name
+            VALUE date_class = rb_const_get(rb_cObject, rb_intern("Date"));
+            if (!rb_obj_is_kind_of(value, date_class)) {
+                rb_raise(rb_eTypeError, "Expected Date object for DATE");
+            }
+            
+            VALUE year = rb_funcall(value, rb_intern("year"), 0);
+            VALUE month = rb_funcall(value, rb_intern("month"), 0);
+            VALUE day = rb_funcall(value, rb_intern("day"), 0);
+            
+            duckdb_date date = rbduckdb_to_duckdb_date_from_value(year, month, day);
+            ((duckdb_date *)vector_data)[index] = date;
+            break;
+        }
         default:
             rb_raise(rb_eArgError, "Unsupported return type for scalar function");
             break;

lib/duckdb/scalar_function.rb

@@ -4,7 +4,7 @@ module DuckDB
   # DuckDB::ScalarFunction encapsulates DuckDB's scalar function
   class ScalarFunction
     # Sets the return type for the scalar function.
-    # Currently supports BIGINT, BLOB, BOOLEAN, DOUBLE, FLOAT, INTEGER, TIMESTAMP, and VARCHAR types.
+    # Currently supports BIGINT, BLOB, BOOLEAN, DATE, DOUBLE, FLOAT, INTEGER, TIMESTAMP, and VARCHAR types.
     #
     # @param logical_type [DuckDB::LogicalType] the return type
     # @return [DuckDB::ScalarFunction] self
@@ -13,9 +13,9 @@ def return_type=(logical_type)
       raise DuckDB::Error, 'logical_type must be a DuckDB::LogicalType' unless logical_type.is_a?(DuckDB::LogicalType)
 
       # Check if the type is supported
-      unless %i[bigint blob boolean double float integer timestamp varchar].include?(logical_type.type)
+      unless %i[bigint blob boolean date double float integer timestamp varchar].include?(logical_type.type)
         raise DuckDB::Error,
-              'Only BIGINT, BLOB, BOOLEAN, DOUBLE, FLOAT, INTEGER, TIMESTAMP, and VARCHAR return types are ' \
+              'Only BIGINT, BLOB, BOOLEAN, DATE, DOUBLE, FLOAT, INTEGER, TIMESTAMP, and VARCHAR return types are ' \
               'currently supported'
       end
 

test/duckdb_test/scalar_function_test.rb

@@ -37,10 +37,10 @@ def test_return_type_setter
 
     def test_return_type_setter_raises_error_for_unsupported_type
       sf = DuckDB::ScalarFunction.new
-      timestamp_type = DuckDB::LogicalType.new(13) # DUCKDB_TYPE_TIMESTAMP (unsupported)
+      time_type = DuckDB::LogicalType.new(14) # DUCKDB_TYPE_TIME (unsupported)
 
       error = assert_raises(DuckDB::Error) do
-        sf.return_type = timestamp_type
+        sf.return_type = time_type
       end
 
       assert_match(/only.*supported/i, error.message)
@@ -271,5 +271,25 @@ def test_scalar_function_timestamp_return_type # rubocop:disable Metrics/AbcSize
       assert_equal Time.new(2024, 1, 15, 11, 30, 0), rows[0][0]
       assert_equal Time.new(2024, 12, 26, 0, 59, 59), rows[1][0]
     end
+
+    def test_scalar_function_date_return_type # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
+      @con.execute('SET threads=1')
+      @con.execute('CREATE TABLE test_table (d DATE)')
+      @con.execute("INSERT INTO test_table VALUES ('2024-01-15'), ('2024-12-25')")
+
+      sf = DuckDB::ScalarFunction.new
+      sf.name = 'add_one_day'
+      sf.add_parameter(DuckDB::LogicalType.new(13)) # DATE (type ID 13)
+      sf.return_type = DuckDB::LogicalType.new(13) # DATE
+      sf.set_function { |date| date + 1 } # Add 1 day
+
+      @con.register_scalar_function(sf)
+      result = @con.execute('SELECT add_one_day(d) FROM test_table ORDER BY d')
+      rows = result.to_a
+
+      assert_equal 2, rows.size
+      assert_equal Date.new(2024, 1, 16), rows[0][0]
+      assert_equal Date.new(2024, 12, 26), rows[1][0]
+    end
   end
 end