red-arrow 3.0.0 → 6.0.0

data/lib/arrow/source-node-options.rb

@@ -0,0 +1,32 @@
+# 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 SourceNodeOptions
+    class << self
+      # @api private
+      def try_convert(value)
+        case value
+        when RecordBatchReader, RecordBatch, Table
+          new(value)
+        else
+          nil
+        end
+      end
+    end
+  end
+end

data/lib/arrow/sparse-union-data-type.rb

@@ -33,7 +33,7 @@ module Arrow
     #   @param type_codes [::Array<Integer>] The IDs that indicates
     #     corresponding fields.
     #
-    #   @example Create a sparse union data type for {2: visible, 9: count}
+    #   @example Create a sparse union data type for `{2: visible, 9: count}`
     #     fields = [
     #       Arrow::Field.new("visible", :boolean),
     #       {
@@ -57,7 +57,7 @@ module Arrow
     #   @option description [::Array<Integer>] :type_codes The IDs
     #     that indicates corresponding fields.
     #
-    #   @example Create a sparse union data type for {2: visible, 9: count}
+    #   @example Create a sparse union data type for `{2: visible, 9: count}`
     #     fields = [
     #       Arrow::Field.new("visible", :boolean),
     #       {

data/lib/arrow/string-dictionary-array-builder.rb

@@ -0,0 +1,27 @@
+# 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 StringDictionaryArrayBuilder
+    include SymbolValuesAppendable
+
+    private
+    def create_values_array_builder
+      StringArrayBuilder.new
+    end
+  end
+end

data/lib/arrow/symbol-values-appendable.rb

@@ -0,0 +1,34 @@
+# 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
+  module SymbolValuesAppendable
+    def append_values(values, is_valids=nil)
+      builder = create_values_array_builder
+      values = values.collect do |value|
+        case value
+        when Symbol
+          value.to_s
+        else
+          value
+        end
+      end
+      builder.append_values(values, is_valids)
+      append_array(builder.finish)
+    end
+  end
+end

data/lib/arrow/table-concatenate-options.rb

@@ -0,0 +1,36 @@
+# 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 TableConcatenateOptions
+    class << self
+      # @api private
+      def try_convert(value)
+        case value
+        when Hash
+          options = new
+          value.each do |k, v|
+            options.public_send("#{k}=", value)
+          end
+          options
+        else
+          nil
+        end
+      end
+    end
+  end
+end

data/lib/arrow/table-formatter.rb

@@ -18,6 +18,125 @@
 module Arrow
   # TODO: Almost codes should be implemented in Apache Arrow C++.
   class TableFormatter
+    # @private
+    class ColumnFormatter
+      attr_reader :column
+      attr_reader :head_values
+      attr_reader :tail_values
+      attr_reader :sample_values
+      def initialize(column, head_values, tail_values)
+        @column = column
+        @head_values = head_values
+        @tail_values = tail_values
+        @sample_values = head_values + tail_values
+        @field_value_widths = {}
+      end
+
+      def data_type
+        @data_type ||= @column.data_type
+      end
+
+      def name
+        @name ||= @column.name
+      end
+
+      def aligned_name
+        @aligned_name ||= format_aligned_name(name, data_type, @sample_values)
+      end
+
+      FLOAT_N_DIGITS = 10
+      FORMATTED_NULL = "(null)"
+
+      def format_value(value, width=0)
+        case value
+        when ::Time
+          value.iso8601
+        when Float
+          "%*f" % [[width, FLOAT_N_DIGITS].max, value]
+        when Integer
+          "%*d" % [width, value]
+        when Hash
+          formatted_values = data_type.fields.collect do |field|
+            field_name = field.name
+            field_value_width = compute_field_value_width(field, @sample_values)
+            formatted_name = format_value(field_name, 0)
+            formatted_value = format_value(value[field_name], field_value_width)
+            "#{formatted_name}: #{formatted_value}"
+          end
+          formatted = "{"
+          formatted << formatted_values.join(", ")
+          formatted << "}"
+          "%-*s" % [width, formatted]
+        when nil
+          "%*s" % [width, FORMATTED_NULL]
+        else
+          "%-*s" % [width, value.to_s]
+        end
+      end
+
+      private
+      def compute_field_value_width(field, sample_values)
+        unless @field_value_widths.key?(field)
+          field_name = field.name
+          field_sample_values = sample_values.collect do |v|
+            (v || {})[field_name]
+          end
+          field_aligned_name = format_aligned_name("",
+                                                   field.data_type,
+                                                   field_sample_values)
+          @field_value_widths[field] = field_aligned_name.size
+        end
+        @field_value_widths[field]
+      end
+
+      def format_aligned_name(name, data_type, sample_values)
+        case data_type
+        when TimestampDataType
+          "%*s" % [::Time.now.iso8601.size, name]
+        when IntegerDataType
+          have_null = false
+          have_negative = false
+          max_value = nil
+          sample_values.each do |value|
+            if value.nil?
+              have_null = true
+            else
+              if max_value.nil?
+                max_value = value.abs
+              else
+                max_value = [value.abs, max_value].max
+              end
+              have_negative = true if value.negative?
+            end
+          end
+          if max_value.nil?
+            width = 0
+          elsif max_value.zero?
+            width = 1
+          else
+            width = (Math.log10(max_value) + 1).truncate
+          end
+          width += 1 if have_negative # Need "-"
+          width = [width, FORMATTED_NULL.size].max if have_null
+          "%*s" % [width, name]
+        when FloatDataType, DoubleDataType
+          "%*s" % [FLOAT_N_DIGITS, name]
+        when StructDataType
+          field_widths = data_type.fields.collect do |field|
+            field_value_width = compute_field_value_width(field, sample_values)
+            field.name.size + ": ".size + field_value_width
+          end
+          width = "{}".size + field_widths.sum
+          if field_widths.size > 0
+            width += (", ".size * (field_widths.size - 1))
+          end
+          "%*s" % [width, name]
+        else
+          name
+        end
+      end
+    end
+
     def initialize(table, options={})
       @table = table
       @options = options
@@ -25,38 +144,43 @@ module Arrow
 
     def format
       text = ""
-      columns = @table.columns
-      format_header(text, columns)
-
       n_rows = @table.n_rows
-      return text if n_rows.zero?
-
       border = @options[:border] || 10
-      n_digits = (Math.log10(n_rows) + 1).truncate
+
       head_limit = [border, n_rows].min
-      head_column_values = columns.collect do |column|
-        column.each.take(head_limit)
+
+      tail_start = [border, n_rows - border].max
+      tail_limit = n_rows - tail_start
+
+      column_formatters = @table.columns.collect do |column|
+        head_values = column.each.take(head_limit)
+        if tail_limit > 0
+          tail_values = column.reverse_each.take(tail_limit).reverse
+        else
+          tail_values = []
+        end
+        ColumnFormatter.new(column, head_values, tail_values)
       end
+
+      format_header(text, column_formatters)
+      return text if n_rows.zero?
+
+      n_digits = (Math.log10(n_rows) + 1).truncate
       format_rows(text,
-                  columns,
-                  head_column_values.transpose,
+                  column_formatters,
+                  column_formatters.collect(&:head_values).transpose,
                   n_digits,
                   0)
       return text if n_rows <= border
 
-      tail_start = [border, n_rows - border].max
-      tail_limit = n_rows - tail_start
-      tail_column_values = columns.collect do |column|
-        column.reverse_each.take(tail_limit).reverse
-      end
 
       if head_limit != tail_start
         format_ellipsis(text)
       end
 
       format_rows(text,
-                  columns,
-                  tail_column_values.transpose,
+                  column_formatters,
+                  column_formatters.collect(&:tail_values).transpose,
                   n_digits,
                   tail_start)
 

data/lib/arrow/table-list-formatter.rb

@@ -22,12 +22,14 @@ module Arrow
     def format_header(text, columns)
     end
 
-    def format_rows(text, columns, rows, n_digits, start_offset)
+    def format_rows(text, column_formatters, rows, n_digits, start_offset)
       rows.each_with_index do |row, nth_row|
         text << ("=" * 20 + " #{start_offset + nth_row} " + "=" * 20 + "\n")
         row.each_with_index do |column_value, nth_column|
-          column = columns[nth_column]
-          text << "#{column.name}: #{column_value}\n"
+          column_formatter = column_formatters[nth_column]
+          formatted_name = column_formatter.name
+          formatted_value = column_formatter.format_value(column_value)
+          text << "#{formatted_name}: #{formatted_value}\n"
         end
       end
     end

data/lib/arrow/table-loader.rb

@@ -15,6 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
+require "uri"
+
 module Arrow
   class TableLoader
     class << self
@@ -31,6 +33,31 @@ module Arrow
     end
 
     def load
+      if @input.is_a?(URI)
+        custom_load_method = "load_from_uri"
+      elsif @input.is_a?(String) and ::File.directory?(@input)
+        custom_load_method = "load_from_directory"
+      else
+        custom_load_method = "load_from_file"
+      end
+      unless respond_to?(custom_load_method, true)
+        available_schemes = []
+        (methods(true) | private_methods(true)).each do |name|
+          match_data = /\Aload_from_/.match(name.to_s)
+          if match_data
+            available_schemes << match_data.post_match
+          end
+        end
+        message = "Arrow::Table load source must be one of ["
+        message << available_schemes.join(", ")
+        message << "]: #{@input.inspect}"
+        raise ArgumentError, message
+      end
+      __send__(custom_load_method)
+    end
+
+    private
+    def load_from_file
       format = @options[:format]
       custom_load_method = "load_as_#{format}"
       unless respond_to?(custom_load_method, true)
@@ -56,21 +83,24 @@ module Arrow
       end
     end
 
-    private
     def fill_options
       if @options[:format] and @options.key?(:compression)
         return
       end
 
-      if @input.is_a?(Buffer)
+      case @input
+      when Buffer
         info = {}
+      when URI
+        extension = PathExtension.new(@input.path)
+        info = extension.extract
       else
         extension = PathExtension.new(@input)
         info = extension.extract
       end
       format = info[:format]
       @options = @options.dup
-      if format and respond_to?("load_as_#{format}", true)
+      if format
         @options[:format] ||= format.to_sym
       else
         @options[:format] ||= :arrow
@@ -183,5 +213,13 @@ module Arrow
       table.instance_variable_set(:@input, input)
       table
     end
+
+    def load_as_json
+      input = open_input_stream
+      reader = JSONReader.new(input)
+      table = reader.read
+      table.instance_variable_set(:@input, input)
+      table
+    end
   end
 end

data/lib/arrow/table-saver.rb

@@ -32,6 +32,29 @@ module Arrow
     end
 
     def save
+      if @output.is_a?(URI)
+        custom_save_method = "save_to_uri"
+      else
+        custom_save_method = "save_to_file"
+      end
+      unless respond_to?(custom_save_method, true)
+        available_schemes = []
+        (methods(true) | private_methods(true)).each do |name|
+          match_data = /\Asave_to_/.match(name.to_s)
+          if match_data
+            available_schemes << match_data.post_match
+          end
+        end
+        message = "Arrow::Table save source must be one of ["
+        message << available_schemes.join(", ")
+        message << "]: #{@output.scheme.inspect}"
+        raise ArgumentError, message
+      end
+      __send__(custom_save_method)
+    end
+
+    private
+    def save_to_file
       format = @options[:format]
       custom_save_method = "save_as_#{format}"
       unless respond_to?(custom_save_method, true)
@@ -57,21 +80,24 @@ module Arrow
       end
     end
 
-    private
     def fill_options
       if @options[:format] and @options.key?(:compression)
         return
       end
 
-      if @output.is_a?(Buffer)
+      case @output
+      when Buffer
         info = {}
+      when URI
+        extension = PathExtension.new(@output.path)
+        info = extension.extract
       else
         extension = PathExtension.new(@output)
         info = extension.extract
       end
       format = info[:format]
       @options = @options.dup
-      if format and respond_to?("save_as_#{format}", true)
+      if format
         @options[:format] ||= format.to_sym
       else
         @options[:format] ||= :arrow

data/lib/arrow/table-table-formatter.rb

@@ -21,51 +21,27 @@ module Arrow
   # TODO: Almost codes should be implemented in Apache Arrow C++.
   class TableTableFormatter < TableFormatter
     private
-    def format_header(text, columns)
-      columns.each do |column|
+    def format_header(text, column_formatters)
+      column_formatters.each do |column_formatter|
         text << "\t"
-        text << format_column_name(column)
+        text << column_formatter.aligned_name
       end
       text << "\n"
     end
 
-    FLOAT_N_DIGITS = 10
-    def format_column_name(column)
-      case column.data_type
-      when TimestampDataType
-        "%*s" % [::Time.now.iso8601.size, column.name]
-      when FloatDataType, DoubleDataType
-        "%*s" % [FLOAT_N_DIGITS, column.name]
-      else
-        column.name
-      end
-    end
-
-    def format_rows(text, columns, rows, n_digits, start_offset)
+    def format_rows(text, column_formatters, rows, n_digits, start_offset)
       rows.each_with_index do |row, nth_row|
         text << ("%*d" % [n_digits, start_offset + nth_row])
         row.each_with_index do |column_value, nth_column|
           text << "\t"
-          column = columns[nth_column]
-          text << format_column_value(column, column_value)
+          column_formatter = column_formatters[nth_column]
+          aligned_name = column_formatter.aligned_name
+          text << column_formatter.format_value(column_value, aligned_name.size)
         end
         text << "\n"
       end
     end
 
-    def format_column_value(column, value)
-      case value
-      when ::Time
-        value.iso8601
-      when Float
-        "%*f" % [[column.name.size, FLOAT_N_DIGITS].max, value]
-      when Integer
-        "%*d" % [column.name.size, value]
-      else
-        "%-*s" % [column.name.size, value.to_s]
-      end
-    end
-
     def format_ellipsis(text)
       text << "...\n"
     end

data/lib/arrow/table.rb

@@ -195,8 +195,6 @@ module Arrow
     alias_method :size, :n_rows
     alias_method :length, :n_rows
 
-    alias_method :[], :find_column
-
     alias_method :slice_raw, :slice
 
     # @overload slice(offset, length)
@@ -236,6 +234,12 @@ module Arrow
     #   @return [Arrow::Table]
     #     The sub `Arrow::Table` that covers only rows of the range of indices.
     #
+    # @overload slice(conditions)
+    #
+    #   @param conditions [Hash] The conditions to select records.
+    #   @return [Arrow::Table]
+    #     The sub `Arrow::Table` that covers only rows matched by condition
+    #
     # @overload slice
     #
     #   @yield [slicer] Gives slicer that constructs condition to select records.
@@ -263,12 +267,37 @@ module Arrow
         expected_n_args = nil
         case args.size
         when 1
-          if args[0].is_a?(Integer)
+          case args[0]
+          when Integer
             index = args[0]
             index += n_rows if index < 0
             return nil if index < 0
             return nil if index >= n_rows
             return Record.new(self, index)
+          when Hash
+            condition_pairs = args[0]
+            slicer = Slicer.new(self)
+            conditions = []
+            condition_pairs.each do |key, value|
+              case value
+              when Range
+                # TODO: Optimize "begin <= key <= end" case by missing "between" kernel
+                # https://issues.apache.org/jira/browse/ARROW-9843
+                unless value.begin.nil?
+                  conditions << (slicer[key] >= value.begin)
+                end
+                unless value.end.nil?
+                  if value.exclude_end?
+                    conditions << (slicer[key] < value.end)
+                  else
+                    conditions << (slicer[key] <= value.end)
+                  end
+                end
+              else
+                conditions << (slicer[key] == value)
+              end
+            end
+            slicers << conditions.inject(:&)
           else
             slicers << args[0]
           end
@@ -397,41 +426,6 @@ module Arrow
       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)
@@ -442,8 +436,8 @@ module Arrow
       RollingWindow.new(self, size)
     end
 
-    def save(path, options={})
-      saver = TableSaver.new(self, path, options)
+    def save(output, options={})
+      saver = TableSaver.new(self, output, options)
       saver.save
     end
 

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

@@ -29,7 +29,7 @@ module Arrow
     #
     #     The unit must be second or millisecond.
     #
-    #   @example Create a time32 data type with {Arrow::TimeUnit}
+    #   @example Create a time32 data type with Arrow::TimeUnit
     #     Arrow::Time32DataType.new(Arrow::TimeUnit::MILLI)
     #
     #   @example Create a time32 data type with Symbol
@@ -45,7 +45,7 @@ module Arrow
     #
     #     The unit must be second or millisecond.
     #
-    #   @example Create a time32 data type with {Arrow::TimeUnit}
+    #   @example Create a time32 data type with Arrow::TimeUnit
     #     Arrow::Time32DataType.new(unit: Arrow::TimeUnit::MILLI)
     #
     #   @example Create a time32 data type with Symbol