red-arrow 0.8.0 → 0.8.1

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

@@ -0,0 +1,29 @@
+# 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 Date32ArrayBuilder
+    private
+    UNIX_EPOCH = Date.new(1970, 1, 1)
+    def convert_to_arrow_value(value)
+      value = value.to_date if value.respond_to?(:to_date)
+
+      if value.is_a?(Date)
+        (value - UNIX_EPOCH).to_i
+      else
+        value
+      end
+    end
+  end
+end

data/lib/arrow/date32-array.rb

@@ -0,0 +1,27 @@
+# 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 Date32Array
+    def get_value(i)
+      to_date(get_raw_value(i))
+    end
+
+    private
+    UNIX_EPOCH = 2440588
+    def to_date(raw_value)
+      Date.jd(UNIX_EPOCH + raw_value)
+    end
+  end
+end

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

@@ -0,0 +1,30 @@
+# 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 Date64ArrayBuilder
+    private
+    def convert_to_arrow_value(value)
+      if value.respond_to?(:to_time) and not value.is_a?(Time)
+        value = value.to_time
+      end
+
+      if value.is_a?(Time)
+        value.to_i * 1_000 + value.usec / 1000
+      else
+        value
+      end
+    end
+  end
+end

data/lib/arrow/date64-array.rb

@@ -0,0 +1,26 @@
+# 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 Date64Array
+    def get_value(i)
+      to_datetime(get_raw_value(i))
+    end
+
+    private
+    def to_datetime(raw_value)
+      Time.at(*raw_value.divmod(1_000)).to_datetime
+    end
+  end
+end

data/lib/arrow/loader.rb

@@ -33,14 +33,23 @@ module Arrow
       require "arrow/buffer"
       require "arrow/chunked-array"
       require "arrow/column"
+      require "arrow/csv-loader"
       require "arrow/csv-reader"
+      require "arrow/date32-array"
+      require "arrow/date32-array-builder"
+      require "arrow/date64-array"
+      require "arrow/date64-array-builder"
       require "arrow/field"
+      require "arrow/record"
       require "arrow/record-batch"
       require "arrow/slicer"
       require "arrow/table"
       require "arrow/table-formatter"
+      require "arrow/table-loader"
+      require "arrow/table-saver"
       require "arrow/tensor"
       require "arrow/timestamp-array"
+      require "arrow/timestamp-array-builder"
 
       require "arrow/record-batch-file-reader"
       require "arrow/record-batch-stream-reader"
@@ -67,7 +76,7 @@ module Arrow
           method_name = "get_value"
         end
         super(info, klass, method_name)
-      when "Arrow::TimestampArray"
+      when "Arrow::TimestampArray", "Arrow::Date32Array", "Arrow::Date64Array"
         case method_name
         when "get_value"
           method_name = "get_raw_value"

data/lib/arrow/record-batch.rb

@@ -12,57 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-require "arrow/record"
+require "arrow/record-containable"
 
 module Arrow
   class RecordBatch
+    include RecordContainable
     include Enumerable
 
-    def each(reuse_record: false)
-      unless block_given?
-        return to_enum(__method__, reuse_record: reuse_record)
-      end
-
-      if reuse_record
-        record = Record.new(self, nil)
-        n_rows.times do |i|
-          record.index = i
-          yield(record)
-        end
-      else
-        n_rows.times do |i|
-          yield(Record.new(self, i))
-        end
-      end
-    end
-
-    def find_column(name_or_index)
-      case name_or_index
-      when String, Symbol
-        name = name_or_index
-        index = resolve_name(name)
-      else
-        index = name_or_index
-      end
-      columns[index]
-    end
+    alias_method :each, :each_record
 
     alias_method :columns_raw, :columns
     def columns
       @columns ||= columns_raw
     end
 
-    private
-    def resolve_name(name)
-      (@name_to_index ||= build_name_to_index)[name.to_s]
+    def respond_to_missing?(name, include_private)
+      return true if find_column(name)
+      super
     end
 
-    def build_name_to_index
-      index = {}
-      schema.fields.each_with_index do |field, i|
-        index[field.name] = i
+    def method_missing(name, *args, &block)
+      if args.empty?
+        column = find_column(name)
+        return column if column
       end
-      index
+      super
     end
   end
 end

data/lib/arrow/record-containable.rb

@@ -0,0 +1,70 @@
+# 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.
+
+module Arrow
+  module RecordContainable
+    def each_column(&block)
+      return to_enum(__method__) unless block_given?
+
+      columns.each(&block)
+    end
+
+    def each_record(reuse_record: false)
+      unless block_given?
+        return to_enum(__method__, reuse_record: reuse_record)
+      end
+
+      if reuse_record
+        record = Record.new(self, nil)
+        n_rows.times do |i|
+          record.index = i
+          yield(record)
+        end
+      else
+        n_rows.times do |i|
+          yield(Record.new(self, i))
+        end
+      end
+    end
+
+    def find_column(name_or_index)
+      case name_or_index
+      when String, Symbol
+        name = name_or_index.to_s
+        index = resolve_column_name(name)
+        return nil if index.nil?
+        columns[index]
+      when Integer
+        index = name_or_index
+        columns[index]
+      else
+        message = "column name or index must be String, Symbol or Integer"
+        raise ArgumentError, message
+      end
+    end
+
+    private
+    def resolve_column_name(name)
+      (@column_name_to_index ||= build_column_name_resolve_table)[name]
+    end
+
+    def build_column_name_resolve_table
+      table = {}
+      schema.fields.each_with_index do |field, i|
+        table[field.name] = i
+      end
+      table
+    end
+  end
+end

data/lib/arrow/record.rb

@@ -15,25 +15,40 @@
 module Arrow
   class Record
     attr_accessor :index
-    def initialize(record_batch, index)
-      @record_batch = record_batch
+    def initialize(record_container, index)
+      @record_container = record_container
       @index = index
     end
 
     def [](column_name_or_column_index)
-      @record_batch.find_column(column_name_or_column_index)[@index]
+      column = @record_container.find_column(column_name_or_column_index)
+      return nil if column.nil?
+      column[@index]
     end
 
     def columns
-      @record_batch.columns
+      @record_container.columns
     end
 
     def to_h
       attributes = {}
-      @record_batch.schema.fields.each_with_index do |field, i|
+      @record_container.schema.fields.each_with_index do |field, i|
         attributes[field.name] = self[i]
       end
       attributes
     end
+
+    def respond_to_missing?(name, include_private)
+      return true if @record_container.find_column(name)
+      super
+    end
+
+    def method_missing(name, *args, &block)
+      if args.empty?
+        column = @record_container.find_column(name)
+        return column[@index] if column
+      end
+      super
+    end
   end
 end

data/lib/arrow/table-loader.rb

@@ -0,0 +1,117 @@
+# 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 TableLoader
+    class << self
+      def load(path, options={})
+        new(path, options).load
+      end
+    end
+
+    def initialize(path, options={})
+      @path = path
+      @options = options
+    end
+
+    def load
+      path = @path
+      path = path.to_path if path.respond_to?(:to_path)
+      format = @options[:format] || guess_format(path) || :arrow
+
+      custom_load_method = "load_as_#{format}"
+      unless respond_to?(custom_load_method, true)
+        available_formats = []
+        (methods(true) | private_methods(true)).each do |name|
+          match_data = /\Aload_as_/.match(name.to_s)
+          if match_data
+            available_formats << match_data.post_match
+          end
+        end
+        message = "Arrow::Table load format must be one of ["
+        message << available_formats.join(", ")
+        message << "]: #{format.inspect}"
+        raise ArgumentError, message
+      end
+      __send__(custom_load_method, path)
+    end
+
+    private
+    def guess_format(path)
+      extension = ::File.extname(path).gsub(/\A\./, "").downcase
+      return nil if extension.empty?
+
+      return extension if respond_to?("load_as_#{extension}", true)
+
+      nil
+    end
+
+    def load_raw(input, reader)
+      schema = reader.schema
+      chunked_arrays = []
+      reader.each do |record_batch|
+        record_batch.columns.each_with_index do |array, i|
+          chunked_array = (chunked_arrays[i] ||= [])
+          chunked_array << array
+        end
+      end
+      columns = schema.fields.collect.with_index do |field, i|
+        Column.new(field, ChunkedArray.new(chunked_arrays[i]))
+      end
+      table = Table.new(schema, columns)
+      table.instance_variable_set(:@input, input)
+      table
+    end
+
+    def load_as_arrow(path)
+      input = nil
+      reader = nil
+      error = nil
+      reader_class_candidates = [
+        RecordBatchFileReader,
+        RecordBatchStreamReader,
+      ]
+      reader_class_candidates.each do |reader_class_candidate|
+        input = MemoryMappedInputStream.new(path)
+        begin
+          reader = reader_class_candidate.new(input)
+        rescue Arrow::Error
+          error = $!
+        else
+          break
+        end
+      end
+      raise error if reader.nil?
+      load_raw(input, reader)
+    end
+
+    def load_as_batch(path)
+      input = MemoryMappedInputStream.new(path)
+      reader = RecordBatchFileReader.new(input)
+      load_raw(input, reader)
+    end
+
+    def load_as_stream(path)
+      input = MemoryMappedInputStream.new(path)
+      reader = RecordBatchStreamReader.new(input)
+      load_raw(input, reader)
+    end
+
+    def load_as_csv(path)
+      options = @options.dup
+      options.delete(:format)
+      CSVLoader.load(Pathname.new(path), options)
+    end
+  end
+end