Merge pull request #1124 from suketa/memory_helper_support_timestamp

Support TIMESTAMP column writing Ruby Time object in DataChunk#set_value

Author: suketa

CHANGELOG.md

@@ -3,6 +3,8 @@
 All notable changes to this project will be documented in this file.
 
 # Unreleased
+- support TIMESTAMP column writing Ruby Time object in `DuckDB::DataChunk#set_value`.
+- add `DuckDB::MemoryHelper.write_timestamp` method to write a Ruby Time object as a DuckDB timestamp to raw memory.
 - add `DuckDB::Connection#expose_as_table` to expose a Ruby object as a queryable DuckDB table function via a registered adapter.
 - add `DuckDB::TableFunction.add_table_adapter` to register a table adapter for a Ruby class.
 - add `DuckDB::TableFunction.table_adapter_for` to look up a registered table adapter by class.

ext/duckdb/memory_helper.c

@@ -234,6 +234,43 @@ static VALUE rbduckdb_memory_helper_write_float(VALUE self, VALUE ptr, VALUE ind
     return Qnil;
 }
 
+/*
+ * call-seq:
+ *   DuckDB::MemoryHelper.write_timestamp(ptr, index, value) -> nil
+ *
+ * Writes a DuckDB timestamp to raw memory.
+ * +value+ must be a Ruby Time object.
+ *
+ *   ptr = vector.get_data
+ *   DuckDB::MemoryHelper.write_timestamp(ptr, 0, Time.new(2024, 3, 15, 10, 30, 45))
+ */
+static VALUE rbduckdb_memory_helper_write_timestamp(VALUE self, VALUE ptr, VALUE index, VALUE value) {
+    duckdb_timestamp *data;
+    idx_t idx;
+
+    (void)self;
+
+    if (!rb_obj_is_kind_of(value, rb_cTime)) {
+        rb_raise(rb_eTypeError, "Expected Time object for TIMESTAMP");
+    }
+
+    data = (duckdb_timestamp *)NUM2ULL(ptr);
+    idx = (idx_t)NUM2ULL(index);
+
+    VALUE local_time = rb_funcall(value, rb_intern("getlocal"), 0);
+    data[idx] = rbduckdb_to_duckdb_timestamp_from_value(
+        rb_funcall(local_time, rb_intern("year"), 0),
+        rb_funcall(local_time, rb_intern("month"), 0),
+        rb_funcall(local_time, rb_intern("day"), 0),
+        rb_funcall(local_time, rb_intern("hour"), 0),
+        rb_funcall(local_time, rb_intern("min"), 0),
+        rb_funcall(local_time, rb_intern("sec"), 0),
+        rb_funcall(local_time, rb_intern("usec"), 0)
+    );
+
+    return Qnil;
+}
+
 void rbduckdb_init_memory_helper(void) {
     mDuckDBMemoryHelper = rb_define_module_under(mDuckDB, "MemoryHelper");
 
@@ -248,4 +285,5 @@ void rbduckdb_init_memory_helper(void) {
     rb_define_singleton_method(mDuckDBMemoryHelper, "write_uinteger", rbduckdb_memory_helper_write_uinteger, 3);
     rb_define_singleton_method(mDuckDBMemoryHelper, "write_ubigint", rbduckdb_memory_helper_write_ubigint, 3);
     rb_define_singleton_method(mDuckDBMemoryHelper, "write_float", rbduckdb_memory_helper_write_float, 3);
+    rb_define_singleton_method(mDuckDBMemoryHelper, "write_timestamp", rbduckdb_memory_helper_write_timestamp, 3);
 }

lib/duckdb/data_chunk.rb

@@ -94,6 +94,9 @@ def set_value(col_idx, row_idx, value)
         vector.assign_string_element(row_idx, value.to_s)
       when :blob
         vector.assign_string_element_len(row_idx, value.to_s)
+      when :timestamp
+        data = vector.get_data
+        MemoryHelper.write_timestamp(data, row_idx, value)
       else
         raise ArgumentError, "Unsupported type for DataChunk#set_value: #{type_id}"
       end

sample/issue922.rb

@@ -35,7 +35,7 @@ def infer_columns(data_frame)
 
 DuckDB::TableFunction.add_table_adapter(Polars::DataFrame, PolarsDataFrameTableAdapter.new)
 
-df = polars::dataframe.new(
+df = Polars::DataFrame.new(
   {
     a: [1, 2, 3],
     b: %w[one two three]

test/duckdb_test/data_chunk_test.rb

@@ -358,5 +358,42 @@ def test_data_chunk_set_value_multiple_columns
       assert_in_delta 87.3, row1[2], 0.001
     end
     # rubocop:enable Metrics/AbcSize, Metrics/MethodLength, Minitest/MultipleAssertions
+
+    # Test 12: DataChunk#set_value with TIMESTAMP
+    def test_data_chunk_set_value_timestamp # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
+      @conn.execute('SET threads=1')
+
+      done = false
+      table_function = DuckDB::TableFunction.new
+      table_function.name = 'test_set_value_timestamp'
+
+      table_function.bind do |bind_info|
+        bind_info.add_result_column('ts', DuckDB::LogicalType::TIMESTAMP)
+      end
+
+      table_function.init { |_init_info| done = false }
+
+      time1 = Time.new(2024, 3, 15, 10, 30, 45, '+00:00')
+      time2 = Time.new(2000, 1, 1, 0, 0, 0, '+00:00')
+
+      table_function.execute do |_func_info, output|
+        if done
+          output.size = 0
+        else
+          output.set_value(0, 0, time1)
+          output.set_value(0, 1, time2)
+          output.size = 2
+          done = true
+        end
+      end
+
+      @conn.register_table_function(table_function)
+      result = @conn.query('SELECT * FROM test_set_value_timestamp()')
+      rows = result.each.to_a
+
+      assert_equal 2, rows.length
+      assert_equal time1.utc, rows[0].first.utc
+      assert_equal time2.utc, rows[1].first.utc
+    end
   end
 end