test/test-chunked-array.rb in red-arrow-0.15.1 vs test/test-chunked-array.rb in red-arrow-0.16.0

@@ -82,6 +82,100 @@
       assert do
         not (@chunked_array == 29)
       end
     end
   end
+
+  sub_test_case("#filter") do
+    def setup
+      arrays = [
+        Arrow::BooleanArray.new([false, true]),
+        Arrow::BooleanArray.new([false, true, false]),
+      ]
+      @chunked_array = Arrow::ChunkedArray.new(arrays)
+    end
+
+    test("Array: boolean") do
+      filter = [nil, true, true, false, true]
+      chunks = [
+        Arrow::BooleanArray.new([nil, true]),
+        Arrow::BooleanArray.new([false, false]),
+      ]
+      filtered_chunked_array = Arrow::ChunkedArray.new(chunks)
+      assert_equal(filtered_chunked_array,
+                   @chunked_array.filter(filter))
+    end
+
+    test("Arrow::BooleanArray") do
+      filter = Arrow::BooleanArray.new([nil, true, true, false, true])
+      chunks = [
+        Arrow::BooleanArray.new([nil, true]),
+        Arrow::BooleanArray.new([false, false]),
+      ]
+      filtered_chunked_array = Arrow::ChunkedArray.new(chunks)
+      assert_equal(filtered_chunked_array,
+                   @chunked_array.filter(filter))
+    end
+
+    test("Arrow::ChunkedArray") do
+      chunks = [
+        Arrow::BooleanArray.new([nil, true]),
+        Arrow::BooleanArray.new([true, false, true]),
+      ]
+      filter = Arrow::ChunkedArray.new(chunks)
+      filtered_chunks = [
+        Arrow::BooleanArray.new([nil, true]),
+        Arrow::BooleanArray.new([false, false]),
+      ]
+      filtered_chunked_array = Arrow::ChunkedArray.new(filtered_chunks)
+      assert_equal(filtered_chunked_array,
+                   @chunked_array.filter(filter))
+    end
+  end
+
+  sub_test_case("#take") do
+    def setup
+      chunks = [
+        Arrow::Int16Array.new([1, 0]),
+        Arrow::Int16Array.new([2]),
+      ]
+      @chunked_array = Arrow::ChunkedArray.new(chunks)
+    end
+
+    test("Arrow: boolean") do
+      chunks = [
+        Arrow::Int16Array.new([0, 1]),
+        Arrow::Int16Array.new([2])
+      ]
+      taken_chunked_array = Arrow::ChunkedArray.new(chunks)
+      indices = [1, 0, 2]
+      assert_equal(taken_chunked_array,
+                   @chunked_array.take(indices))
+    end
+
+    test("Arrow::Array") do
+      chunks = [
+        Arrow::Int16Array.new([0, 1]),
+        Arrow::Int16Array.new([2])
+      ]
+      taken_chunked_array = Arrow::ChunkedArray.new(chunks)
+      indices = Arrow::Int16Array.new([1, 0, 2])
+      assert_equal(taken_chunked_array,
+                   @chunked_array.take(indices))
+    end
+
+    test("Arrow::ChunkedArray") do
+      taken_chunks = [
+        Arrow::Int16Array.new([0, 1]),
+        Arrow::Int16Array.new([2])
+      ]
+      taken_chunked_array = Arrow::ChunkedArray.new(taken_chunks)
+      indices_chunks = [
+        Arrow::Int16Array.new([1, 0]),
+        Arrow::Int16Array.new([2])
+      ]
+      indices = Arrow::ChunkedArray.new(indices_chunks)
+      assert_equal(taken_chunked_array,
+                   @chunked_array.take(indices))
+    end
+  end
 end