Wrap return_fixed_size_list in Ruby

To allow setting true/false/nil.

Author: stenlarsson

ruby/red-arrow/lib/arrow/list-slice-options.rb

@@ -17,6 +17,37 @@
 
 module Arrow
   class ListSliceOptions
+    alias_method :return_fixed_size_list_raw, :return_fixed_size_list
+    private :return_fixed_size_list_raw
+
+    RETURN_FIXED_SIZE_GLIB_TO_RUBY = {
+      ListSliceReturnFixedSizeList::AUTO => nil,
+      ListSliceReturnFixedSizeList::TRUE => true,
+      ListSliceReturnFixedSizeList::FALSE => false,
+    }.freeze
+
+    RETURN_FIXED_SIZE_RUBY_TO_GLIB = RETURN_FIXED_SIZE_GLIB_TO_RUBY.invert.freeze
+
+    # Whether to return a FixedSizeListArray. If true _and_ stop is after a
+    # list element’s length, nil values will be appended to create the requested
+    # slice size. The default of nil will return the same type which was passed in.
+    def return_fixed_size_list
+      RETURN_FIXED_SIZE_GLIB_TO_RUBY.fetch(
+        return_fixed_size_list_raw,
+        return_fixed_size_list_raw)
+    end
+
+    # Whether to return a FixedSizeListArray. If true _and_ stop is after a
+    # list element’s length, nil values will be appended to create the requested
+    # slice size. The default of nil will return the same type which was passed in.
+    def return_fixed_size_list=(return_fixed_size_list)
+      set_property(
+        :return_fixed_size_list,
+        RETURN_FIXED_SIZE_RUBY_TO_GLIB.fetch(
+          return_fixed_size_list,
+          return_fixed_size_list))
+    end
+
     alias_method :stop_raw, :stop
     private :stop_raw
 

ruby/red-arrow/test/test-list-slice-options.rb

@@ -28,6 +28,16 @@ def test_stop_property
     assert_equal(nil, @options.stop)
   end
 
+  def test_return_fixed_size_list_property
+    assert_equal(nil, @options.return_fixed_size_list)
+    @options.return_fixed_size_list = true
+    assert_equal(true, @options.return_fixed_size_list)
+    @options.return_fixed_size_list = false
+    assert_equal(false, @options.return_fixed_size_list)
+    @options.return_fixed_size_list = nil
+    assert_equal(nil, @options.return_fixed_size_list)
+  end
+
   def test_list_slice_function_without_stop
     args = [
       Arrow::ArrayDatum.new(Arrow::ListArray.new([:list, :int8], [[1, 2, 3], [4, 5]])),
@@ -39,4 +49,42 @@ def test_list_slice_function_without_stop
     assert_equal(Arrow::ListArray.new([:list, :int8], [[2, 3], [5]]),
                  result)
   end
+
+  def test_list_slice_function_with_return_fixed_size_list_auto
+    args = [
+      Arrow::ArrayDatum.new(Arrow::FixedSizeListArray.new([:fixed_size_list, :int8, 2],
+                                                          [[1, 2], [3, 4]])),
+    ]
+    @options.start = 1
+    list_slice_function = Arrow::Function.find("list_slice")
+    result = list_slice_function.execute(args, @options).value
+    assert_equal(Arrow::FixedSizeListArray.new([:fixed_size_list, :int8, 1], [[2], [4]]),
+                 result)
+  end
+
+  def test_list_slice_function_with_return_fixed_size_list_true
+    args = [
+      Arrow::ArrayDatum.new(Arrow::ListArray.new([:list, :int8], [[1, 2, 3], [4, 5]])),
+    ]
+    @options.start = 1
+    @options.stop = 3
+    @options.return_fixed_size_list = true
+    list_slice_function = Arrow::Function.find("list_slice")
+    result = list_slice_function.execute(args, @options).value
+    assert_equal(Arrow::FixedSizeListArray.new([:fixed_size_list, :int8, 2], [[2, 3], [5, nil]]),
+                 result)
+  end
+
+  def test_list_slice_function_with_return_fixed_size_list_false
+    args = [
+      Arrow::ArrayDatum.new(Arrow::FixedSizeListArray.new([:fixed_size_list, :int8, 2],
+                                                          [[1, 2], [3, 4]])),
+    ]
+    @options.start = 1
+    @options.return_fixed_size_list = false
+    list_slice_function = Arrow::Function.find("list_slice")
+    result = list_slice_function.execute(args, @options).value
+    assert_equal(Arrow::ListArray.new([:list, :int8], [[2], [4]]),
+                 result)
+  end
 end