red-arrow 0.11.0 → 0.12.0

data/test/test-record-batch.rb

@@ -16,42 +16,97 @@
 # under the License.
 
 class RecordBatchTest < Test::Unit::TestCase
-  sub_test_case(".each") do
-    setup do
-      fields = [
-        Arrow::Field.new("count", :uint32),
+  sub_test_case(".new") do
+    def setup
+      @schema = Arrow::Schema.new(visible: :boolean,
+                                  count: :uint32)
+    end
+
+    test("[Schema, records]") do
+      records = [
+        {visible: true, count: 1},
+        nil,
+        [false, 3],
       ]
-      @schema = Arrow::Schema.new(fields)
-      @counts = Arrow::UInt32Array.new([1, 2, 4, 8])
-      @record_batch = Arrow::RecordBatch.new(@schema, @counts.length, [@counts])
+      record_batch = Arrow::RecordBatch.new(@schema, records)
+      assert_equal([
+                     {"visible" => true,  "count" => 1},
+                     {"visible" => nil,   "count" => nil},
+                     {"visible" => false, "count" => 3},
+                   ],
+                   record_batch.each_record.collect(&:to_h))
     end
 
-    test("default") do
-      records = []
-      @record_batch.each do |record|
-        records << [record, record.index]
-      end
+    test("[Schema, columns]") do
+      columns = {
+        visible: [true, nil, false],
+        count: [1, 2, nil],
+      }
+      record_batch = Arrow::RecordBatch.new(@schema, columns)
       assert_equal([
-                     [0, 0],
-                     [1, 1],
-                     [2, 2],
-                     [3, 3],
+                     {"visible" => true,  "count" => 1},
+                     {"visible" => nil,   "count" => 2},
+                     {"visible" => false, "count" => nil},
                    ],
-                   records.collect {|record, i| [record.index, i]})
+                   record_batch.each_record.collect(&:to_h))
     end
 
-    test("reuse_record: true") do
-      records = []
-      @record_batch.each(reuse_record: true) do |record|
-        records << [record, record.index]
-      end
+    test("[Schema, n_rows, columns]") do
+      columns = [
+        Arrow::BooleanArray.new([true, nil, false]),
+        Arrow::UInt32Array.new([1, 2, nil]),
+      ]
+      n_rows = columns[0].length
+      record_batch = Arrow::RecordBatch.new(@schema, n_rows, columns)
       assert_equal([
-                     [3, 0],
-                     [3, 1],
-                     [3, 2],
-                     [3, 3],
+                     {"visible" => true,  "count" => 1},
+                     {"visible" => nil,   "count" => 2},
+                     {"visible" => false, "count" => nil},
                    ],
-                   records.collect {|record, i| [record.index, i]})
+                   record_batch.each_record.collect(&:to_h))
+    end
+  end
+
+  sub_test_case("instance methods") do
+    def setup
+      @schema = Arrow::Schema.new(count: :uint32)
+      @counts = Arrow::UInt32Array.new([1, 2, 4, 8])
+      @record_batch = Arrow::RecordBatch.new(@schema, @counts.length, [@counts])
+    end
+
+    sub_test_case("#each") do
+      test("default") do
+        records = []
+        @record_batch.each do |record|
+          records << [record, record.index]
+        end
+        assert_equal([
+                       [0, 0],
+                       [1, 1],
+                       [2, 2],
+                       [3, 3],
+                     ],
+                     records.collect {|record, i| [record.index, i]})
+      end
+
+      test("reuse_record: true") do
+        records = []
+        @record_batch.each(reuse_record: true) do |record|
+          records << [record, record.index]
+        end
+        assert_equal([
+                       [3, 0],
+                       [3, 1],
+                       [3, 2],
+                       [3, 3],
+                     ],
+                     records.collect {|record, i| [record.index, i]})
+      end
+    end
+
+    test("#to_table") do
+      assert_equal(Arrow::Table.new(@schema, [@counts]),
+                   @record_batch.to_table)
     end
   end
 end

data/test/test-schema.rb

@@ -0,0 +1,104 @@
+# 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.
+
+class SchemaTest < Test::Unit::TestCase
+  def setup
+    @count_field = Arrow::Field.new("count", :uint32)
+    @visible_field = Arrow::Field.new("visible", :boolean)
+  end
+
+  sub_test_case(".new") do
+    test("[Arrow::Field]") do
+      fields = [
+        @count_field,
+        @visible_field,
+      ]
+      assert_equal("count: uint32\n" +
+                   "visible: bool",
+                   Arrow::Schema.new(fields).to_s)
+    end
+
+    test("[Arrow::Field, Hash]") do
+      fields = [
+        @count_field,
+        {name: "visible", type: :boolean},
+      ]
+      assert_equal("count: uint32\n" +
+                   "visible: bool",
+                   Arrow::Schema.new(fields).to_s)
+    end
+
+    test("{String, Symbol => Arrow::DataType}") do
+      fields = {
+        "count" => Arrow::UInt32DataType.new,
+        :visible => :boolean,
+      }
+      assert_equal("count: uint32\n" +
+                   "visible: bool",
+                   Arrow::Schema.new(fields).to_s)
+    end
+
+    test("{String, Symbol => Hash}") do
+      fields = {
+        "count" => {type: :uint32},
+        :tags => {
+          type: :list,
+          field: {
+            name: "tag",
+            type: :string,
+          },
+        },
+      }
+      assert_equal("count: uint32\n" +
+                   "tags: list<tag: string>",
+                   Arrow::Schema.new(fields).to_s)
+    end
+  end
+
+  sub_test_case("instance methods") do
+    def setup
+      super
+      @schema = Arrow::Schema.new([@count_field, @visible_field])
+    end
+
+    sub_test_case("#[]") do
+      test("[String]") do
+        assert_equal([@count_field, @visible_field],
+                     [@schema["count"], @schema["visible"]])
+      end
+
+      test("[Symbol]") do
+        assert_equal([@count_field, @visible_field],
+                     [@schema[:count], @schema[:visible]])
+      end
+
+      test("[Integer]") do
+        assert_equal([@count_field, @visible_field],
+                     [@schema[0], @schema[1]])
+      end
+
+      test("[invalid]") do
+        invalid = []
+        message = "field name or index must be String, Symbol or Integer"
+        message << ": <#{invalid.inspect}>"
+        assert_raise(ArgumentError.new(message)) do
+          @schema[invalid]
+        end
+      end
+    end
+  end
+end

data/test/test-sparse-union-data-type.rb

@@ -0,0 +1,41 @@
+# 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.
+
+class SparseUnionDataTypeTest < Test::Unit::TestCase
+  sub_test_case(".new") do
+    def setup
+      @fields = [
+        Arrow::Field.new("visible", :boolean),
+        {
+          name: "count",
+          type: :int32,
+        },
+      ]
+    end
+
+    test("ordered arguments") do
+      assert_equal("union[sparse]<visible: bool=2, count: int32=9>",
+                   Arrow::SparseUnionDataType.new(@fields, [2, 9]).to_s)
+    end
+
+    test("description") do
+      assert_equal("union[sparse]<visible: bool=2, count: int32=9>",
+                   Arrow::SparseUnionDataType.new(fields: @fields,
+                                                  type_codes: [2, 9]).to_s)
+    end
+  end
+end

data/test/test-struct-array-builder.rb

@@ -0,0 +1,180 @@
+# 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.
+
+class StructArrayBuilderTest < Test::Unit::TestCase
+  def setup
+    @data_type = Arrow::StructDataType.new(visible: {type: :boolean},
+                                           count: {type: :uint64})
+    @builder = Arrow::StructArrayBuilder.new(@data_type)
+  end
+
+  sub_test_case("#append_value") do
+    test("nil") do
+      @builder.append_value(nil)
+      array = @builder.finish
+      assert_equal([
+                     [nil],
+                     [nil],
+                   ],
+                   [
+                     array.find_field(0).to_a,
+                     array.find_field(1).to_a,
+                   ])
+    end
+
+    test("Array") do
+      @builder.append_value([true, 1])
+      array = @builder.finish
+      assert_equal([
+                     [true],
+                     [1],
+                   ],
+                   [
+                     array.find_field(0).to_a,
+                     array.find_field(1).to_a,
+                   ])
+    end
+
+    test("Arrow::Struct") do
+      source_array = Arrow::StructArray.new(@data_type, [[true, 1]])
+      struct = source_array.get_value(0)
+      @builder.append_value(struct)
+      array = @builder.finish
+      assert_equal([
+                     [true],
+                     [1],
+                   ],
+                   [
+                     array.find_field(0).to_a,
+                     array.find_field(1).to_a,
+                   ])
+    end
+
+    test("Hash") do
+      @builder.append_value(count: 1, visible: true)
+      array = @builder.finish
+      assert_equal([
+                     [true],
+                     [1],
+                   ],
+                   [
+                     array.find_field(0).to_a,
+                     array.find_field(1).to_a,
+                   ])
+    end
+  end
+
+  sub_test_case("#append_values") do
+    test("[nil]") do
+      @builder.append_values([nil])
+      array = @builder.finish
+      assert_equal([
+                     [nil],
+                     [nil],
+                   ],
+                   [
+                     array.find_field(0).to_a,
+                     array.find_field(1).to_a,
+                   ])
+    end
+
+    test("[Array]") do
+      @builder.append_values([[true, 1]])
+      array = @builder.finish
+      assert_equal([
+                     [true],
+                     [1],
+                   ],
+                   [
+                     array.find_field(0).to_a,
+                     array.find_field(1).to_a,
+                   ])
+    end
+
+    test("[Hash]") do
+      @builder.append_values([{count: 1, visible: true}])
+      array = @builder.finish
+      assert_equal([
+                     [true],
+                     [1],
+                   ],
+                   [
+                     array.find_field(0).to_a,
+                     array.find_field(1).to_a,
+                   ])
+    end
+
+    test("[nil, Array, Hash]") do
+      @builder.append_values([
+                               nil,
+                               [true, 1],
+                               {count: 2, visible: false},
+                             ])
+      array = @builder.finish
+      assert_equal([
+                     [nil, true, false],
+                     [nil, 1, 2],
+                   ],
+                   [
+                     array.find_field(0).to_a,
+                     array.find_field(1).to_a,
+                   ])
+    end
+
+    test("is_valids") do
+      @builder.append_values([
+                               [true, 1],
+                               [false, 2],
+                               [true, 3],
+                             ],
+                             [
+                               true,
+                               false,
+                               true,
+                             ])
+      array = @builder.finish
+      assert_equal([
+                     [true, nil, true],
+                     [1, nil, 3],
+                   ],
+                   [
+                     array.find_field(0).to_a,
+                     array.find_field(1).to_a,
+                   ])
+    end
+  end
+
+  sub_test_case("#append") do
+    test("backward compatibility") do
+      @builder.append
+      @builder.get_field_builder(0).append(true)
+      @builder.get_field_builder(1).append(1)
+      @builder.append
+      @builder.get_field_builder(0).append(false)
+      @builder.get_field_builder(1).append(2)
+      array = @builder.finish
+      assert_equal([
+                     [true, 1],
+                     [false, 2],
+                   ],
+                   [
+                     array.get_value(0).values,
+                     array.get_value(1).values,
+                   ])
+    end
+  end
+end