red-arrow 0.8.0 → 0.8.1

data/lib/arrow/table-saver.rb

@@ -0,0 +1,93 @@
+# 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 TableSaver
+    class << self
+      def save(table, path, options={})
+        new(table, path, options).save
+      end
+    end
+
+    def initialize(table, path, options={})
+      @table = table
+      @path = path
+      @options = options
+    end
+
+    def save
+      path = @path
+      path = path.to_path if path.respond_to?(:to_path)
+      format = @options[:format] || guess_format(path) || :arrow
+
+      custom_save_method = "save_as_#{format}"
+      unless respond_to?(custom_save_method, true)
+        available_formats = []
+        (methods(true) | private_methods(true)).each do |name|
+          match_data = /\Asave_as_/.match(name.to_s)
+          if match_data
+            available_formats << match_data.post_match
+          end
+        end
+        message = "Arrow::Table save format must be one of ["
+        message << available_formats.join(", ")
+        message << "]: #{format.inspect}"
+        raise ArgumentError, message
+      end
+      __send__(custom_save_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?("save_as_#{extension}", true)
+
+      nil
+    end
+
+    def save_raw(writer_class, path)
+      FileOutputStream.open(path, false) do |output|
+        writer_class.open(output, @table.schema) do |writer|
+          writer.write_table(@table)
+        end
+      end
+    end
+
+    def save_as_arrow(path)
+      save_as_batch(path)
+    end
+
+    def save_as_batch(path)
+      save_raw(RecordBatchFileWriter, path)
+    end
+
+    def save_as_stream(path)
+      save_raw(RecordBatchStreamWriter, path)
+    end
+
+    def save_as_csv(path)
+      CSV.open(path, "w") do |csv|
+        names = @table.schema.fields.collect(&:name)
+        csv << names
+        @table.each_record(reuse_record: true) do |record|
+          csv << names.collect do |name|
+            record[name]
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow/table.rb

@@ -12,8 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+require "arrow/record-containable"
+
 module Arrow
   class Table
+    include RecordContainable
+
+    class << self
+      def load(path, options={})
+        TableLoader.load(path, options)
+      end
+    end
+
     alias_method :initialize_raw, :initialize
     def initialize(schema_or_raw_table_or_columns, columns=nil)
       if columns.nil?
@@ -38,16 +48,8 @@ module Arrow
       initialize_raw(schema, columns)
     end
 
-    def each_column
-      return to_enum(__method__) unless block_given?
-
-      n_columns.times do |i|
-        yield(get_column(i))
-      end
-    end
-
     def columns
-      @columns ||= each_column.to_a
+      @columns ||= n_columns.times.collect {|i| get_column(i)}
     end
 
     def each_record_batch
@@ -64,19 +66,12 @@ module Arrow
     # @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
+        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}"
+            message = "unknown column: <#{column_name.inspect}>: #{inspect}"
             raise ArgumentError, message
           end
           column
@@ -246,6 +241,11 @@ module Arrow
       self.class.new(selected_columns)
     end
 
+    def save(path, options={})
+      saver = TableSaver.new(self, path, options)
+      saver.save
+    end
+
     def to_s(options={})
       formatter = TableFormatter.new(self, options)
       formatter.format
@@ -269,13 +269,6 @@ module Arrow
     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 = []

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

@@ -0,0 +1,59 @@
+# 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 TimestampArrayBuilder
+    # TODO: Workaround for Apache Arrow GLib 0.8.0
+    alias_method :initialize_raw, :initialize
+    def initialize(data_type)
+      initialize_raw(data_type)
+      data_type_name = data_type.to_s
+      if data_type_name.end_with?("[s]")
+        @unit_id = :second
+      elsif data_type_name.end_with?("[ms]")
+        @unit_id = :milli
+      elsif data_type_name.end_with?("[us]")
+        @unit_id = :micro
+      else
+        @unit_id = :nano
+      end
+    end
+
+    private
+    def unit_id
+      @unit_id ||= unit.nick.to_sym
+    end
+
+    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)
+        case unit_id
+        when :second
+          value.to_i
+        when :milli
+          value.to_i * 1_000 + value.usec / 1000
+        when :micro
+          value.to_i * 1_000_000 + value.usec
+        else
+          value.to_i * 1_000_000_000 + value.nsec
+        end
+      else
+        value
+      end
+    end
+  end
+end

data/lib/arrow/version.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.
@@ -13,5 +13,5 @@
 # limitations under the License.
 
 module Arrow
-  VERSION = "0.8.0"
+  VERSION = "0.8.1"
 end

data/test/test-array-builder.rb

@@ -14,58 +14,96 @@
 
 class ArrayBuilderTest < Test::Unit::TestCase
   sub_test_case(".build") do
-    test("empty") do
-      array = Arrow::Int32ArrayBuilder.build([])
-      assert_equal([],
-                   array.to_a)
+    def assert_build(builder_class, raw_array)
+      array = builder_class.build(raw_array)
+      assert_equal(raw_array, array.to_a)
     end
 
-    test("values") do
-      array = Arrow::Int32ArrayBuilder.build([1, -2])
-      assert_equal([1, -2],
-                   array.to_a)
-    end
+    sub_test_case("generic builder") do
+      test("strings") do
+        assert_build(Arrow::ArrayBuilder,
+                     ["Hello", nil, "World"])
+      end
 
-    test("values, nils") do
-      array = Arrow::Int32ArrayBuilder.build([1, -2, nil, nil])
-      assert_equal([1, -2, nil, nil],
-                   array.to_a)
-    end
+      test("boolean") do
+        assert_build(Arrow::ArrayBuilder,
+                     [true, nil, false])
+      end
 
-    test("values, nils, values") do
-      array = Arrow::Int32ArrayBuilder.build([1, -2, nil, nil, 3, -4])
-      assert_equal([1, -2, nil, nil, 3, -4],
-                   array.to_a)
-    end
+      test("positive integers") do
+        assert_build(Arrow::ArrayBuilder,
+                     [1, nil, 2, nil, 3])
+      end
 
-    test("values, nils, values, nils") do
-      array = Arrow::Int32ArrayBuilder.build([1, -2, nil, nil, 3, -4, nil, nil])
-      assert_equal([1, -2, nil, nil, 3, -4, nil, nil],
-                   array.to_a)
-    end
+      test("negative integers") do
+        assert_build(Arrow::ArrayBuilder,
+                     [nil, -1, nil, -2, nil, -3])
+      end
 
-    test("nils") do
-      array = Arrow::Int32ArrayBuilder.build([nil, nil])
-      assert_equal([nil, nil],
-                   array.to_a)
-    end
+      test("times") do
+        assert_build(Arrow::ArrayBuilder,
+                     [Time.at(0), Time.at(1), Time.at(2)])
+      end
 
-    test("nils, values") do
-      array = Arrow::Int32ArrayBuilder.build([nil, nil, 3, -4])
-      assert_equal([nil, nil, 3, -4],
-                   array.to_a)
-    end
+      test("dates") do
+        assert_build(Arrow::ArrayBuilder,
+                     [Date.new(2018, 1, 4), Date.new(2018, 1, 5)])
+      end
 
-    test("nils, values, nil") do
-      array = Arrow::Int32ArrayBuilder.build([nil, nil, 3, -4, nil, nil])
-      assert_equal([nil, nil, 3, -4, nil, nil],
-                   array.to_a)
+      test("datetimes") do
+        assert_build(Arrow::ArrayBuilder,
+                     [
+                       DateTime.new(2018, 1, 4, 23, 18, 23),
+                       DateTime.new(2018, 1, 5, 0, 23, 21),
+                     ])
+      end
     end
 
-    test("nils, values, nil, values") do
-      array = Arrow::Int32ArrayBuilder.build([nil, nil, 3, -4, nil, nil, 5, -6])
-      assert_equal([nil, nil, 3, -4, nil, nil, 5, -6],
-                   array.to_a)
+    sub_test_case("specific builder") do
+      test("empty") do
+        assert_build(Arrow::Int32ArrayBuilder,
+                     [])
+      end
+
+      test("values") do
+        assert_build(Arrow::Int32ArrayBuilder,
+                     [1, -2])
+      end
+
+      test("values, nils") do
+        assert_build(Arrow::Int32ArrayBuilder,
+                     [1, -2, nil, nil])
+      end
+
+      test("values, nils, values") do
+        assert_build(Arrow::Int32ArrayBuilder,
+                     [1, -2, nil, nil, 3, -4])
+      end
+
+      test("values, nils, values, nils") do
+        assert_build(Arrow::Int32ArrayBuilder,
+                     [1, -2, nil, nil, 3, -4, nil, nil])
+      end
+
+      test("nils") do
+        assert_build(Arrow::Int32ArrayBuilder,
+                     [nil, nil])
+      end
+
+      test("nils, values") do
+        assert_build(Arrow::Int32ArrayBuilder,
+                     [nil, nil, 3, -4])
+      end
+
+      test("nils, values, nil") do
+        assert_build(Arrow::Int32ArrayBuilder,
+                     [nil, nil, 3, -4, nil, nil])
+      end
+
+      test("nils, values, nil, values") do
+        assert_build(Arrow::Int32ArrayBuilder,
+                     [nil, nil, 3, -4, nil, nil, 5, -6])
+      end
     end
   end
 end

data/test/test-csv-loader.rb

@@ -0,0 +1,79 @@
+# 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.
+
+class CSVLoaderTest < Test::Unit::TestCase
+  include Helper::Fixture
+
+  sub_test_case(".load") do
+    test("String: data: with header") do
+      data = fixture_path("with-header.csv").read
+      assert_equal(<<-TABLE, Arrow::CSVLoader.load(data).to_s)
+	name	score
+0	alice	   10
+1	bob 	   29
+2	chris	   -1
+      TABLE
+    end
+
+    test("String: data: without header") do
+      data = fixture_path("without-header.csv").read
+      assert_equal(<<-TABLE, Arrow::CSVLoader.load(data).to_s)
+	0	1
+0	alice	10
+1	bob	29
+2	chris	-1
+      TABLE
+    end
+
+    test("String: path: with header") do
+      path = fixture_path("with-header.csv").to_s
+      assert_equal(<<-TABLE, Arrow::CSVLoader.load(path).to_s)
+	name	score
+0	alice	   10
+1	bob 	   29
+2	chris	   -1
+      TABLE
+    end
+
+    test("String: path: without header") do
+      path = fixture_path("without-header.csv").to_s
+      assert_equal(<<-TABLE, Arrow::CSVLoader.load(path).to_s)
+	0	1
+0	alice	10
+1	bob	29
+2	chris	-1
+      TABLE
+    end
+
+    test("Pathname: with header") do
+      path = fixture_path("with-header.csv")
+      assert_equal(<<-TABLE, Arrow::CSVLoader.load(path).to_s)
+	name	score
+0	alice	   10
+1	bob 	   29
+2	chris	   -1
+      TABLE
+    end
+
+    test("Pathname: without header") do
+      path = fixture_path("without-header.csv")
+      assert_equal(<<-TABLE, Arrow::CSVLoader.load(path).to_s)
+	0	1
+0	alice	10
+1	bob	29
+2	chris	-1
+      TABLE
+    end
+  end
+end