red-arrow 4.0.0 → 6.0.1

data/lib/arrow/group.rb

@@ -16,157 +16,149 @@
 # under the License.
 
 module Arrow
-  # Experimental
-  #
-  # TODO: Almost codes should be implemented in Apache Arrow C++.
   class Group
     def initialize(table, keys)
       @table = table
       @keys = keys
     end
 
-    def count
-      key_names = @keys.collect(&:to_s)
-      target_columns = @table.columns.reject do |column|
-        key_names.include?(column.name)
-      end
-      aggregate(target_columns) do |column, indexes|
-        n = 0
-        indexes.each do |index|
-          n += 1 unless column.null?(index)
-        end
-        n
-      end
+    def count(*target_names)
+      aggregate(*build_aggregations("hash_count", target_names))
     end
 
-    def sum
-      key_names = @keys.collect(&:to_s)
-      target_columns = @table.columns.reject do |column|
-        key_names.include?(column.name) or
-          not column.data_type.is_a?(NumericDataType)
-      end
-      aggregate(target_columns) do |column, indexes|
-        n = 0
-        indexes.each do |index|
-          value = column[index]
-          n += value unless value.nil?
-        end
-        n
-      end
+    def sum(*target_names)
+      aggregate(*build_aggregations("hash_sum", target_names))
     end
 
-    def average
-      key_names = @keys.collect(&:to_s)
-      target_columns = @table.columns.reject do |column|
-        key_names.include?(column.name) or
-          not column.data_type.is_a?(NumericDataType)
-      end
-      aggregate(target_columns) do |column, indexes|
-        average = 0.0
-        n = 0
-        indexes.each do |index|
-          value = column[index]
-          unless value.nil?
-            n += 1
-            average += (value - average) / n
-          end
-        end
-        average
-      end
+    def product(*target_names)
+      aggregate(*build_aggregations("hash_product", target_names))
     end
 
-    def min
-      key_names = @keys.collect(&:to_s)
-      target_columns = @table.columns.reject do |column|
-        key_names.include?(column.name) or
-          not column.data_type.is_a?(NumericDataType)
-      end
-      aggregate(target_columns) do |column, indexes|
-        n = nil
-        indexes.each do |index|
-          value = column[index]
-          next if value.nil?
-          n ||= value
-          n = value if value < n
-        end
-        n
-      end
+    def mean(*target_names)
+      aggregate(*build_aggregations("hash_mean", target_names))
     end
 
-    def max
-      key_names = @keys.collect(&:to_s)
-      target_columns = @table.columns.reject do |column|
-        key_names.include?(column.name) or
-          not column.data_type.is_a?(NumericDataType)
-      end
-      aggregate(target_columns) do |column, indexes|
-        n = nil
-        indexes.each do |index|
-          value = column[index]
-          next if value.nil?
-          n ||= value
-          n = value if value > n
-        end
-        n
-      end
+    def min(*target_names)
+      aggregate(*build_aggregations("hash_min", target_names))
+    end
+
+    def max(*target_names)
+      aggregate(*build_aggregations("hash_max", target_names))
+    end
+
+    def stddev(*target_names)
+      aggregate(*build_aggregations("hash_stddev", target_names))
+    end
+
+    def variance(*target_names)
+      aggregate(*build_aggregations("hash_variance", target_names))
+    end
+
+    def aggregate(aggregation, *more_aggregations)
+      aggregations = [aggregation] + more_aggregations
+      normalized_aggregations = normalize_aggregations(aggregations)
+      plan = ExecutePlan.new
+      source_node = plan.build_source_node(@table)
+      aggregate_node =
+        plan.build_aggregate_node(source_node,
+                                  {
+                                    aggregations: normalized_aggregations,
+                                    keys: @keys
+                                  })
+      sink_node_options = SinkNodeOptions.new
+      plan.build_sink_node(aggregate_node, sink_node_options)
+      plan.validate
+      plan.start
+      plan.wait
+      reader = sink_node_options.get_reader(aggregate_node.output_schema)
+      reader.read_all
     end
 
     private
-    def aggregate(target_columns)
-      sort_values = @table.n_rows.times.collect do |i|
-        key_values = @keys.collect do |key|
-          @table[key][i]
+    def build_aggregations(function_name, target_names)
+      if target_names.empty?
+        [function_name]
+      else
+        target_names.collect do |name|
+          "#{function_name}(#{name})"
         end
-        [key_values, i]
-      end
-      sorted = sort_values.sort_by do |key_values, i|
-        key_values
       end
+    end
 
-      grouped_keys = []
-      aggregated_arrays_raw = []
-      target_columns.size.times do
-        aggregated_arrays_raw << []
-      end
-      indexes = []
-      sorted.each do |key_values, i|
-        if grouped_keys.empty?
-          grouped_keys << key_values
-          indexes.clear
-          indexes << i
-        else
-          if key_values == grouped_keys.last
-            indexes << i
-          else
-            grouped_keys << key_values
-            target_columns.each_with_index do |column, j|
-              aggregated_arrays_raw[j] << yield(column, indexes)
-            end
-            indexes.clear
-            indexes << i
+    def normalize_aggregations(aggregations)
+      normalized_aggregations = []
+      aggregations.each do |aggregation|
+        case aggregation
+        when :all
+          all_functions = [
+            "hash_count",
+            "hash_sum",
+            "hash_product",
+            "hash_mean",
+            "hash_stddev",
+            "hash_variance",
+            # "hash_tdigest",
+            "hash_min",
+            "hash_max",
+            "hash_any",
+            "hash_all",
+          ]
+          normalized_aggregations.concat(normalize_aggregations(all_functions))
+        when /\A([a-zA-Z0-9_].+?)\((.+?)\)\z/
+          function = $1
+          input = $2.strip
+          normalized_aggregations << {function: function, input: input}
+        when "count", "hash_count"
+          function = aggregation
+          target_columns.each do |column|
+            normalized_aggregations << {function: function, input: column.name}
           end
+        when "any", "hash_any", "all", "hash_all"
+          function = aggregation
+          boolean_target_columns.each do |column|
+            normalized_aggregations << {function: function, input: column.name}
+          end
+        when String
+          function = aggregation
+          numeric_target_columns.each do |column|
+            normalized_aggregations << {function: function, input: column.name}
+          end
+        else
+          normalized_aggregations << aggregation
         end
       end
-      target_columns.each_with_index do |column, j|
-        aggregated_arrays_raw[j] << yield(column, indexes)
+      normalized_aggregations
+    end
+
+    def target_columns
+      @target_columns ||= find_target_columns
+    end
+
+    def find_target_columns
+      key_names = @keys.collect(&:to_s)
+      @table.columns.find_all do |column|
+        not key_names.include?(column.name)
       end
+    end
 
-      grouped_key_arrays_raw = grouped_keys.transpose
-      fields = []
-      arrays = []
-      @keys.each_with_index do |key, i|
-        key_column = @table[key]
-        key_column_array_raw = grouped_key_arrays_raw[i]
-        key_column_array = key_column.data_type.build_array(key_column_array_raw)
-        fields << key_column.field
-        arrays << key_column_array
+    def boolean_target_columns
+      @boolean_target_columns ||= find_boolean_target_columns
+    end
+
+    def find_boolean_target_columns
+      target_columns.find_all do |column|
+        column.data_type.is_a?(BooleanDataType)
       end
-      target_columns.each_with_index do |column, i|
-        array = ArrayBuilder.build(aggregated_arrays_raw[i])
-        arrays << array
-        fields << Field.new(column.field.name, array.value_data_type)
+    end
+
+    def numeric_target_columns
+      @numeric_target_columns ||= find_numeric_target_columns
+    end
+
+    def find_numeric_target_columns
+      target_columns.find_all do |column|
+        column.data_type.is_a?(NumericDataType)
       end
-      Table.new(fields, arrays)
     end
   end
 end

data/lib/arrow/loader.rb

@@ -29,6 +29,7 @@ module Arrow
     def post_load(repository, namespace)
       require_libraries
       require_extension_library
+      gc_guard
     end
 
     def require_libraries
@@ -37,10 +38,14 @@ module Arrow
       require "arrow/generic-filterable"
       require "arrow/generic-takeable"
       require "arrow/record-containable"
+      require "arrow/symbol-values-appendable"
 
+      require "arrow/aggregate-node-options"
+      require "arrow/aggregation"
       require "arrow/array"
       require "arrow/array-builder"
       require "arrow/bigdecimal-extension"
+      require "arrow/binary-dictionary-array-builder"
       require "arrow/buffer"
       require "arrow/chunked-array"
       require "arrow/column"
@@ -52,6 +57,7 @@ module Arrow
       require "arrow/date32-array-builder"
       require "arrow/date64-array"
       require "arrow/date64-array-builder"
+      require "arrow/datum"
       require "arrow/decimal128"
       require "arrow/decimal128-array"
       require "arrow/decimal128-array-builder"
@@ -63,13 +69,19 @@ module Arrow
       require "arrow/dense-union-data-type"
       require "arrow/dictionary-array"
       require "arrow/dictionary-data-type"
+      require "arrow/equal-options"
+      require "arrow/expression"
       require "arrow/field"
       require "arrow/file-output-stream"
+      require "arrow/file-system"
       require "arrow/fixed-size-binary-array"
       require "arrow/fixed-size-binary-array-builder"
       require "arrow/group"
       require "arrow/list-array-builder"
       require "arrow/list-data-type"
+      require "arrow/map-array"
+      require "arrow/map-array-builder"
+      require "arrow/map-data-type"
       require "arrow/null-array"
       require "arrow/null-array-builder"
       require "arrow/path-extension"
@@ -78,17 +90,22 @@ module Arrow
       require "arrow/record-batch-builder"
       require "arrow/record-batch-file-reader"
       require "arrow/record-batch-iterator"
+      require "arrow/record-batch-reader"
       require "arrow/record-batch-stream-reader"
       require "arrow/rolling-window"
+      require "arrow/scalar"
       require "arrow/schema"
       require "arrow/slicer"
       require "arrow/sort-key"
       require "arrow/sort-options"
+      require "arrow/source-node-options"
       require "arrow/sparse-union-data-type"
+      require "arrow/string-dictionary-array-builder"
       require "arrow/struct-array"
       require "arrow/struct-array-builder"
       require "arrow/struct-data-type"
       require "arrow/table"
+      require "arrow/table-concatenate-options"
       require "arrow/table-formatter"
       require "arrow/table-list-formatter"
       require "arrow/table-table-formatter"
@@ -112,6 +129,27 @@ module Arrow
       require "arrow.so"
     end
 
+    def gc_guard
+      require "arrow/constructor-arguments-gc-guardable"
+
+      [
+        @base_module::BinaryScalar,
+        @base_module::Buffer,
+        @base_module::DenseUnionScalar,
+        @base_module::FixedSizeBinaryScalar,
+        @base_module::LargeBinaryScalar,
+        @base_module::LargeListScalar,
+        @base_module::LargeStringScalar,
+        @base_module::ListScalar,
+        @base_module::MapScalar,
+        @base_module::SparseUnionScalar,
+        @base_module::StringScalar,
+        @base_module::StructScalar,
+      ].each do |klass|
+        klass.prepend(ConstructorArgumentsGCGuardable)
+      end
+    end
+
     def load_object_info(info)
       super
 
@@ -164,6 +202,12 @@ module Arrow
           method_name = "dup"
         end
         super(info, klass, method_name)
+      when "Arrow::BooleanScalar"
+        case method_name
+        when "value?"
+          method_name = "value"
+        end
+        super(info, klass, method_name)
       else
         super
       end

data/lib/arrow/map-array-builder.rb

@@ -0,0 +1,109 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+module Arrow
+  class MapArrayBuilder
+    class << self
+      def build(data_type, values)
+        builder = new(data_type)
+        builder.build(values)
+      end
+    end
+
+    alias_method :append_value_raw, :append_value
+
+    # @overload append_value
+    #
+    #   Starts appending a map record. You need to append
+    #   values of map by {#key_builder} and {#item_builder}.
+    #
+    # @overload append_value(value)
+    #
+    #   Appends a map record including key and item values.
+    #
+    #   @param value [nil, #each] The map record.
+    #
+    #     If this is `nil`, the map record is null.
+    #
+    #     If this is an `Object` that has `#each`, each value is a pair of key and item.
+    #
+    # @since 6.0.0
+    def append_value(*args)
+      n_args = args.size
+
+      case n_args
+      when 0
+        append_value_raw
+      when 1
+        value = args[0]
+        case value
+        when nil
+          append_null
+        else
+          unless value.respond_to?(:each)
+            message = "map value must be nil, Hash or Object that has #each: #{value.inspect}"
+            raise ArgumentError, message
+          end
+          append_value_raw
+          @key_builder ||= key_builder
+          @item_builder ||= item_builder
+          case value
+          when Hash
+            keys = value.keys
+            values = value.values
+          else
+            keys = []
+            values = []
+            value.each do |key, item|
+              keys << key
+              values << item
+            end
+          end
+          @key_builder.append(*keys)
+          @item_builder.append(*values)
+        end
+      else
+        message = "wrong number of arguments (given #{n_args}, expected 0..1)"
+        raise ArgumentError, message
+      end
+    end
+
+    alias_method :append_values_raw, :append_values
+
+    def append_values(values, is_valids=nil)
+      value = values[0]
+      case value
+      when Integer
+        append_values_raw(values, is_valids)
+      else
+        if is_valids
+          is_valids.each_with_index do |is_valid, i|
+            if is_valid
+              append_value(values[i])
+            else
+              append_null
+            end
+          end
+        else
+          values.each do |value|
+            append_value(value)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow/map-array.rb

@@ -0,0 +1,26 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+module Arrow
+  class MapArray
+    def get_value(i)
+      super.each_with_object({}) do |item, result|
+        result[item["key"]] = item["value"]
+      end
+    end
+  end
+end

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

@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+module Arrow
+  class MapDataType
+    alias_method :initialize_raw, :initialize
+    private :initialize_raw
+
+    # Creates a new {Arrow::MapDataType}.
+    #
+    # @overload initialize(key, item)
+    #
+    #   @param key [Arrow::DataType, Hash, String, Symbol]
+    #     The key data type of the map data type.
+    #
+    #     You can specify data type as a description by `Hash`.
+    #
+    #     See {Arrow::DataType.resolve} how to specify data type
+    #     description.
+    #
+    #   @param item [Arrow::DataType, Hash, String, Symbol]
+    #     The item data type of the map data type.
+    #
+    #     You can specify data type as a description by `Hash`.
+    #
+    #     See {Arrow::DataType.resolve} how to specify data type
+    #     description.
+    #
+    #   @example Create a map data type for `{0: "Hello", 1: "World"}`
+    #     key = :int8
+    #     item = :string
+    #     Arrow::MapDataType.new(key, item)
+    #
+    # @overload initialize(description)
+    #
+    #   @param description [Hash] The description of the map data
+    #     type. It must have `:key`, `:item` values.
+    #
+    #   @option description [Arrow::DataType, Hash, String, Symbol]
+    #     :key The key data type of the map data type.
+    #
+    #     You can specify data type as a description by `Hash`.
+    #
+    #     See {Arrow::DataType.resolve} how to specify data type
+    #     description.
+    #
+    #   @option description [Arrow::DataType, Hash, String, Symbol]
+    #     :item  The item data type of the map data type.
+    #
+    #     You can specify data type as a description by `Hash`.
+    #
+    #     See {Arrow::DataType.resolve} how to specify data type
+    #     description.
+    #
+    #   @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
+      case n_args
+      when 1
+        description = args[0]
+        key = description[:key]
+        item = description[:item]
+      when 2
+        key, item = args
+      else
+        message = "wrong number of arguments (given, #{n_args}, expected 1..2)"
+        raise ArgumentError, message
+      end
+      key = DataType.resolve(key)
+      item = DataType.resolve(item)
+      initialize_raw(key, item)
+    end
+  end
+end

data/lib/arrow/path-extension.rb

@@ -24,7 +24,7 @@ module Arrow
     def extract
       basename = ::File.basename(@path)
       components = basename.split(".")
-      return {} if components.size == 1
+      return {} if components.size < 2
 
       extension = components.last.downcase
       if components.size > 2

data/lib/arrow/record-batch-reader.rb

@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+module Arrow
+  class RecordBatchReader
+    class << self
+      # @api private
+      def try_convert(value)
+        case value
+        when ::Array
+          return nil if value.empty?
+          if value.all? {|v| v.is_a?(RecordBatch)}
+            new(value)
+          else
+            nil
+          end
+        when RecordBatch
+          new([value])
+        when Table
+          TableBatchReader.new(value)
+        else
+          nil
+        end
+      end
+    end
+  end
+end

data/lib/arrow/record-batch.rb

@@ -50,8 +50,6 @@ module Arrow
     alias_method :size, :n_rows
     alias_method :length, :n_rows
 
-    alias_method :[], :find_column
-
     # Converts the record batch to {Arrow::Table}.
     #
     # @return [Arrow::Table]