red-arrow 6.0.0 → 8.0.0

data/test/test-function.rb

@@ -53,6 +53,14 @@ class FunctionTest < Test::Unit::TestCase
                    or_function.execute(args).value.to_a)
     end
 
+    test("Arrow::Column") do
+      or_function = Arrow::Function.find("or")
+      table = Arrow::Table.new(a: [true, false, false],
+                               b: [true, false, true])
+      assert_equal([true, false, true],
+                   or_function.execute([table.a, table.b]).value.to_a)
+    end
+
     test("Arrow::Scalar") do
       add_function = Arrow::Function.find("add")
       args = [
@@ -116,12 +124,13 @@ class FunctionTest < Test::Unit::TestCase
       cast_function = Arrow::Function.find("cast")
       date = Date.new(2021, 6, 12)
       args = [date]
-      options = Arrow::CastOptions.new
-      options.to_data_type = Arrow::TimestampDataType.new(:second)
+      options = {
+        to_data_type: Arrow::TimestampDataType.new(:second),
+      }
       time = Time.utc(date.year,
                       date.month,
                       date.day)
-      assert_equal(Arrow::TimestampScalar.new(options.to_data_type,
+      assert_equal(Arrow::TimestampScalar.new(options[:to_data_type],
                                               time.to_i),
                    cast_function.execute(args, options).value)
     end
@@ -132,9 +141,10 @@ class FunctionTest < Test::Unit::TestCase
                                    # 00:10:00
                                    60 * 10)
       args = [arrow_time]
-      options = Arrow::CastOptions.new
-      options.to_data_type = Arrow::Time64DataType.new(:micro)
-      assert_equal(Arrow::Time64Scalar.new(options.to_data_type,
+      options = {
+        to_data_type: Arrow::Time64DataType.new(:micro),
+      }
+      assert_equal(Arrow::Time64Scalar.new(options[:to_data_type],
                                            # 00:10:00.000000
                                            60 * 10 * 1000 * 1000),
                    cast_function.execute(args, options).value)
@@ -146,10 +156,11 @@ class FunctionTest < Test::Unit::TestCase
                                    # 00:10:00.000000
                                    60 * 10 * 1000 * 1000)
       args = [arrow_time]
-      options = Arrow::CastOptions.new
-      options.to_data_type = Arrow::Time32DataType.new(:second)
-      options.allow_time_truncate = true
-      assert_equal(Arrow::Time32Scalar.new(options.to_data_type,
+      options = {
+        to_data_type: Arrow::Time32DataType.new(:second),
+        allow_time_truncate: true,
+      }
+      assert_equal(Arrow::Time32Scalar.new(options[:to_data_type],
                                            # 00:10:00
                                            60 * 10),
                    cast_function.execute(args, options).value)
@@ -159,18 +170,41 @@ class FunctionTest < Test::Unit::TestCase
       cast_function = Arrow::Function.find("cast")
       time = Time.utc(2021, 6, 12, 1, 2, 3, 1)
       args = [time]
-      options = Arrow::CastOptions.new
-      options.to_data_type = Arrow::TimestampDataType.new(:second)
-      options.allow_time_truncate = true
+      options = {
+        to_data_type: Arrow::TimestampDataType.new(:second),
+        allow_time_truncate: true,
+      }
       time = Time.utc(time.year,
                       time.month,
                       time.day,
                       time.hour,
                       time.min,
                       time.sec)
-      assert_equal(Arrow::TimestampScalar.new(options.to_data_type,
+      assert_equal(Arrow::TimestampScalar.new(options[:to_data_type],
                                               time.to_i),
                    cast_function.execute(args, options).value)
     end
+
+    test("SetLookupOptions") do
+      is_in_function = Arrow::Function.find("is_in")
+      args = [
+        Arrow::Int16Array.new([1, 0, 1, 2]),
+      ]
+      options = {
+        value_set: Arrow::Int16Array.new([2, 0]),
+      }
+      assert_equal(Arrow::BooleanArray.new([false, true, false, true]),
+                   is_in_function.execute(args, options).value)
+    end
+  end
+
+  def test_call
+      or_function = Arrow::Function.find("or")
+      args = [
+        Arrow::BooleanArray.new([true, false, false]),
+        Arrow::BooleanArray.new([true, false, true]),
+      ]
+      assert_equal([true, false, true],
+                   or_function.call(args).value.to_a)
   end
 end

data/test/test-table.rb

@@ -186,7 +186,12 @@ class TableTest < Test::Unit::TestCase
     end
 
     test("{key: Range}: beginless include end") do
-      assert_equal(<<-TABLE, @table.slice(count: ..8).to_s)
+      begin
+        range = eval("..8")
+      rescue SyntaxError
+        omit("beginless range isn't supported")
+      end
+      assert_equal(<<-TABLE, @table.slice(count: range).to_s)
 	count	visible
 0	    1	true   
 1	    2	false  
@@ -196,7 +201,12 @@ class TableTest < Test::Unit::TestCase
     end
 
     test("{key: Range}: beginless exclude end") do
-      assert_equal(<<-TABLE, @table.slice(count: ...8).to_s)
+      begin
+        range = eval("...8")
+      rescue SyntaxError
+        omit("beginless range isn't supported")
+      end
+      assert_equal(<<-TABLE, @table.slice(count: range).to_s)
 	count	visible
 0	    1	true   
 1	    2	false  
@@ -205,7 +215,12 @@ class TableTest < Test::Unit::TestCase
     end
 
     test("{key: Range}: endless") do
-      assert_equal(<<-TABLE, @table.slice(count: 16..).to_s)
+      begin
+        range = eval("16..")
+      rescue SyntaxError
+        omit("endless range isn't supported")
+      end
+      assert_equal(<<-TABLE, @table.slice(count: range).to_s)
 	count	visible
 0	   16	true   
 1	   32	false  
@@ -573,14 +588,20 @@ class TableTest < Test::Unit::TestCase
         assert_equal(@table, Arrow::Table.load(output, format: :batch))
       end
 
+      def test_arrows
+        output = create_output(".arrows")
+        @table.save(output, format: :arrows)
+        assert_equal(@table, Arrow::Table.load(output, format: :arrows))
+      end
+
       def test_arrow_streaming
-        output = create_output(".arrow")
+        output = create_output(".arrows")
         @table.save(output, format: :arrow_streaming)
         assert_equal(@table, Arrow::Table.load(output, format: :arrow_streaming))
       end
 
       def test_stream
-        output = create_output(".arrow")
+        output = create_output(".arrows")
         @table.save(output, format: :stream)
         assert_equal(@table, Arrow::Table.load(output, format: :stream))
       end
@@ -626,6 +647,24 @@ class TableTest < Test::Unit::TestCase
         end
 
         sub_test_case("save: auto detect") do
+          test("arrow") do
+            output = create_output(".arrow")
+            @table.save(output)
+            assert_equal(@table,
+                         Arrow::Table.load(output,
+                                           format: :arrow,
+                                           schema: @table.schema))
+          end
+
+          test("arrows") do
+            output = create_output(".arrows")
+            @table.save(output)
+            assert_equal(@table,
+                         Arrow::Table.load(output,
+                                           format: :arrows,
+                                           schema: @table.schema))
+          end
+
           test("csv") do
             output = create_output(".csv")
             @table.save(output)
@@ -664,7 +703,13 @@ class TableTest < Test::Unit::TestCase
 
           test("arrow: streaming") do
             output = create_output(".arrow")
-            @table.save(output, format: :arrow_streaming)
+            @table.save(output, format: :arrows)
+            assert_equal(@table, Arrow::Table.load(output))
+          end
+
+          test("arrows") do
+            output = create_output(".arrows")
+            @table.save(output, format: :arrows)
             assert_equal(@table, Arrow::Table.load(output))
           end
 
@@ -728,6 +773,47 @@ chris\t-1
         end
       end
     end
+
+    sub_test_case("URI") do
+      def start_web_server(path, data, content_type)
+        http_server = WEBrick::HTTPServer.new(:Port => 0)
+        http_server.mount_proc(path) do |request, response|
+          response.body = data
+          response.content_type = content_type
+        end
+        http_server_thread = Thread.new do
+          http_server.start
+        end
+        begin
+          Timeout.timeout(1) do
+            yield(http_server[:Port])
+          end
+        ensure
+          http_server.shutdown
+          http_server_thread.join
+        end
+      end
+
+      data("Arrow File",
+           ["arrow", "application/vnd.apache.arrow.file"])
+      data("Arrow Stream",
+           ["arrows", "application/vnd.apache.arrow.stream"])
+      data("CSV",
+           ["csv", "text/csv"])
+      def test_http(data)
+        extension, content_type = data
+        output = Arrow::ResizableBuffer.new(1024)
+        @table.save(output, format: extension.to_sym)
+        path = "/data.#{extension}"
+        start_web_server(path,
+                         output.data.to_s,
+                         content_type) do |port|
+          input = URI("http://127.0.0.1:#{port}#{path}")
+          loaded_table = Arrow::Table.load(input)
+          assert_equal(@table.to_s, loaded_table.to_s)
+        end
+      end
+    end
   end
 
   test("#pack") do
@@ -922,4 +1008,116 @@ visible: false
       TABLE
     end
   end
+
+  sub_test_case("#join") do
+    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: Symbol") 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],
+                                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", [1, 2]],
+                                      ["key2", [100, 20]],
+                                      ["number", [1100, 2020]],
+                                      ["key1", [1, 2]],
+                                      ["key2", [100, 20]],
+                                      ["string", ["1-100", "2-20"]],
+                                    ]),
+                   table1.join(table2, ["key1", :key2]))
+    end
+
+    test("keys: {left: String, right: Symbol}") do
+      table1 = Arrow::Table.new(left_key: [1, 2, 3],
+                                number: [10, 20, 30])
+      table2 = Arrow::Table.new(right_key: [3, 1],
+                                string: ["three", "one"])
+      assert_equal(Arrow::Table.new([
+                                      ["left_key", [1, 3]],
+                                      ["number", [10, 30]],
+                                      ["right_key", [1, 3]],
+                                      ["string", ["one", "three"]],
+                                    ]),
+                   table1.join(table2, {left: "left_key", right: :right_key}))
+    end
+
+    test("keys: {left: [String, Symbol], right: [Symbol, String]}") do
+      table1 = Arrow::Table.new(left_key1: [1, 1, 2, 2],
+                                left_key2: [10, 100, 20, 200],
+                                number: [1010, 1100, 2020, 2200])
+      table2 = Arrow::Table.new(right_key1: [1, 2, 2],
+                                right_key2: [100, 20, 50],
+                                string: ["1-100", "2-20", "2-50"])
+      assert_equal(Arrow::Table.new([
+                                      ["left_key1", [1, 2]],
+                                      ["left_key2", [100, 20]],
+                                      ["number", [1100, 2020]],
+                                      ["right_key1", [1, 2]],
+                                      ["right_key2", [100, 20]],
+                                      ["string", ["1-100", "2-20"]],
+                                    ]),
+                   table1.join(table2,
+                               {
+                                 left: ["left_key1", :left_key2],
+                                 right: [:right_key1, "right_key2"],
+                               }))
+    end
+
+    test("type:") 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, "key", type: :left_outer))
+    end
+
+    test("left_outputs: & right_outputs:") 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",
+                               left_outputs: ["key", "number"],
+                               right_outputs: ["string"]))
+    end
+  end
 end

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

@@ -276,6 +276,36 @@ module ValuesBasicArraysTests
     target = build(Arrow::Decimal256Array.new(data_type, values))
     assert_equal(values, target.values)
   end
+
+  def test_month_interval
+    values = [
+      1,
+      nil,
+      12,
+    ]
+    target = build(Arrow::MonthIntervalArray.new(values))
+    assert_equal(values, target.values)
+  end
+
+  def test_day_time_interval
+    values = [
+      {day: 1, millisecond: 100},
+      nil,
+      {day: 2, millisecond: 300},
+    ]
+    target = build(Arrow::DayTimeIntervalArray.new(values))
+    assert_equal(values, target.values)
+  end
+
+  def test_month_day_nano_interval
+    values = [
+      {month: 1, day: 1, nanosecond: 100},
+      nil,
+      {month: 2, day: 3, nanosecond: 400},
+    ]
+    target = build(Arrow::MonthDayNanoIntervalArray.new(values))
+    assert_equal(values, target.values)
+  end
 end
 
 class ValuesArrayBasicArraysTest < Test::Unit::TestCase

data/test/values/test-dense-union-array.rb

@@ -347,6 +347,33 @@ module ValuesDenseUnionArrayTests
     assert_equal(values, target.values)
   end
 
+  def test_month_interval
+    values = [
+      {"0" => 1},
+      {"1" => nil},
+    ]
+    target = build(:month_interval, values)
+    assert_equal(values, target.values)
+  end
+
+  def test_day_time_interval
+    values = [
+      {"0" => {day: 1, millisecond: 100}},
+      {"1" => nil},
+    ]
+    target = build(:day_time_interval, values)
+    assert_equal(values, target.values)
+  end
+
+  def test_month_day_nano_interval
+    values = [
+      {"0" => {month: 1, day: 1, nanosecond: 100}},
+      {"1" => nil},
+    ]
+    target = build(:month_day_nano_interval, values)
+    assert_equal(values, target.values)
+  end
+
   def test_list
     values = [
       {"0" => [true, nil, false]},