red-arrow 0.8.1 → 0.8.2

data/lib/arrow/loader.rb

@@ -1,4 +1,4 @@
-# Copyright 2017 Kouhei Sutou <kou@clear-code.com>
+# Copyright 2017-2018 Kouhei Sutou <kou@clear-code.com>
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -35,6 +35,7 @@ module Arrow
       require "arrow/column"
       require "arrow/csv-loader"
       require "arrow/csv-reader"
+      require "arrow/data-type"
       require "arrow/date32-array"
       require "arrow/date32-array-builder"
       require "arrow/date64-array"
@@ -45,6 +46,8 @@ module Arrow
       require "arrow/slicer"
       require "arrow/table"
       require "arrow/table-formatter"
+      require "arrow/table-list-formatter"
+      require "arrow/table-table-formatter"
       require "arrow/table-loader"
       require "arrow/table-saver"
       require "arrow/tensor"

data/lib/arrow/slicer.rb

@@ -110,26 +110,19 @@ module Arrow
 
       def evaluate
         data = @column.data
-        if data.n_chunks == 1
-          array = data.get_chunk(0)
-          if array.is_a?(BooleanArray)
-            array
-          else
-            array.cast(BooleanDataType.new)
-          end
+
+        case @column.data_type
+        when BooleanDataType
+          data
         else
-          raw_array = []
-          data.each_chunk do |chunk|
-            if chunk.is_a?(BooleanArray)
-              boolean_array = chunk
-            else
-              boolean_array = chunk.cast(BooleanDataType.new)
-            end
-            boolean_array.each do |value|
-              raw_array << value
+          if data.n_chunks == 1
+            data.get_chunk(0).cast(BooleanDataType.new)
+          else
+            arrays = data.each_chunk.collect do |chunk|
+              chunk.cast(BooleanDataType.new)
             end
+            ChunkedArray.new(arrays)
           end
-          BooleanArray.new(raw_array)
         end
       end
 
@@ -141,6 +134,10 @@ module Arrow
         self == nil
       end
 
+      def valid?
+        self != nil
+      end
+
       def ==(value)
         EqualCondition.new(@column, value)
       end
@@ -165,6 +162,10 @@ module Arrow
         GreaterEqualCondition.new(@column, value)
       end
 
+      def in?(values)
+        InCondition.new(@column, values)
+      end
+
       def select(&block)
         SelectCondition.new(@column, block)
       end
@@ -245,8 +246,12 @@ module Arrow
       def evaluate
         case @value
         when nil
-          raw_array = @column.collect do |value|
-            not value.nil?
+          if @column.n_nulls.zero?
+            raw_array = [true] * @column.length
+          else
+            raw_array = @column.length.times.collect do |i|
+              @column.valid?(i)
+            end
           end
           BooleanArray.new(raw_array)
         else
@@ -350,6 +355,58 @@ module Arrow
       end
     end
 
+    class InCondition < Condition
+      def initialize(column, values)
+        @column = column
+        @values = values
+      end
+
+      def !@
+        NotInCondition.new(@column, @values)
+      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)
+      end
+    end
+
+    class NotInCondition < Condition
+      def initialize(column, values)
+        @column = column
+        @values = values
+      end
+
+      def !@
+        InCondition.new(@column, @values)
+      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)
+      end
+    end
+
     class SelectCondition < Condition
       def initialize(column, block)
         @column = column

data/lib/arrow/table-formatter.rb

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-require "time"
-
 module Arrow
   class TableFormatter
     def initialize(table, options={})
@@ -24,65 +22,41 @@ module Arrow
     def format
       text = ""
       columns = @table.columns
-      columns.each do |column|
-        text << "\t"
-        text << format_column_name(column)
-      end
-      text << "\n"
+      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
-      [border, n_rows].min.times do |i|
-        format_row(text, columns, i, n_digits)
+      head_limit = [border, n_rows].min
+      head_column_values = columns.collect do |column|
+        column.each.take(head_limit)
       end
+      format_rows(text,
+                  columns,
+                  head_column_values.transpose,
+                  n_digits,
+                  0)
       return text if n_rows <= border
 
-      text << "...\n"
-      [border, n_rows - border].max.upto(n_rows - 1) do |i|
-        format_row(text, columns, i, n_digits)
+      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
 
-      text
-    end
-
-    private
-    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
+      if head_limit != tail_start
+        format_ellipsis(text)
       end
-    end
 
-    def format_row(text, columns, i, n_digits)
-      text << ("%*d" % [n_digits, i])
-      columns.each do |column|
-        text << "\t"
-        text << format_column_value(column, i)
-      end
-      text << "\n"
-      text
-    end
+      format_rows(text,
+                  columns,
+                  tail_column_values.transpose,
+                  n_digits,
+                  tail_start)
 
-    def format_column_value(column, i)
-      value = column[i]
-      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
+      text
     end
   end
 end

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

@@ -0,0 +1,35 @@
+# Copyright 2018 Kouhei Sutou <kou@clear-code.com>
+#
+# Licensed 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 TableListFormatter < TableFormatter
+    private
+    def format_header(text, columns)
+    end
+
+    def format_rows(text, columns, 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"
+        end
+      end
+    end
+
+    def format_ellipsis(text)
+      text << "...\n"
+    end
+  end
+end

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

@@ -0,0 +1,69 @@
+# Copyright 2017-2018 Kouhei Sutou <kou@clear-code.com>
+#
+# Licensed 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.
+
+require "time"
+
+module Arrow
+  class TableTableFormatter < TableFormatter
+    private
+    def format_header(text, columns)
+      columns.each do |column|
+        text << "\t"
+        text << format_column_name(column)
+      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)
+      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)
+        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
+  end
+end

data/lib/arrow/table.rb

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+require "arrow/group"
 require "arrow/record-containable"
 
 module Arrow
@@ -61,24 +62,10 @@ module Arrow
       end
     end
 
-    # TODO
-    #
-    # @return [Arrow::Column, Array<Arrow::Column>, nil]
-    def [](*args)
-      if args.size == 1
-        find_column(args[0])
-      else
-        new_columns = args.collect do |column_name|
-          column = find_column(column_name)
-          if column.nil?
-            message = "unknown column: <#{column_name.inspect}>: #{inspect}"
-            raise ArgumentError, message
-          end
-          column
-        end
-        self.class.new(schema, new_columns)
-      end
-    end
+    alias_method :size, :n_rows
+    alias_method :length, :n_rows
+
+    alias_method :[], :find_column
 
     # TODO
     #
@@ -100,39 +87,32 @@ module Arrow
         slicer = slicer.evaluate if slicer.respond_to?(:evaluate)
         case slicer
         when Integer
+          slicer += n_rows if slicer < 0
           ranges << [slicer, slicer]
         when Range
           from = slicer.first
           to = slicer.last
           to -= 1 if slicer.exclude_end?
+          from += n_rows if from < 0
+          to += n_rows if to < 0
           ranges << [from, to]
         when ::Array
           from = slicer[0]
+          from += n_rows if from < 0
           to = from + slicer[1] - 1
           ranges << [from, to]
-        when BooleanArray
-          in_target = false
-          target_start = nil
-          slicer.each_with_index do |is_target, i|
-            if is_target
-              unless in_target
-                target_start = i
-                in_target = true
-              end
-            else
-              if in_target
-                ranges << [target_start, i - 1]
-                target_start = nil
-                in_target = false
-              end
-            end
-          end
-          if in_target
-            ranges << [target_start, slicer.length - 1]
+        when ChunkedArray
+          offset = 0
+          slicer.each_chunk do |array|
+            boolean_array_to_slice_ranges(array, offset, ranges)
+            offset += array.length
           end
+        when BooleanArray
+          boolean_array_to_slice_ranges(slicer, 0, ranges)
         else
-          message = "slicer must be Integer, Range, [from, to] or " +
-            "Arrow::BooleanArray, Arrow::Slicer::Condition: #{slicer.inspect}"
+          message = "slicer must be Integer, Range, [from, to], " +
+            "Arrow::ChunkedArray of Arrow::BooleanArray, " +
+            "Arrow::BooleanArray or Arrow::Slicer::Condition: #{slicer.inspect}"
           raise ArgumentError, message
         end
       end
@@ -209,6 +189,9 @@ 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?
@@ -241,13 +224,30 @@ module Arrow
       self.class.new(selected_columns)
     end
 
+    def group(*keys)
+      Group.new(self, keys)
+    end
+
     def save(path, options={})
       saver = TableSaver.new(self, path, options)
       saver.save
     end
 
+    def pack
+      packed_columns = columns.collect do |column|
+        column.pack
+      end
+      self.class.new(schema, packed_columns)
+    end
+
     def to_s(options={})
-      formatter = TableFormatter.new(self, options)
+      case options[:format]
+      when :list
+        formatter_class = TableListFormatter
+      else
+        formatter_class = TableTableFormatter
+      end
+      formatter = formatter_class.new(self, options)
       formatter.format
     end
 
@@ -269,6 +269,28 @@ module Arrow
     end
 
     private
+    def boolean_array_to_slice_ranges(array, offset, ranges)
+      in_target = false
+      target_start = nil
+      array.each_with_index do |is_target, i|
+        if is_target
+          unless in_target
+            target_start = offset + i
+            in_target = true
+          end
+        else
+          if in_target
+            ranges << [target_start, offset + i - 1]
+            target_start = nil
+            in_target = false
+          end
+        end
+      end
+      if in_target
+        ranges << [target_start, offset + array.length - 1]
+      end
+    end
+
     def slice_by_ranges(ranges)
       sliced_columns = columns.collect do |column|
         chunks = []