Merge pull request #1076 from suketa/add-timestamp-support-to-scalar-function

Add TIMESTAMP support to scalar functions

Author: suketa

ext/duckdb/scalar_function.c

@@ -259,6 +259,25 @@ static void vector_set_value_at(duckdb_vector vector, duckdb_logical_type elemen
             duckdb_vector_assign_string_element_len(vector, index, str_ptr, str_len);
             break;
         }
+        case DUCKDB_TYPE_TIMESTAMP: {
+            // Convert Ruby Time to DuckDB timestamp (microseconds since epoch)
+            if (!rb_obj_is_kind_of(value, rb_cTime)) {
+                rb_raise(rb_eTypeError, "Expected Time object for TIMESTAMP");
+            }
+            
+            duckdb_timestamp_struct ts_struct;
+            ts_struct.date.year = NUM2INT(rb_funcall(value, rb_intern("year"), 0));
+            ts_struct.date.month = NUM2INT(rb_funcall(value, rb_intern("month"), 0));
+            ts_struct.date.day = NUM2INT(rb_funcall(value, rb_intern("day"), 0));
+            ts_struct.time.hour = NUM2INT(rb_funcall(value, rb_intern("hour"), 0));
+            ts_struct.time.min = NUM2INT(rb_funcall(value, rb_intern("min"), 0));
+            ts_struct.time.sec = NUM2INT(rb_funcall(value, rb_intern("sec"), 0));
+            ts_struct.time.micros = NUM2INT(rb_funcall(value, rb_intern("usec"), 0));
+            
+            duckdb_timestamp ts = duckdb_to_timestamp(ts_struct);
+            ((duckdb_timestamp *)vector_data)[index] = ts;
+            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 BOOLEAN, INTEGER, BIGINT, FLOAT, DOUBLE, VARCHAR, and BLOB types.
+    # Currently supports BIGINT, BLOB, BOOLEAN, DOUBLE, FLOAT, INTEGER, TIMESTAMP, and VARCHAR types.
     #
     # @param logical_type [DuckDB::LogicalType] the return type
     # @return [DuckDB::ScalarFunction] self
@@ -13,9 +13,10 @@ 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 varchar].include?(logical_type.type)
+      unless %i[bigint blob boolean double float integer timestamp varchar].include?(logical_type.type)
         raise DuckDB::Error,
-              'Only BIGINT, BLOB, BOOLEAN, DOUBLE, FLOAT, INTEGER, and VARCHAR return types are currently supported'
+              'Only BIGINT, BLOB, BOOLEAN, DOUBLE, FLOAT, INTEGER, TIMESTAMP, and VARCHAR return types are ' \
+              'currently supported'
       end
 
       _set_return_type(logical_type)

test/duckdb_test/scalar_function_test.rb

@@ -251,5 +251,25 @@ def test_scalar_function_blob_return_type # rubocop:disable Metrics/AbcSize, Met
       assert_equal "\xFF\x00\x01\x02\x03".b, rows[0][0]
       assert_equal "\xFF\x00\xAA\xBB\xCC".b, rows[1][0]
     end
+
+    def test_scalar_function_timestamp_return_type # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
+      @con.execute('SET threads=1')
+      @con.execute('CREATE TABLE test_table (ts TIMESTAMP)')
+      @con.execute("INSERT INTO test_table VALUES ('2024-01-15 10:30:00'), ('2024-12-25 23:59:59')")
+
+      sf = DuckDB::ScalarFunction.new
+      sf.name = 'add_one_hour'
+      sf.add_parameter(DuckDB::LogicalType.new(12)) # TIMESTAMP (type ID 12)
+      sf.return_type = DuckDB::LogicalType.new(12) # TIMESTAMP
+      sf.set_function { |ts| ts + 3600 } # Add 1 hour (3600 seconds)
+
+      @con.register_scalar_function(sf)
+      result = @con.execute('SELECT add_one_hour(ts) FROM test_table ORDER BY ts')
+      rows = result.to_a
+
+      assert_equal 2, rows.size
+      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
   end
 end