red-arrow 5.0.0 → 6.0.0

data/lib/arrow/slicer.rb

@@ -16,9 +16,6 @@
 # under the License.
 
 module Arrow
-  # Experimental
-  #
-  # TODO: Almost codes should be implemented in Apache Arrow C++.
   class Slicer
     def initialize(table)
       @table = table
@@ -43,6 +40,21 @@ module Arrow
       super
     end
 
+    module Helper
+      class << self
+        def ensure_boolean(column)
+          case column.data_type
+          when Arrow::BooleanDataType
+            column.data
+          else
+            options = CastOptions.new
+            options.to_data_type = Arrow::BooleanDataType.new
+            Function.find("cast").execute([column.data], options).value
+          end
+        end
+      end
+    end
+
     class Condition
       def evaluate
         message = "Slicer::Condition must define \#evaluate: #{inspect}"
@@ -69,43 +81,28 @@ module Arrow
       end
 
       def evaluate
-        values1 = @condition1.evaluate.each
-        values2 = @condition2.evaluate.each
-        raw_array = []
-        begin
-          loop do
-            value1 = values1.next
-            value2 = values2.next
-            if value1.nil? or value2.nil?
-              raw_array << nil
-            else
-              raw_array << evaluate_value(value1, value2)
-            end
-          end
-        rescue StopIteration
-        end
-        BooleanArray.new(raw_array)
+        function.execute([@condition1.evaluate, @condition2.evaluate]).value
       end
     end
 
     class AndCondition < LogicalCondition
       private
-      def evaluate_value(value1, value2)
-        value1 and value2
+      def function
+        Function.find("and")
       end
     end
 
     class OrCondition < LogicalCondition
       private
-      def evaluate_value(value1, value2)
-        value1 or value2
+      def function
+        Function.find("or")
       end
     end
 
     class XorCondition < LogicalCondition
       private
-      def evaluate_value(value1, value2)
-        value1 ^ value2
+      def function
+        Function.find("xor")
       end
     end
 
@@ -115,21 +112,7 @@ module Arrow
       end
 
       def evaluate
-        data = @column.data
-
-        case @column.data_type
-        when BooleanDataType
-          data
-        else
-          if data.n_chunks == 1
-            data.get_chunk(0).cast(BooleanDataType.new, nil)
-          else
-            arrays = data.each_chunk.collect do |chunk|
-              chunk.cast(BooleanDataType.new, nil)
-            end
-            ChunkedArray.new(arrays)
-          end
-        end
+        Helper.ensure_boolean(@column)
       end
 
       def !@
@@ -187,23 +170,8 @@ module Arrow
       end
 
       def evaluate
-        data = @column.data
-        raw_array = []
-        data.each_chunk do |chunk|
-          if chunk.is_a?(BooleanArray)
-            boolean_array = chunk
-          else
-            boolean_array = chunk.cast(BooleanDataType.new, nil)
-          end
-          boolean_array.each do |value|
-            if value.nil?
-              raw_array << value
-            else
-              raw_array << !value
-            end
-          end
-        end
-        BooleanArray.new(raw_array)
+        data = Helper.ensure_boolean(@column)
+        Function.find("invert").execute([data]).value
       end
 
       def !@
@@ -222,19 +190,10 @@ module Arrow
       end
 
       def evaluate
-        case @value
-        when nil
-          raw_array = @column.collect(&:nil?)
-          BooleanArray.new(raw_array)
+        if @value.nil?
+          Function.find("is_null").execute([@column.data]).value
         else
-          raw_array = @column.collect do |value|
-            if value.nil?
-              nil
-            else
-              @value == value
-            end
-          end
-          BooleanArray.new(raw_array)
+          Function.find("equal").execute([@column.data, @value]).value
         end
       end
     end
@@ -250,25 +209,10 @@ module Arrow
       end
 
       def evaluate
-        case @value
-        when nil
-          if @column.n_nulls.zero?
-            raw_array = [true] * @column.n_rows
-          else
-            raw_array = @column.n_rows.times.collect do |i|
-              @column.valid?(i)
-            end
-          end
-          BooleanArray.new(raw_array)
+        if @value.nil?
+          Function.find("is_valid").execute([@column.data]).value
         else
-          raw_array = @column.collect do |value|
-            if value.nil?
-              nil
-            else
-              @value != value
-            end
-          end
-          BooleanArray.new(raw_array)
+          Function.find("not_equal").execute([@column.data, @value]).value
         end
       end
     end
@@ -284,14 +228,7 @@ module Arrow
       end
 
       def evaluate
-        raw_array = @column.collect do |value|
-          if value.nil?
-            nil
-          else
-            @value > value
-          end
-        end
-        BooleanArray.new(raw_array)
+        Function.find("less").execute([@column.data, @value]).value
       end
     end
 
@@ -306,14 +243,7 @@ module Arrow
       end
 
       def evaluate
-        raw_array = @column.collect do |value|
-          if value.nil?
-            nil
-          else
-            @value >= value
-          end
-        end
-        BooleanArray.new(raw_array)
+        Function.find("less_equal").execute([@column.data, @value]).value
       end
     end
 
@@ -328,14 +258,7 @@ module Arrow
       end
 
       def evaluate
-        raw_array = @column.collect do |value|
-          if value.nil?
-            nil
-          else
-            @value < value
-          end
-        end
-        BooleanArray.new(raw_array)
+        Function.find("greater").execute([@column.data, @value]).value
       end
     end
 
@@ -350,14 +273,7 @@ module Arrow
       end
 
       def evaluate
-        raw_array = @column.collect do |value|
-          if value.nil?
-            nil
-          else
-            @value <= value
-          end
-        end
-        BooleanArray.new(raw_array)
+        Function.find("greater_equal").execute([@column.data, @value]).value
       end
     end
 
@@ -372,18 +288,10 @@ module Arrow
       end
 
       def evaluate
-        values_index = {}
-        @values.each do |value|
-          values_index[value] = true
-        end
-        raw_array = @column.collect do |value|
-          if value.nil?
-            nil
-          else
-            values_index.key?(value)
-          end
-        end
-        BooleanArray.new(raw_array)
+        values = @values
+        values = Array.new(values) unless values.is_a?(Array)
+        options = SetLookupOptions.new(values)
+        Function.find("is_in").execute([@column.data], options).value
       end
     end
 
@@ -398,18 +306,11 @@ module Arrow
       end
 
       def evaluate
-        values_index = {}
-        @values.each do |value|
-          values_index[value] = true
-        end
-        raw_array = @column.collect do |value|
-          if value.nil?
-            nil
-          else
-            not values_index.key?(value)
-          end
-        end
-        BooleanArray.new(raw_array)
+        values = @values
+        values = Array.new(values) unless values.is_a?(Array)
+        options = SetLookupOptions.new(values)
+        booleans = Function.find("is_in").execute([@column.data], options).value
+        Function.find("invert").execute([booleans]).value
       end
     end
 

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/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