diff options
Diffstat (limited to '')
| -rw-r--r-- | src/arrow/ruby/red-arrow/test/test-chunked-array.rb | 183 |
1 files changed, 183 insertions, 0 deletions
diff --git a/src/arrow/ruby/red-arrow/test/test-chunked-array.rb b/src/arrow/ruby/red-arrow/test/test-chunked-array.rb new file mode 100644 index 000000000..3785e9868 --- /dev/null +++ b/src/arrow/ruby/red-arrow/test/test-chunked-array.rb @@ -0,0 +1,183 @@ +# 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 ChunkedArrayTest < Test::Unit::TestCase + test("#each") do + arrays = [ + Arrow::BooleanArray.new([true, false]), + Arrow::BooleanArray.new([nil, true]), + ] + chunked_array = Arrow::ChunkedArray.new(arrays) + assert_equal([true, false, nil, true], + chunked_array.to_a) + end + + sub_test_case("#pack") do + test("basic array") do + arrays = [ + Arrow::BooleanArray.new([true, false]), + Arrow::BooleanArray.new([nil, true]), + ] + chunked_array = Arrow::ChunkedArray.new(arrays) + packed_chunked_array = chunked_array.pack + assert_equal([ + Arrow::BooleanArray, + [true, false, nil, true], + ], + [ + packed_chunked_array.class, + packed_chunked_array.to_a, + ]) + end + + test("TimestampArray") do + type = Arrow::TimestampDataType.new(:nano) + arrays = [ + Arrow::TimestampArrayBuilder.new(type).build([Time.at(0)]), + Arrow::TimestampArrayBuilder.new(type).build([Time.at(1)]), + ] + chunked_array = Arrow::ChunkedArray.new(arrays) + packed_chunked_array = chunked_array.pack + assert_equal([ + Arrow::TimestampArray, + [Time.at(0), Time.at(1)], + ], + [ + packed_chunked_array.class, + packed_chunked_array.to_a, + ]) + end + end + + sub_test_case("#==") do + def setup + arrays = [ + Arrow::BooleanArray.new([true]), + Arrow::BooleanArray.new([false, true]), + ] + @chunked_array = Arrow::ChunkedArray.new(arrays) + end + + test("Arrow::ChunkedArray") do + assert do + @chunked_array == @chunked_array + end + end + + test("not Arrow::ChunkedArray") do + 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) + @options = Arrow::FilterOptions.new + @options.null_selection_behavior = :emit_null + 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, @options)) + 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, @options)) + 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, @options)) + 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 |