ARROW-13910: [Ruby] Arrow::Table#[]/Arrow::RecordBatch#[] accepts Ran…

…ge and selectors

Closes #11090 from kou/ruby-table-array-ref

Authored-by: Sutou Kouhei <kou@clear-code.com>
Signed-off-by: Sutou Kouhei <kou@clear-code.com>

ruby/red-arrow/lib/arrow/column-containable.rb

@@ -27,6 +27,17 @@ def each_column(&block)
       columns.each(&block)
     end
 
+    # @overload [](name)
+    #   Find a column that has the given name.
+    #
+    #   @param name [String, Symbol] The column name to be found.
+    #   @return [Column] The found column.
+    #
+    # @overload [](index)
+    #   Find the `index`-th column.
+    #
+    #   @param index [Integer] The index to be found.
+    #   @return [Column] The found column.
     def find_column(name_or_index)
       case name_or_index
       when String, Symbol
@@ -40,9 +51,97 @@ def find_column(name_or_index)
         return nil if index < 0 or index >= n_columns
         Column.new(self, index)
       else
-        message = "column name or index must be String, Symbol or Integer"
+        message = "column name or index must be String, Symbol or Integer: "
+        message << name_or_index.inspect
         raise ArgumentError, message
       end
     end
+
+    # Selects columns that are selected by `selectors` and/or `block`
+    # and creates a new container only with the selected columns.
+    #
+    # @param selectors [Array<String, Symbol, Integer, Range>]
+    #   If a selector is `String`, `Symbol` or `Integer`, the selector
+    #   selects a column by {#find_column}.
+    #
+    #   If a selector is `Range`, the selector selects columns by `::Array#[]`.
+    # @yield [column] Gives a column to the block to select columns.
+    #   This uses `::Array#select`.
+    # @yieldparam column [Column] A target column.
+    # @yieldreturn [Boolean] Whether the given column is selected or not.
+    # @return [self.class] The newly created container that only has selected
+    #   columns.
+    def select_columns(*selectors, &block)
+      if selectors.empty?
+        return to_enum(__method__) unless block_given?
+        selected_columns = columns.select(&block)
+      else
+        selected_columns = []
+        selectors.each do |selector|
+          case selector
+          when Range
+            selected_columns.concat(columns[selector])
+          else
+            column = find_column(selector)
+            if column.nil?
+              case selector
+              when String, Symbol
+                message = "unknown column: #{selector.inspect}: #{inspect}"
+                raise KeyError.new(message)
+              else
+                message = "out of index (0..#{n_columns - 1}): "
+                message << "#{selector.inspect}: #{inspect}"
+                raise IndexError.new(message)
+              end
+            end
+            selected_columns << column
+          end
+        end
+        selected_columns = selected_columns.select(&block) if block_given?
+      end
+      self.class.new(selected_columns)
+    end
+
+    # @overload [](name)
+    #   Find a column that has the given name.
+    #
+    #   @param name [String, Symbol] The column name to be found.
+    #   @return [Column] The found column.
+    #   @see #find_column
+    #
+    # @overload [](index)
+    #   Find the `index`-th column.
+    #
+    #   @param index [Integer] The index to be found.
+    #   @return [Column] The found column.
+    #   @see #find_column
+    #
+    # @overload [](range)
+    #   Selects columns that are in `range` and creates a new container
+    #   only with the selected columns.
+    #
+    #   @param range [Range] The range to be selected.
+    #   @return [self.class] The newly created container that only has selected
+    #     columns.
+    #   @see #select_columns
+    #
+    # @overload [](selectors)
+    #   Selects columns that are selected by `selectors` and creates a
+    #   new container only with the selected columns.
+    #
+    #   @param selectors [Array] The selectors that are used to select columns.
+    #   @return [self.class] The newly created container that only has selected
+    #     columns.
+    #   @see #select_columns
+    def [](selector)
+      case selector
+      when ::Array
+        select_columns(*selector)
+      when Range
+        select_columns(selector)
+      else
+        find_column(selector)
+      end
+    end
   end
 end

ruby/red-arrow/lib/arrow/map-data-type.rb

@@ -40,7 +40,7 @@ class MapDataType
     #     See {Arrow::DataType.resolve} how to specify data type
     #     description.
     #
-    #   @example Create a map data type for {0: "Hello", 1: "World"}
+    #   @example Create a map data type for `{0: "Hello", 1: "World"}`
     #     key = :int8
     #     item = :string
     #     Arrow::MapDataType.new(key, item)
@@ -66,7 +66,7 @@ class MapDataType
     #     See {Arrow::DataType.resolve} how to specify data type
     #     description.
     #
-    #   @example Create a maap data type for {0: "Hello", 1: "World"}
+    #   @example Create a map data type for `{0: "Hello", 1: "World"}`
     #     Arrow::MapDataType.new(key: :int8, item: :string)
     def initialize(*args)
       n_args = args.size

ruby/red-arrow/lib/arrow/record-batch.rb

@@ -50,8 +50,6 @@ def new(*args)
     alias_method :size, :n_rows
     alias_method :length, :n_rows
 
-    alias_method :[], :find_column
-
     # Converts the record batch to {Arrow::Table}.
     #
     # @return [Arrow::Table]

ruby/red-arrow/lib/arrow/table.rb

@@ -195,8 +195,6 @@ def each_record_batch
     alias_method :size, :n_rows
     alias_method :length, :n_rows
 
-    alias_method :[], :find_column
-
     alias_method :slice_raw, :slice
 
     # @overload slice(offset, length)
@@ -397,41 +395,6 @@ def remove_column(name_or_index)
       remove_column_raw(index)
     end
 
-    # TODO
-    #
-    # @return [Arrow::Table]
-    def select_columns(*selectors, &block)
-      if selectors.empty?
-        return to_enum(__method__) unless block_given?
-        selected_columns = columns.select(&block)
-      else
-        selected_columns = []
-        selectors.each do |selector|
-          case selector
-          when String, Symbol
-            column = find_column(selector)
-            if column.nil?
-              message = "unknown column: #{selector.inspect}: #{inspect}"
-              raise KeyError.new(message)
-            end
-            selected_columns << column
-          when Range
-            selected_columns.concat(columns[selector])
-          else
-            column = columns[selector]
-            if column.nil?
-              message = "out of index (0..#{n_columns - 1}): " +
-              "#{selector.inspect}: #{inspect}"
-              raise IndexError.new(message)
-            end
-            selected_columns << column
-          end
-        end
-        selected_columns = selected_columns.select(&block) if block_given?
-      end
-      self.class.new(selected_columns)
-    end
-
     # Experimental
     def group(*keys)
       Group.new(self, keys)

ruby/red-arrow/test/test-record-batch.rb

@@ -136,5 +136,47 @@ def setup
         end
       end
     end
+
+    sub_test_case("#[]") do
+      def setup
+        @record_batch = Arrow::RecordBatch.new(a: [true],
+                                               b: [true],
+                                               c: [true],
+                                               d: [true],
+                                               e: [true],
+                                               f: [true],
+                                               g: [true])
+      end
+
+      test("[String]") do
+        assert_equal(Arrow::Column.new(@record_batch, 0),
+                     @record_batch["a"])
+      end
+
+      test("[Symbol]") do
+        assert_equal(Arrow::Column.new(@record_batch, 1),
+                     @record_batch[:b])
+      end
+
+      test("[Integer]") do
+        assert_equal(Arrow::Column.new(@record_batch, 6),
+                     @record_batch[-1])
+      end
+
+      test("[Range]") do
+        assert_equal(Arrow::RecordBatch.new(d: [true],
+                                            e: [true]),
+                     @record_batch[3..4])
+      end
+
+      test("[[Symbol, String, Integer, Range]]") do
+        assert_equal(Arrow::RecordBatch.new(c: [true],
+                                            a: [true],
+                                            g: [true],
+                                            d: [true],
+                                            e: [true]),
+                     @record_batch[[:c, "a", -1, 3..4]])
+      end
+    end
   end
 end

ruby/red-arrow/test/test-table.rb

@@ -190,20 +190,45 @@ def setup
   end
 
   sub_test_case("#[]") do
+    def setup
+      @table = Arrow::Table.new(a: [true],
+                                b: [true],
+                                c: [true],
+                                d: [true],
+                                e: [true],
+                                f: [true],
+                                g: [true])
+    end
+
     test("[String]") do
       assert_equal(Arrow::Column.new(@table, 0),
-                   @table["count"])
+                   @table["a"])
     end
 
     test("[Symbol]") do
       assert_equal(Arrow::Column.new(@table, 1),
-                   @table[:visible])
+                   @table[:b])
     end
 
     test("[Integer]") do
-      assert_equal(Arrow::Column.new(@table, 1),
+      assert_equal(Arrow::Column.new(@table, 6),
                    @table[-1])
     end
+
+    test("[Range]") do
+      assert_equal(Arrow::Table.new(d: [true],
+                                    e: [true]),
+                   @table[3..4])
+    end
+
+    test("[[Symbol, String, Integer, Range]]") do
+      assert_equal(Arrow::Table.new(c: [true],
+                                    a: [true],
+                                    g: [true],
+                                    d: [true],
+                                    e: [true]),
+                   @table[[:c, "a", -1, 3..4]])
+    end
   end
 
   sub_test_case("#merge") do