red-arrow 0.4.1 → 0.8.0

data/lib/arrow/table.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.
@@ -14,7 +14,31 @@
 
 module Arrow
   class Table
-    def each_column(&block)
+    alias_method :initialize_raw, :initialize
+    def initialize(schema_or_raw_table_or_columns, columns=nil)
+      if columns.nil?
+        if schema_or_raw_table_or_columns[0].is_a?(Column)
+          columns = schema_or_raw_table_or_columns
+          fields = columns.collect(&:field)
+          schema = Schema.new(fields)
+        else
+          raw_table = schema_or_raw_table_or_columns
+          fields = []
+          columns = []
+          raw_table.each do |name, array|
+            field = Field.new(name.to_s, array.value_data_type)
+            fields << field
+            columns << Column.new(field, array)
+          end
+          schema = Schema.new(fields)
+        end
+      else
+        schema = schema_or_raw_table_or_columns
+      end
+      initialize_raw(schema, columns)
+    end
+
+    def each_column
       return to_enum(__method__) unless block_given?
 
       n_columns.times do |i|
@@ -23,7 +47,289 @@ module Arrow
     end
 
     def columns
-      each_column.to_a
+      @columns ||= each_column.to_a
+    end
+
+    def each_record_batch
+      return to_enum(__method__) unless block_given?
+
+      reader = TableBatchReader.new(self)
+      while record_batch = reader.read_next
+        yield(record_batch)
+      end
+    end
+
+    # TODO
+    #
+    # @return [Arrow::Column, Array<Arrow::Column>, nil]
+    def [](*args)
+      if args.size == 1
+        case args[0]
+        when String, Symbol
+          find_column(args[0])
+        else
+          message = "#{self.class}\#[#{args[0].inspect}]: " +
+            "Must be String or Symbol"
+          raise ArgumentError, message
+        end
+      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
+
+    # TODO
+    #
+    # @return [Arrow::Table]
+    def slice(*slicers)
+      if block_given?
+        block_slicer = yield(Slicer.new(self))
+        case block_slicer
+        when nil
+          # Ignore
+        when ::Array
+          slicers.concat(block_slicer)
+        else
+          slicers << block_slicer
+        end
+      end
+      ranges = []
+      slicers.each do |slicer|
+        slicer = slicer.evaluate if slicer.respond_to?(:evaluate)
+        case slicer
+        when Integer
+          ranges << [slicer, slicer]
+        when Range
+          from = slicer.first
+          to = slicer.last
+          to -= 1 if slicer.exclude_end?
+          ranges << [from, to]
+        when ::Array
+          from = slicer[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]
+          end
+        else
+          message = "slicer must be Integer, Range, [from, to] or " +
+            "Arrow::BooleanArray, Arrow::Slicer::Condition: #{slicer.inspect}"
+          raise ArgumentError, message
+        end
+      end
+      slice_by_ranges(ranges)
+    end
+
+    # TODO
+    #
+    # @return [Arrow::Table]
+    def merge(other)
+      added_columns = {}
+      removed_columns = {}
+
+      case other
+      when Hash
+        other.each do |name, value|
+          name = name.to_s
+          if value
+            added_columns[name] = ensure_column(name, value)
+          else
+            removed_columns[name] = true
+          end
+        end
+      when Table
+        added_columns = {}
+        other.columns.each do |column|
+          added_columns[column.name] = column
+        end
+      else
+        message = "merge target must be Hash or Arrow::Table: " +
+          "<#{other.inspect}>: #{inspect}"
+        raise ArgumentError, message
+      end
+
+      new_columns = []
+      columns.each do |column|
+        column_name = column.name
+        new_column = added_columns.delete(column_name)
+        if new_column
+          new_columns << new_column
+          next
+        end
+        next if removed_columns.key?(column_name)
+        new_columns << column
+      end
+      added_columns.each do |name, new_column|
+        new_columns << new_column
+      end
+      new_fields = new_columns.collect do |new_column|
+        new_column.field
+      end
+      self.class.new(Schema.new(new_fields), new_columns)
+    end
+
+    alias_method :remove_column_raw, :remove_column
+    def remove_column(name_or_index)
+      case name_or_index
+      when String, Symbol
+        name = name_or_index.to_s
+        index = columns.index {|column| column.name == name}
+        if index.nil?
+          message = "unknown column: #{name_or_index.inspect}: #{inspect}"
+          raise KeyError.new(message)
+        end
+      else
+        index = name_or_index
+        index += n_columns if index < 0
+        if index < 0 or index >= n_columns
+          message = "out of index (0..#{n_columns - 1}): " +
+            "#{name_or_index.inspect}: #{inspect}"
+          raise IndexError.new(message)
+        end
+      end
+      remove_column_raw(index)
+    end
+
+    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
+
+    def to_s(options={})
+      formatter = TableFormatter.new(self, options)
+      formatter.format
+    end
+
+    def inspect
+      "#{super}\n#{to_s}"
+    end
+
+    def respond_to_missing?(name, include_private)
+      return true if find_column(name)
+      super
+    end
+
+    def method_missing(name, *args, &block)
+      if args.empty?
+        column = find_column(name)
+        return column if column
+      end
+      super
+    end
+
+    private
+    def find_column(name)
+      name = name.to_s
+      columns.find do |column|
+        column.name == name
+      end
+    end
+
+    def slice_by_ranges(ranges)
+      sliced_columns = columns.collect do |column|
+        chunks = []
+        arrays = column.data.each_chunk.to_a
+        offset = 0
+        offset_in_array = 0
+        ranges.each do |from, to|
+          range_size = to - from + 1
+          while range_size > 0
+            while offset + arrays.first.length - offset_in_array < from
+              offset += arrays.first.length - offset_in_array
+              arrays.shift
+              offset_in_array = 0
+            end
+            if offset < from
+              skipped_size = from - offset
+              offset += skipped_size
+              offset_in_array += skipped_size
+            end
+            array = arrays.first
+            array_length = array.length
+            rest_length = array_length - offset_in_array
+            if rest_length <= range_size
+              chunks << array.slice(offset_in_array, array_length)
+              offset += rest_length
+              range_size -= rest_length
+              offset_in_array = 0
+              arrays.shift
+            else
+              chunks << array.slice(offset_in_array, range_size)
+              offset += range_size
+              offset_in_array += range_size
+              range_size = 0
+            end
+          end
+        end
+        Column.new(column.field, ChunkedArray.new(chunks))
+      end
+
+      self.class.new(schema, sliced_columns)
+    end
+
+    def ensure_column(name, data)
+      case data
+      when Array
+        field = Field.new(name, data.value_data_type)
+        Column.new(field, data)
+      when Column
+        data
+      else
+        message = "column must be Arrow::Array or Arrow::Column: " +
+          "<#{name}>: <#{data.inspect}>: #{inspect}"
+        raise ArgumentError, message
+      end
     end
   end
 end

data/lib/arrow/timestamp-array.rb

@@ -0,0 +1,69 @@
+# Copyright 2017 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 TimestampArray
+    class << self
+      def new(unit, values)
+        data_type = TimestampDataType.new(unit)
+        builder = TimestampArrayBuilder.new(data_type)
+        builder.build(values)
+      end
+    end
+
+    def get_value(i)
+      to_time(get_raw_value(i))
+    end
+
+    def unit
+      @unit ||= value_data_type_unit
+    end
+
+    private
+    def unit_id
+      @unit_id ||= unit.nick.to_sym
+    end
+
+    def value_data_type_unit
+      data_type = value_data_type
+      if data_type.respond_to?(:unit)
+        data_type.unit
+      else
+        data_type_name = data_type.to_s
+        if data_type_name.end_with?("[s]")
+          TimeUnit::SECOND
+        elsif data_type_name.end_with?("[ms]")
+          TimeUnit::MILLI
+        elsif data_type_name.end_with?("[us]")
+          TimeUnit::MICRO
+        else
+          TimeUnit::NANO
+        end
+      end
+    end
+
+    def to_time(raw_value)
+      case unit_id
+      when :second
+        Time.at(raw_value)
+      when :milli
+        Time.at(*raw_value.divmod(1_000))
+      when :micro
+        Time.at(*raw_value.divmod(1_000_000))
+      else
+        Time.at(raw_value / 1_000_000_000.0)
+      end
+    end
+  end
+end

data/lib/arrow/version.rb

@@ -13,5 +13,5 @@
 # limitations under the License.
 
 module Arrow
-  VERSION = "0.4.1"
+  VERSION = "0.8.0"
 end

data/red-arrow.gemspec

@@ -38,6 +38,7 @@ Gem::Specification.new do |spec|
   spec.files = ["README.md", "Rakefile", "Gemfile", "#{spec.name}.gemspec"]
   spec.files += [".yardopts"]
   spec.files += Dir.glob("lib/**/*.rb")
+  spec.files += Dir.glob("image/*.*")
   spec.files += Dir.glob("doc/text/*")
   spec.test_files += Dir.glob("test/**/*")
   spec.extensions = ["dependency-check/Rakefile"]

data/test/fixture/with-header.csv

@@ -0,0 +1,4 @@
+name,score
+alice,10
+bob,29
+chris,-1

data/test/fixture/without-header.csv

@@ -0,0 +1,3 @@
+alice,10
+bob,29
+chris,-1

data/test/helper.rb

@@ -14,6 +14,9 @@
 
 require "arrow"
 
+require "pathname"
 require "tempfile"
 
 require "test-unit"
+
+require_relative "helper/fixture"

data/test/helper/fixture.rb

@@ -0,0 +1,25 @@
+# Copyright 2017 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 Helper
+  module Fixture
+    def fixture_dir
+      Pathname.new(__dir__).join("..", "fixture").expand_path
+    end
+
+    def fixture_path(*components)
+      fixture_dir.join(*components)
+    end
+  end
+end