ceph/src/arrow/c_glib/test/test-take.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 TestTake < Test::Unit::TestCase
  19   include Helper::Buildable
  20   include Helper::Omittable
  21
  22   sub_test_case("Array") do
  23     def test_no_null
  24       indices = build_int16_array([1, 0, 2])
  25       assert_equal(build_int16_array([0, 1, 2]),
  26                    build_int16_array([1, 0 ,2]).take(indices))
  27     end
  28
  29     def test_null
  30       indices = build_int16_array([2, nil, 0])
  31       assert_equal(build_int16_array([2, nil, 1]),
  32                    build_int16_array([1, 0, 2]).take(indices))
  33     end
  34
  35     def test_out_of_index
  36       indices = build_int16_array([1, 2, 3])
  37       assert_raise(Arrow::Error::Index) do
  38         build_int16_array([0, 1, 2]).take(indices)
  39       end
  40     end
  41
  42     def test_chunked_array
  43       taken_chunks = [
  44         build_int16_array([0, 1]),
  45         build_int16_array([2])
  46       ]
  47       taken_chunked_array = Arrow::ChunkedArray.new(taken_chunks)
  48       indices_chunks = [
  49         build_int16_array([1, 0]),
  50         build_int16_array([2])
  51       ]
  52       indices = Arrow::ChunkedArray.new(indices_chunks)
  53       assert_equal(taken_chunked_array,
  54                    build_int16_array([1, 0, 2]).take_chunked_array(indices))
  55     end
  56   end
  57
  58   sub_test_case("Table") do
  59     def setup
  60       fields = [
  61         Arrow::Field.new("field1", Arrow::Int16DataType.new),
  62         Arrow::Field.new("field2", Arrow::Int16DataType.new)
  63       ]
  64       @schema = Arrow::Schema.new(fields)
  65       arrays = [
  66         build_int16_array([0, 1, 2]),
  67         build_int16_array([3, 5, 4])
  68       ]
  69       @table = Arrow::Table.new(@schema, arrays)
  70     end
  71
  72     def test_no_null
  73       arrays = [
  74         build_int16_array([1, 0, 2]),
  75         build_int16_array([5, 3, 4])
  76       ]
  77       taken_table = Arrow::Table.new(@schema, arrays)
  78       indices = build_int16_array([1, 0, 2])
  79       assert_equal(taken_table,
  80                    @table.take(indices))
  81     end
  82
  83     def test_null
  84       arrays = [
  85         build_int16_array([2, nil, 0]),
  86         build_int16_array([4, nil, 3])
  87       ]
  88       taken_table = Arrow::Table.new(@schema, arrays)
  89       indices = build_int16_array([2, nil, 0])
  90       assert_equal(taken_table,
  91                    @table.take(indices))
  92     end
  93
  94     def test_out_of_index
  95       indices = build_int16_array([1, 2, 3])
  96       assert_raise(Arrow::Error::Index) do
  97         @table.take(indices)
  98       end
  99     end
 100
 101     def test_chunked_array
 102       arrays = [
 103         build_int16_array([1, 0, 2]),
 104         build_int16_array([5, 3, 4])
 105       ]
 106       taken_table = Arrow::Table.new(@schema, arrays)
 107       chunks = [
 108         build_int16_array([1, 0]),
 109         build_int16_array([2])
 110       ]
 111       indices = Arrow::ChunkedArray.new(chunks)
 112       assert_equal(taken_table,
 113                    @table.take_chunked_array(indices))
 114     end
 115   end
 116
 117   sub_test_case("ChunkedArray") do
 118     def setup
 119       chunks = [
 120         build_int16_array([1, 0]),
 121         build_int16_array([2]),
 122       ]
 123       @chunked_array = Arrow::ChunkedArray.new(chunks)
 124     end
 125
 126     def test_no_null
 127       chunks = [
 128         build_int16_array([0, 1]),
 129         build_int16_array([2])
 130       ]
 131       taken_chunked_array = Arrow::ChunkedArray.new(chunks)
 132       indices = build_int16_array([1, 0, 2])
 133       assert_equal(taken_chunked_array,
 134                    @chunked_array.take(indices))
 135     end
 136
 137     def test_null
 138       chunks = [
 139         build_int16_array([2, nil]),
 140         build_int16_array([1])
 141       ]
 142       taken_chunked_array = Arrow::ChunkedArray.new(chunks)
 143       indices = build_int16_array([2, nil, 0])
 144       assert_equal(taken_chunked_array,
 145                    @chunked_array.take(indices))
 146     end
 147
 148     def test_out_of_index
 149       indices = build_int16_array([1, 2, 3])
 150       assert_raise(Arrow::Error::Index) do
 151         @chunked_array.take(indices)
 152       end
 153     end
 154
 155     def test_chunked_array
 156       taken_chunks = [
 157         build_int16_array([0, 1]),
 158         build_int16_array([2])
 159       ]
 160       taken_chunked_array = Arrow::ChunkedArray.new(taken_chunks)
 161       indices_chunks = [
 162         build_int16_array([1, 0]),
 163         build_int16_array([2])
 164       ]
 165       indices = Arrow::ChunkedArray.new(indices_chunks)
 166       assert_equal(taken_chunked_array,
 167                    @chunked_array.take_chunked_array(indices))
 168     end
 169   end
 170
 171   sub_test_case("RecordBatch") do
 172     def setup
 173       fields = [
 174         Arrow::Field.new("field1", Arrow::Int16DataType.new),
 175         Arrow::Field.new("field2", Arrow::Int16DataType.new)
 176       ]
 177       @schema = Arrow::Schema.new(fields)
 178       columns = [
 179         build_int16_array([1, 0, 2]),
 180         build_int16_array([3, 5, 4])
 181       ]
 182       @record_batch = Arrow::RecordBatch.new(@schema, 3, columns)
 183     end
 184
 185     def test_no_null
 186       columns = [
 187         build_int16_array([0, 1, 2]),
 188         build_int16_array([5, 3, 4])
 189       ]
 190       taken_record_batch = Arrow::RecordBatch.new(@schema, 3, columns)
 191       indices = build_int16_array([1, 0, 2])
 192       assert_equal(taken_record_batch,
 193                    @record_batch.take(indices))
 194     end
 195
 196     def test_null
 197       columns = [
 198         build_int16_array([2, nil, 1]),
 199         build_int16_array([4, nil, 3])
 200       ]
 201       taken_record_batch = Arrow::RecordBatch.new(@schema, 3, columns)
 202       indices = build_int16_array([2, nil, 0])
 203       assert_equal(taken_record_batch,
 204                    @record_batch.take(indices))
 205     end
 206
 207     def test_out_of_index
 208       indices = build_int16_array([1, 2, 3])
 209       assert_raise(Arrow::Error::Index) do
 210         @record_batch.take(indices)
 211       end
 212     end
 213   end
 214 end