red-arrow 5.0.0 → 6.0.0

data/test/test-group.rb

@@ -42,9 +42,9 @@ class GroupTest < Test::Unit::TestCase
       }
       table = Arrow::Table.new(raw_table)
       assert_equal(<<-TABLE, table.group(:time).count.to_s)
-	                     time	int
-0	#{time_values[0].iso8601}	  1
-1	#{time_values[1].iso8601}	  1
+	count(int)	                     time
+0	         1	#{time_values[0].iso8601}
+1	         1	#{time_values[1].iso8601}
       TABLE
     end
   end
@@ -52,20 +52,31 @@ class GroupTest < Test::Unit::TestCase
   sub_test_case("#count") do
     test("single") do
       assert_equal(<<-TABLE, @table.group(:group_key1).count.to_s)
-	group_key1	group_key2	int	uint	float	string
-0	         1	         2	  2	   1	    1	     2
-1	         2	         1	  0	   1	    1	     1
-2	         3	         3	  3	   3	    3	     2
+	count(group_key2)	count(int)	count(uint)	count(float)	count(string)	group_key1
+0	                2	         2	          1	           1	            2	         1
+1	                1	         0	          1	           1	            1	         2
+2	                3	         3	          3	           3	            2	         3
       TABLE
     end
 
     test("multiple") do
       assert_equal(<<-TABLE, @table.group(:group_key1, :group_key2).count.to_s)
-	group_key1	group_key2	int	uint	float	string
-0	         1	         1	  2	   1	    1	     2
-1	         2	         1	  0	   1	    1	     1
-2	         3	         1	  1	   1	    1	     0
-3	         3	         2	  2	   2	    2	     2
+	count(int)	count(uint)	count(float)	count(string)	group_key1	group_key2
+0	         2	          1	           1	            2	         1	         1
+1	         0	          1	           1	            1	         2	         1
+2	         1	          1	           1	            0	         3	         1
+3	         2	          2	           2	            2	         3	         2
+      TABLE
+    end
+
+    test("column") do
+      group = @table.group(:group_key1, :group_key2)
+      assert_equal(<<-TABLE, group.count(:int, :uint).to_s)
+	count(int)	count(uint)	group_key1	group_key2
+0	         2	          1	         1	         1
+1	         0	          1	         2	         1
+2	         1	          1	         3	         1
+3	         2	          2	         3	         2
       TABLE
     end
   end
@@ -73,41 +84,41 @@ class GroupTest < Test::Unit::TestCase
   sub_test_case("#sum") do
     test("single") do
       assert_equal(<<-TABLE, @table.group(:group_key1).sum.to_s)
-	group_key1	group_key2	int	uint	     float
-0	         1	         2	 -3	   1	  2.200000
-1	         2	         1	  0	   3	  3.300000
-2	         3	         5	-15	  15	 16.500000
+	sum(group_key2)	sum(int)	sum(uint)	sum(float)	group_key1
+0	              2	      -3	        1	  2.200000	         1
+1	              1	  (null)	        3	  3.300000	         2
+2	              5	     -15	       15	 16.500000	         3
       TABLE
     end
 
     test("multiple") do
       assert_equal(<<-TABLE, @table.group(:group_key1, :group_key2).sum.to_s)
-	group_key1	group_key2	int	uint	     float
-0	         1	         1	 -3	   1	  2.200000
-1	         2	         1	  0	   3	  3.300000
-2	         3	         1	 -4	   4	  4.400000
-3	         3	         2	-11	  11	 12.100000
+	sum(int)	sum(uint)	sum(float)	group_key1	group_key2
+0	      -3	        1	  2.200000	         1	         1
+1	  (null)	        3	  3.300000	         2	         1
+2	      -4	        4	  4.400000	         3	         1
+3	     -11	       11	 12.100000	         3	         2
       TABLE
     end
   end
 
-  sub_test_case("#average") do
+  sub_test_case("#mean") do
     test("single") do
-      assert_equal(<<-TABLE, @table.group(:group_key1).average.to_s)
-	group_key1	group_key2	       int	      uint	     float
-0	         1	  1.000000	 -1.500000	  1.000000	  2.200000
-1	         2	  1.000000	  0.000000	  3.000000	  3.300000
-2	         3	  1.666667	 -5.000000	  5.000000	  5.500000
+      assert_equal(<<-TABLE, @table.group(:group_key1).mean.to_s)
+	mean(group_key2)	 mean(int)	mean(uint)	mean(float)	group_key1
+0	        1.000000	 -1.500000	  1.000000	   2.200000	         1
+1	        1.000000	    (null)	  3.000000	   3.300000	         2
+2	        1.666667	 -5.000000	  5.000000	   5.500000	         3
       TABLE
     end
 
     test("multiple") do
-      assert_equal(<<-TABLE, @table.group(:group_key1, :group_key2).average.to_s)
-	group_key1	group_key2	       int	      uint	     float
-0	         1	         1	 -1.500000	  1.000000	  2.200000
-1	         2	         1	  0.000000	  3.000000	  3.300000
-2	         3	         1	 -4.000000	  4.000000	  4.400000
-3	         3	         2	 -5.500000	  5.500000	  6.050000
+      assert_equal(<<-TABLE, @table.group(:group_key1, :group_key2).mean.to_s)
+	 mean(int)	mean(uint)	mean(float)	group_key1	group_key2
+0	 -1.500000	  1.000000	   2.200000	         1	         1
+1	    (null)	  3.000000	   3.300000	         2	         1
+2	 -4.000000	  4.000000	   4.400000	         3	         1
+3	 -5.500000	  5.500000	   6.050000	         3	         2
       TABLE
     end
   end
@@ -115,20 +126,20 @@ class GroupTest < Test::Unit::TestCase
   sub_test_case("#min") do
     test("single") do
       assert_equal(<<-TABLE, @table.group(:group_key1).min.to_s)
-	group_key1	group_key2	int	uint	     float
-0	         1	         1	 -2	   1	  2.200000
-1	         2	         1	   	   3	  3.300000
-2	         3	         1	 -6	   4	  4.400000
+	min(group_key2)	min(int)	min(uint)	min(float)	group_key1
+0	              1	      -2	        1	  2.200000	         1
+1	              1	  (null)	        3	  3.300000	         2
+2	              1	      -6	        4	  4.400000	         3
       TABLE
     end
 
     test("multiple") do
       assert_equal(<<-TABLE, @table.group(:group_key1, :group_key2).min.to_s)
-	group_key1	group_key2	int	uint	     float
-0	         1	         1	 -2	   1	  2.200000
-1	         2	         1	   	   3	  3.300000
-2	         3	         1	 -4	   4	  4.400000
-3	         3	         2	 -6	   5	  5.500000
+	min(int)	min(uint)	min(float)	group_key1	group_key2
+0	      -2	        1	  2.200000	         1	         1
+1	  (null)	        3	  3.300000	         2	         1
+2	      -4	        4	  4.400000	         3	         1
+3	      -6	        5	  5.500000	         3	         2
       TABLE
     end
   end
@@ -136,20 +147,33 @@ class GroupTest < Test::Unit::TestCase
   sub_test_case("#max") do
     test("single") do
       assert_equal(<<-TABLE, @table.group(:group_key1).max.to_s)
-	group_key1	group_key2	int	uint	     float
-0	         1	         1	 -1	   1	  2.200000
-1	         2	         1	   	   3	  3.300000
-2	         3	         2	 -4	   6	  6.600000
+	max(group_key2)	max(int)	max(uint)	max(float)	group_key1
+0	              1	      -1	        1	  2.200000	         1
+1	              1	  (null)	        3	  3.300000	         2
+2	              2	      -4	        6	  6.600000	         3
       TABLE
     end
 
     test("multiple") do
       assert_equal(<<-TABLE, @table.group(:group_key1, :group_key2).max.to_s)
-	group_key1	group_key2	int	uint	     float
-0	         1	         1	 -1	   1	  2.200000
-1	         2	         1	   	   3	  3.300000
-2	         3	         1	 -4	   4	  4.400000
-3	         3	         2	 -5	   6	  6.600000
+	max(int)	max(uint)	max(float)	group_key1	group_key2
+0	      -1	        1	  2.200000	         1	         1
+1	  (null)	        3	  3.300000	         2	         1
+2	      -4	        4	  4.400000	         3	         1
+3	      -5	        6	  6.600000	         3	         2
+      TABLE
+    end
+  end
+
+  sub_test_case("#aggregate") do
+    test("function()") do
+      group = @table.group(:group_key1, :group_key2)
+      assert_equal(<<-TABLE, group.aggregate("count(int)", "sum(uint)").to_s)
+	count(int)	sum(uint)	group_key1	group_key2
+0	         2	        1	         1	         1
+1	         0	        3	         2	         1
+2	         1	        4	         3	         1
+3	         2	       11	         3	         2
       TABLE
     end
   end

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

@@ -0,0 +1,110 @@
+# 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 MapArrayBuilderTest < Test::Unit::TestCase
+  def setup
+    key_type = Arrow::StringDataType.new
+    item_type = Arrow::Int16DataType.new
+    data_type = Arrow::MapDataType.new(key_type, item_type)
+    @builder = Arrow::MapArrayBuilder.new(data_type)
+  end
+
+  sub_test_case("#append_value") do
+    test("nil") do
+      @builder.append_value(nil)
+      array = @builder.finish
+      assert_equal([nil], array.collect {|value| value})
+    end
+
+    test("Hash") do
+      @builder.append_value({"a" => 0, "b" => 1})
+      @builder.append_value({"c" => 0, "d" => 1})
+      array = @builder.finish
+      assert_equal([
+                     {"a" => 0, "b" => 1},
+                     {"c" => 0, "d" => 1}
+                   ],
+                   array.collect {|value| value})
+    end
+
+    test("#each") do
+      @builder.append_value([["a", 0], ["b", 1]])
+      @builder.append_value([["c", 0], ["d", 1]])
+      array = @builder.finish
+      assert_equal([
+                     {"a" => 0, "b" => 1},
+                     {"c" => 0, "d" => 1}
+                   ],
+                   array.collect {|value| value})
+    end
+  end
+
+  sub_test_case("#append_values") do
+    test("[nil]") do
+      @builder.append_values([nil])
+      array = @builder.finish
+      assert_equal([nil], array.collect {|value| value})
+    end
+
+    test("[Hash]") do
+      @builder.append_values([{"a" => 0, "b" => 1}, {"c" => 0, "d" => 1}])
+      array = @builder.finish
+      assert_equal([
+                     {"a" => 0, "b" => 1},
+                     {"c" => 0, "d" => 1}
+                   ],
+                   array.collect {|value| value})
+    end
+
+    test("[#each]") do
+      @builder.append_values([[["a", 0], ["b", 1]], [["c", 0], ["d", 1]]])
+      array = @builder.finish
+      assert_equal([
+                     {"a" => 0, "b" => 1},
+                     {"c" => 0, "d" => 1}
+                   ],
+                   array.collect {|value| value})
+    end
+
+    test("[nil, Hash, #each]") do
+      @builder.append_values([nil, {"a" => 0, "b" => 1}, [["c", 0], ["d", 1]]])
+      array = @builder.finish
+      assert_equal([
+                     nil,
+                     {"a" => 0, "b" => 1},
+                     {"c" => 0, "d" => 1}
+                   ],
+                   array.collect {|value| value})
+    end
+
+    test("is_valids") do
+      @builder.append_values([
+                               {"a" => 0, "b" => 1},
+                               {"c" => 0, "d" => 1},
+                               {"e" => 0, "f" => 1}
+                             ],
+                             [true, false, true])
+      array = @builder.finish
+      assert_equal([
+                     {"a" => 0, "b" => 1},
+                     nil,
+                     {"e" => 0, "f" => 1}
+                   ],
+                   array.collect {|value| value})
+    end
+  end
+end

data/test/test-map-array.rb

@@ -0,0 +1,33 @@
+# 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 MapArrayTest < Test::Unit::TestCase
+  sub_test_case(".new") do
+    test("build") do
+      key_type = Arrow::StringDataType.new
+      item_type = Arrow::Int16DataType.new
+      data_type = Arrow::MapDataType.new(key_type, item_type)
+      values = [
+        {"a" => 0, "b" => 1},
+        nil,
+        {"c" => 0, "d" => 1}
+      ]
+      array = Arrow::MapArray.new(data_type, values)
+      assert_equal(values, array.collect {|value| value})
+    end
+  end
+end

data/test/test-map-data-type.rb

@@ -0,0 +1,36 @@
+# 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 MapDataTypeTest < Test::Unit::TestCase
+  sub_test_case(".new") do
+    def setup
+      @key = :int8
+      @item = :string
+    end
+
+    test("ordered arguments") do
+      assert_equal("map<int8, string>",
+                   Arrow::MapDataType.new(@key, @item).to_s)
+    end
+
+    test("description") do
+      assert_equal("map<int8, string>",
+                   Arrow::MapDataType.new(key: @key,
+                                          item: @item).to_s)
+    end
+  end
+end

data/test/test-record-batch-reader.rb

@@ -0,0 +1,46 @@
+# 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 TestRecordBatchReader < Test::Unit::TestCase
+  sub_test_case(".try_convert") do
+    test("Arrow::RecordBatch") do
+      record_batch =
+        Arrow::RecordBatch.new("count" => [1, 2, 3],
+                               "private" => [true, false, true])
+      reader = Arrow::RecordBatchReader.try_convert(record_batch)
+      assert_equal(record_batch,
+                   reader.read_next)
+    end
+
+    test("[Arrow::RecordBatch]") do
+      record_batch =
+        Arrow::RecordBatch.new("count" => [1, 2, 3],
+                               "private" => [true, false, true])
+      reader = Arrow::RecordBatchReader.try_convert([record_batch])
+      assert_equal(record_batch,
+                   reader.read_next)
+    end
+
+    test("Arrow::Table") do
+      table = Arrow::Table.new("count" => [1, 2, 3],
+                               "private" => [true, false, true])
+      reader = Arrow::RecordBatchReader.try_convert(table)
+      assert_equal(table,
+                   reader.read_all)
+    end
+  end
+end

data/test/test-record-batch.rb

@@ -136,5 +136,47 @@ class RecordBatchTest < Test::Unit::TestCase
         end
       end
     end
+
+    sub_test_case("#[]") do
+      def setup
+        @record_batch = Arrow::RecordBatch.new(a: [true],
+                                               b: [true],
+                                               c: [true],
+                                               d: [true],
+                                               e: [true],
+                                               f: [true],
+                                               g: [true])
+      end
+
+      test("[String]") do
+        assert_equal(Arrow::Column.new(@record_batch, 0),
+                     @record_batch["a"])
+      end
+
+      test("[Symbol]") do
+        assert_equal(Arrow::Column.new(@record_batch, 1),
+                     @record_batch[:b])
+      end
+
+      test("[Integer]") do
+        assert_equal(Arrow::Column.new(@record_batch, 6),
+                     @record_batch[-1])
+      end
+
+      test("[Range]") do
+        assert_equal(Arrow::RecordBatch.new(d: [true],
+                                            e: [true]),
+                     @record_batch[3..4])
+      end
+
+      test("[[Symbol, String, Integer, Range]]") do
+        assert_equal(Arrow::RecordBatch.new(c: [true],
+                                            a: [true],
+                                            g: [true],
+                                            d: [true],
+                                            e: [true]),
+                     @record_batch[[:c, "a", -1, 3..4]])
+      end
+    end
   end
 end