red-arrow 10.0.0 → 16.0.0

data/test/test-table.rb

@@ -41,8 +41,25 @@ class TableTest < Test::Unit::TestCase
   end
 
   sub_test_case(".new") do
+    test("{Symbol: Arrow::Array}") do
+      schema = Arrow::Schema.new(numbers: :int64)
+      assert_equal(Arrow::Table.new(schema,
+                                    [Arrow::Int64Array.new([1, 2, 3])]),
+                   Arrow::Table.new(numbers: Arrow::Int64Array.new([1, 2, 3])))
+    end
+
+    test("{Symbol: Arrow::ChunkedArray}") do
+      chunked_array = Arrow::ChunkedArray.new([Arrow::Int64Array.new([1, 2, 3])])
+      schema = Arrow::Schema.new(numbers: :int64)
+      assert_equal(Arrow::Table.new(schema,
+                                    [Arrow::Int64Array.new([1, 2, 3])]),
+                   Arrow::Table.new(numbers: chunked_array))
+    end
+
     test("{Symbol: Arrow::Tensor}") do
-      assert_equal(Arrow::Table.new(numbers: Arrow::UInt8Array.new([1, 2, 3])),
+      schema = Arrow::Schema.new(numbers: :uint8)
+      assert_equal(Arrow::Table.new(schema,
+                                    [Arrow::UInt8Array.new([1, 2, 3])]),
                    Arrow::Table.new(numbers: Arrow::Tensor.new([1, 2, 3])))
     end
 
@@ -51,7 +68,8 @@ class TableTest < Test::Unit::TestCase
       def array_like.to_ary
         [1, 2, 3]
       end
-      assert_equal(Arrow::Table.new(numbers: Arrow::UInt8Array.new([1, 2, 3])),
+      schema = Arrow::Schema.new(numbers: :uint8)
+      assert_equal(Arrow::Table.new(schema, [Arrow::UInt8Array.new([1, 2, 3])]),
                    Arrow::Table.new(numbers: array_like))
     end
   end
@@ -69,26 +87,24 @@ class TableTest < Test::Unit::TestCase
       target_rows_raw = [nil, true, true, false, true, false, true, true]
       target_rows = Arrow::BooleanArray.new(target_rows_raw)
       assert_equal(<<-TABLE, @table.slice(target_rows).to_s)
-	 count	visible
-0	(null)	 (null)
-1	     2	false  
-2	     4	 (null)
-3	    16	true   
-4	    64	 (null)
-5	   128	 (null)
+	count	visible
+0	    2	false  
+1	    4	 (null)
+2	   16	true   
+3	   64	 (null)
+4	  128	 (null)
       TABLE
     end
 
     test("Array: boolean") do
       target_rows_raw = [nil, true, true, false, true, false, true, true]
       assert_equal(<<-TABLE, @table.slice(target_rows_raw).to_s)
-	 count	visible
-0	(null)	 (null)
-1	     2	false  
-2	     4	 (null)
-3	    16	true   
-4	    64	 (null)
-5	   128	 (null)
+	count	visible
+0	    2	false  
+1	    4	 (null)
+2	   16	true   
+3	   64	 (null)
+4	  128	 (null)
       TABLE
     end
 
@@ -180,24 +196,18 @@ class TableTest < Test::Unit::TestCase
 
     test("{key: true}") do
       assert_equal(<<-TABLE, @table.slice(visible: true).to_s)
-	 count	visible
-0	     1	true   
-1	(null)	 (null)
-2	     8	true   
-3	    16	true   
-4	(null)	 (null)
-5	(null)	 (null)
+	count	visible
+0	    1	true   
+1	    8	true   
+2	   16	true   
       TABLE
     end
 
     test("{key: false}") do
       assert_equal(<<-TABLE, @table.slice(visible: false).to_s)
-	 count	visible
-0	     2	false  
-1	(null)	 (null)
-2	    32	false  
-3	(null)	 (null)
-4	(null)	 (null)
+	count	visible
+0	    2	false  
+1	   32	false  
       TABLE
     end
 
@@ -268,11 +278,8 @@ class TableTest < Test::Unit::TestCase
 
     test("{key1: Range, key2: true}") do
       assert_equal(<<-TABLE, @table.slice(count: 0..8, visible: false).to_s)
-	 count	visible
-0	     2	false  
-1	(null)	 (null)
-2	(null)	 (null)
-3	(null)	 (null)
+	count	visible
+0	    2	false  
       TABLE
     end
 
@@ -582,6 +589,25 @@ class TableTest < Test::Unit::TestCase
 0	1
       TABLE
     end
+
+    test("empty result") do
+      selected_table = @table.filter([false] * @table.size).select_columns(:a)
+      assert_equal(<<-TABLE, selected_table.to_s)
+	a
+      TABLE
+    end
+  end
+
+  sub_test_case("#column_names") do
+    test("unique") do
+      table = Arrow::Table.new(a: [1], b: [2], c: [3])
+      assert_equal(%w[a b c], table.column_names)
+    end
+
+    test("duplicated") do
+      table = Arrow::Table.new([["a", [1, 2, 3]], ["a", [4, 5, 6]]])
+      assert_equal(%w[a a], table.column_names)
+    end
   end
 
   sub_test_case("#save and .load") do
@@ -651,6 +677,31 @@ class TableTest < Test::Unit::TestCase
                                        format: :tsv,
                                        schema: @table.schema))
       end
+
+      def test_json
+        output = create_output(".json")
+        # TODO: Implement this.
+        # @table.save(output, format: :json)
+        columns = ""
+        @table.each_record.each do |record|
+          column = {
+            "count" => record.count,
+            "visible" => record.visible,
+          }
+          columns << column.to_json
+          columns << "\n"
+        end
+        if output.is_a?(String)
+          File.write(output, columns)
+        else
+          output.resize(columns.bytesize)
+          output.set_data(0, columns)
+        end
+        assert_equal(@table,
+                     Arrow::Table.load(output,
+                                       format: :json,
+                                       schema: @table.schema))
+      end
     end
 
     sub_test_case("path") do
@@ -690,6 +741,11 @@ class TableTest < Test::Unit::TestCase
                                            schema: @table.schema))
           end
 
+          test("csv, return value") do
+            output = create_output(".csv")
+            assert_equal(@table, @table.save(output))
+          end
+
           test("csv.gz") do
             output = create_output(".csv.gz")
             @table.save(output)
@@ -830,6 +886,76 @@ chris\t-1
         end
       end
     end
+
+    sub_test_case("GC") do
+      def setup
+        table = Arrow::Table.new(integer: [1, 2, 3],
+                                 string: ["a", "b", "c"])
+        @buffer = Arrow::ResizableBuffer.new(1024)
+        table.save(@buffer, format: :arrow)
+        @loaded_table = Arrow::Table.load(@buffer)
+      end
+
+      def test_chunked_array
+        chunked_array = @loaded_table[0].data
+        assert_equal(@buffer,
+                     chunked_array.instance_variable_get(:@input).buffer)
+      end
+
+      def test_array
+        array = @loaded_table[0].data.chunks[0]
+        assert_equal(@buffer,
+                     array.instance_variable_get(:@input).buffer)
+      end
+
+      def test_record_batch
+        record_batch = @loaded_table.each_record_batch.first
+        assert_equal(@buffer,
+                     record_batch.instance_variable_get(:@input).buffer)
+      end
+
+      def test_record_batch_array
+        array = @loaded_table.each_record_batch.first[0].data
+        assert_equal(@buffer,
+                     array.instance_variable_get(:@input).buffer)
+      end
+
+      def test_record_batch_table
+        table = @loaded_table.each_record_batch.first.to_table
+        assert_equal(@buffer,
+                     table.instance_variable_get(:@input).buffer)
+      end
+
+      def test_slice
+        table = @loaded_table.slice(0..-1)
+        assert_equal(@buffer,
+                     table.instance_variable_get(:@input).buffer)
+      end
+
+      def test_merge
+        table = @loaded_table.merge({})
+        assert_equal(@buffer,
+                     table.instance_variable_get(:@input).buffer)
+      end
+
+      def test_remove_column
+        table = @loaded_table.remove_column(0)
+        assert_equal(@buffer,
+                     table.instance_variable_get(:@input).buffer)
+      end
+
+      def test_pack
+        table = @loaded_table.pack
+        assert_equal(@buffer,
+                     table.instance_variable_get(:@input).buffer)
+      end
+
+      def test_join
+        table = @loaded_table.join(@loaded_table, :integer)
+        assert_equal(@buffer,
+                     table.instance_variable_get(:@input).buffer)
+      end
+    end
   end
 
   test("#pack") do
@@ -1026,7 +1152,7 @@ visible: false
   end
 
   sub_test_case("#join") do
-    test("keys: String") do
+    test("keys: nil (natural join)") do
       table1 = Arrow::Table.new(key: [1, 2, 3],
                                 number: [10, 20, 30])
       table2 = Arrow::Table.new(key: [3, 1],
@@ -1034,7 +1160,19 @@ visible: false
       assert_equal(Arrow::Table.new([
                                       ["key", [1, 3]],
                                       ["number", [10, 30]],
+                                      ["string", ["one", "three"]],
+                                    ]),
+                   table1.join(table2))
+    end
+
+    test("keys: String") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
                                       ["key", [1, 3]],
+                                      ["number", [10, 30]],
                                       ["string", ["one", "three"]],
                                     ]),
                    table1.join(table2, "key"))
@@ -1048,12 +1186,25 @@ visible: false
       assert_equal(Arrow::Table.new([
                                       ["key", [1, 3]],
                                       ["number", [10, 30]],
-                                      ["key", [1, 3]],
                                       ["string", ["one", "three"]],
                                     ]),
                    table1.join(table2, :key))
     end
 
+    test("keys: [String]") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["key", [1, 3]],
+                                      ["number", [10, 30]],
+                                      ["key", [1, 3]],
+                                      ["string", ["one", "three"]],
+                                    ]),
+                   table1.join(table2, ["key"]))
+    end
+
     test("keys: [String, Symbol]") do
       table1 = Arrow::Table.new(key1: [1, 1, 2, 2],
                                 key2: [10, 100, 20, 200],
@@ -1083,7 +1234,9 @@ visible: false
                                       ["right_key", [1, 3]],
                                       ["string", ["one", "three"]],
                                     ]),
-                   table1.join(table2, {left: "left_key", right: :right_key}))
+                   table1.join(table2,
+                               {left: "left_key", right: :right_key},
+                               type: :inner))
     end
 
     test("keys: {left: [String, Symbol], right: [Symbol, String]}") do
@@ -1105,10 +1258,11 @@ visible: false
                                {
                                  left: ["left_key1", :left_key2],
                                  right: [:right_key1, "right_key2"],
-                               }))
+                               },
+                               type: :inner))
     end
 
-    test("type:") do
+    test("type: :left_outer") do
       table1 = Arrow::Table.new(key: [1, 2, 3],
                                 number: [10, 20, 30])
       table2 = Arrow::Table.new(key: [3, 1],
@@ -1116,12 +1270,85 @@ visible: false
       assert_equal(Arrow::Table.new([
                                       ["key", [1, 3, 2]],
                                       ["number", [10, 30, 20]],
-                                      ["key", [1, 3, nil]],
                                       ["string", ["one", "three", nil]],
                                     ]),
                    table1.join(table2, "key", type: :left_outer))
     end
 
+    test("type: :right_outer") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["key", [1, 3]],
+                                      ["number", [10, 30]],
+                                      ["string", ["one", "three"]],
+                                    ]),
+                   table1.join(table2, "key", type: :right_outer))
+    end
+
+    test("type: :full_outer") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["key", [1, 3, 2]],
+                                      ["number", [10, 30, 20]],
+                                      ["string", ["one", "three", nil]],
+                                    ]),
+                   table1.join(table2, "key", type: :full_outer))
+    end
+
+    test("type: :left_semi") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["key", [1, 3]],
+                                      ["number", [10, 30]],
+                                    ]),
+                   table1.join(table2, "key", type: :left_semi))
+    end
+
+    test("type: :right_semi") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["key", [3, 1]],
+                                      ["string", ["three", "one"]],
+                                    ]),
+                   table1.join(table2, "key", type: :right_semi))
+    end
+
+    test("type: :left_anti") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["key", [2]],
+                                      ["number", [20]],
+                                    ]),
+                   table1.join(table2, "key", type: :left_anti))
+    end
+
+    test("type: :right_anti") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["key", Arrow::ChunkedArray.new(:uint8)],
+                                      ["string", Arrow::ChunkedArray.new(:string)],
+                                    ]),
+                   table1.join(table2, "key", type: :right_anti))
+    end
+
     test("left_outputs: & right_outputs:") do
       table1 = Arrow::Table.new(key: [1, 2, 3],
                                 number: [10, 20, 30])
@@ -1135,5 +1362,94 @@ visible: false
                                left_outputs: ["key", "number"],
                                right_outputs: ["string"]))
     end
+
+    test("left_outputs: & type: :inner") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["key", [1, 3]],
+                                      ["number", [10, 30]],
+                                      ["key", [1, 3]],
+                                      ["string", ["one", "three"]]
+                                    ]),
+                   table1.join(table2,
+                               type: :inner,
+                               left_outputs: table1.column_names,
+                               right_outputs: table2.column_names))
+    end
+
+    test("left_outputs: & type: :left_outer") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["key", [1, 3, 2]],
+                                      ["number", [10, 30, 20]],
+                                      ["key", [1, 3, nil]],
+                                      ["string", ["one", "three", nil]],
+                                    ]),
+                   table1.join(table2,
+                               type: :left_outer,
+                               left_outputs: table1.column_names,
+                               right_outputs: table2.column_names))
+    end
+
+    test("left_outputs: & type: :right_outer") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["key", [1, 3]],
+                                      ["number", [10, 30]],
+                                      ["key", [1, 3]],
+                                      ["string", ["one", "three"]],
+                                    ]),
+                   table1.join(table2,
+                               type: :right_outer,
+                               left_outputs: table1.column_names,
+                               right_outputs: table2.column_names))
+    end
+
+    test("left_outputs: & type: :full_outer") do
+      table1 = Arrow::Table.new(key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["key", [1, 3, 2]],
+                                      ["number", [10, 30, 20]],
+                                      ["key", [1, 3, nil]],
+                                      ["string", ["one", "three", nil]],
+                                    ]),
+                   table1.join(table2,
+                               type: :full_outer,
+                               left_outputs: table1.column_names,
+                               right_outputs: table2.column_names))
+    end
+
+    test("left_suffix: & keys: [String]") do
+      table1 = Arrow::Table.new(key1: [1, 1, 2, 2],
+                                key2: [10, 100, 20, 200],
+                                number: [1010, 1100, 2020, 2200])
+      table2 = Arrow::Table.new(key1: [1, 2, 2],
+                                key2: [100, 20, 50],
+                                string: ["1-100", "2-20", "2-50"])
+      assert_equal(Arrow::Table.new([
+                                      ["key1_left", [1, 2]],
+                                      ["key2_left", [100, 20]],
+                                      ["number", [1100, 2020]],
+                                      ["key1_right", [1, 2]],
+                                      ["key2_right", [100, 20]],
+                                      ["string", ["1-100", "2-20"]],
+                                    ]),
+                    table1.join(table2,
+                                ["key1", "key2"],
+                                left_suffix: "_left",
+                                right_suffix: "_right"))
+    end
   end
 end

data/test/values/test-basic-arrays.rb

@@ -107,6 +107,16 @@ module ValuesBasicArraysTests
     assert_equal(values, target.values)
   end
 
+  def test_half_float
+    values = [
+      -1.5,
+      nil,
+      1.5,
+    ]
+    target = build(Arrow::HalfFloatArray.new(values))
+    assert_equal(values, target.values)
+  end
+
   def test_float
     values = [
       -1.0,