ceph/src/arrow/ruby/red-arrow/test/test-chunked-array.rb

   1 # Licensed to the Apache Software Foundation (ASF) under one
   2 # or more contributor license agreements.  See the NOTICE file
   3 # distributed with this work for additional information
   4 # regarding copyright ownership.  The ASF licenses this file
   5 # to you under the Apache License, Version 2.0 (the
   6 # "License"); you may not use this file except in compliance
   7 # with the License.  You may obtain a copy of the License at
   8 #
   9 #   http://www.apache.org/licenses/LICENSE-2.0
  10 #
  11 # Unless required by applicable law or agreed to in writing,
  12 # software distributed under the License is distributed on an
  13 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14 # KIND, either express or implied.  See the License for the
  15 # specific language governing permissions and limitations
  16 # under the License.
  17
  18 class ChunkedArrayTest < Test::Unit::TestCase
  19   test("#each") do
  20     arrays = [
  21       Arrow::BooleanArray.new([true, false]),
  22       Arrow::BooleanArray.new([nil, true]),
  23     ]
  24     chunked_array = Arrow::ChunkedArray.new(arrays)
  25     assert_equal([true, false, nil, true],
  26                  chunked_array.to_a)
  27   end
  28
  29   sub_test_case("#pack") do
  30     test("basic array") do
  31       arrays = [
  32         Arrow::BooleanArray.new([true, false]),
  33         Arrow::BooleanArray.new([nil, true]),
  34       ]
  35       chunked_array = Arrow::ChunkedArray.new(arrays)
  36       packed_chunked_array = chunked_array.pack
  37       assert_equal([
  38                      Arrow::BooleanArray,
  39                      [true, false, nil, true],
  40                    ],
  41                    [
  42                      packed_chunked_array.class,
  43                      packed_chunked_array.to_a,
  44                    ])
  45     end
  46
  47     test("TimestampArray") do
  48       type = Arrow::TimestampDataType.new(:nano)
  49       arrays = [
  50         Arrow::TimestampArrayBuilder.new(type).build([Time.at(0)]),
  51         Arrow::TimestampArrayBuilder.new(type).build([Time.at(1)]),
  52       ]
  53       chunked_array = Arrow::ChunkedArray.new(arrays)
  54       packed_chunked_array = chunked_array.pack
  55       assert_equal([
  56                      Arrow::TimestampArray,
  57                      [Time.at(0), Time.at(1)],
  58                    ],
  59                    [
  60                      packed_chunked_array.class,
  61                      packed_chunked_array.to_a,
  62                    ])
  63     end
  64   end
  65
  66   sub_test_case("#==") do
  67     def setup
  68       arrays = [
  69         Arrow::BooleanArray.new([true]),
  70         Arrow::BooleanArray.new([false, true]),
  71       ]
  72       @chunked_array = Arrow::ChunkedArray.new(arrays)
  73     end
  74
  75     test("Arrow::ChunkedArray") do
  76       assert do
  77         @chunked_array == @chunked_array
  78       end
  79     end
  80
  81     test("not Arrow::ChunkedArray") do
  82       assert do
  83         not (@chunked_array == 29)
  84       end
  85     end
  86   end
  87
  88   sub_test_case("#filter") do
  89     def setup
  90       arrays = [
  91         Arrow::BooleanArray.new([false, true]),
  92         Arrow::BooleanArray.new([false, true, false]),
  93       ]
  94       @chunked_array = Arrow::ChunkedArray.new(arrays)
  95       @options = Arrow::FilterOptions.new
  96       @options.null_selection_behavior = :emit_null
  97     end
  98
  99     test("Array: boolean") do
 100       filter = [nil, true, true, false, true]
 101       chunks = [
 102         Arrow::BooleanArray.new([nil, true]),
 103         Arrow::BooleanArray.new([false, false]),
 104       ]
 105       filtered_chunked_array = Arrow::ChunkedArray.new(chunks)
 106       assert_equal(filtered_chunked_array,
 107                    @chunked_array.filter(filter, @options))
 108     end
 109
 110     test("Arrow::BooleanArray") do
 111       filter = Arrow::BooleanArray.new([nil, true, true, false, true])
 112       chunks = [
 113         Arrow::BooleanArray.new([nil, true]),
 114         Arrow::BooleanArray.new([false, false]),
 115       ]
 116       filtered_chunked_array = Arrow::ChunkedArray.new(chunks)
 117       assert_equal(filtered_chunked_array,
 118                    @chunked_array.filter(filter, @options))
 119     end
 120
 121     test("Arrow::ChunkedArray") do
 122       chunks = [
 123         Arrow::BooleanArray.new([nil, true]),
 124         Arrow::BooleanArray.new([true, false, true]),
 125       ]
 126       filter = Arrow::ChunkedArray.new(chunks)
 127       filtered_chunks = [
 128         Arrow::BooleanArray.new([nil, true]),
 129         Arrow::BooleanArray.new([false, false]),
 130       ]
 131       filtered_chunked_array = Arrow::ChunkedArray.new(filtered_chunks)
 132       assert_equal(filtered_chunked_array,
 133                    @chunked_array.filter(filter, @options))
 134     end
 135   end
 136
 137   sub_test_case("#take") do
 138     def setup
 139       chunks = [
 140         Arrow::Int16Array.new([1, 0]),
 141         Arrow::Int16Array.new([2]),
 142       ]
 143       @chunked_array = Arrow::ChunkedArray.new(chunks)
 144     end
 145
 146     test("Arrow: boolean") do
 147       chunks = [
 148         Arrow::Int16Array.new([0, 1]),
 149         Arrow::Int16Array.new([2])
 150       ]
 151       taken_chunked_array = Arrow::ChunkedArray.new(chunks)
 152       indices = [1, 0, 2]
 153       assert_equal(taken_chunked_array,
 154                    @chunked_array.take(indices))
 155     end
 156
 157     test("Arrow::Array") do
 158       chunks = [
 159         Arrow::Int16Array.new([0, 1]),
 160         Arrow::Int16Array.new([2])
 161       ]
 162       taken_chunked_array = Arrow::ChunkedArray.new(chunks)
 163       indices = Arrow::Int16Array.new([1, 0, 2])
 164       assert_equal(taken_chunked_array,
 165                    @chunked_array.take(indices))
 166     end
 167
 168     test("Arrow::ChunkedArray") do
 169       taken_chunks = [
 170         Arrow::Int16Array.new([0, 1]),
 171         Arrow::Int16Array.new([2])
 172       ]
 173       taken_chunked_array = Arrow::ChunkedArray.new(taken_chunks)
 174       indices_chunks = [
 175         Arrow::Int16Array.new([1, 0]),
 176         Arrow::Int16Array.new([2])
 177       ]
 178       indices = Arrow::ChunkedArray.new(indices_chunks)
 179       assert_equal(taken_chunked_array,
 180                    @chunked_array.take(indices))
 181     end
 182   end
 183 end