Merge pull request #1073 from suketa/add-float-support-to-scalar-function

Add FLOAT support to ScalarFunction return type

Author: suketa

CHANGELOG.md

@@ -3,6 +3,7 @@
 All notable changes to this project will be documented in this file.
 
 # Unreleased
+- add FLOAT support to DuckDB::ScalarFunction return type.
 - add BOOLEAN support to DuckDB::ScalarFunction return type.
 - add DOUBLE support to DuckDB::ScalarFunction return type.
 - add BIGINT support to DuckDB::ScalarFunction return type.

ext/duckdb/scalar_function.c

@@ -237,6 +237,9 @@ static void vector_set_value_at(duckdb_vector vector, duckdb_logical_type elemen
         case DUCKDB_TYPE_BIGINT:
             ((int64_t *)vector_data)[index] = NUM2LL(value);
             break;
+        case DUCKDB_TYPE_FLOAT:
+            ((float *)vector_data)[index] = (float)NUM2DBL(value);
+            break;
         case DUCKDB_TYPE_DOUBLE:
             ((double *)vector_data)[index] = NUM2DBL(value);
             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 INTEGER, BIGINT, DOUBLE, and BOOLEAN types.
+    # Currently supports BOOLEAN, INTEGER, BIGINT, FLOAT, and DOUBLE types.
     #
     # @param logical_type [DuckDB::LogicalType] the return type
     # @return [DuckDB::ScalarFunction] self
@@ -13,8 +13,8 @@ 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[boolean integer bigint double].include?(logical_type.type)
-        raise DuckDB::Error, 'Only BOOLEAN, INTEGER, BIGINT, and DOUBLE return types are currently supported'
+      unless %i[boolean integer bigint float double].include?(logical_type.type)
+        raise DuckDB::Error, 'Only BOOLEAN, INTEGER, BIGINT, FLOAT, and DOUBLE return types are currently supported'
       end
 
       _set_return_type(logical_type)

test/duckdb_test/scalar_function_test.rb

@@ -197,5 +197,22 @@ def test_scalar_function_boolean_return_type # rubocop:disable Metrics/MethodLen
 
       assert_equal [[false], [false], [true]], result.to_a
     end
+
+    def test_scalar_function_float_return_type # rubocop:disable Metrics/MethodLength
+      @con.execute('SET threads=1')
+      @con.execute('CREATE TABLE test_table (value FLOAT)')
+      @con.execute('INSERT INTO test_table VALUES (2.5)')
+
+      sf = DuckDB::ScalarFunction.new
+      sf.name = 'add_half'
+      sf.add_parameter(DuckDB::LogicalType.new(10)) # FLOAT (type ID 10)
+      sf.return_type = DuckDB::LogicalType.new(10) # FLOAT
+      sf.set_function { |v| v + 0.5 }
+
+      @con.register_scalar_function(sf)
+      result = @con.execute('SELECT add_half(value) FROM test_table')
+
+      assert_in_delta 3.0, result.first.first, 0.0001
+    end
   end
 end