ceph/src/arrow/cpp/src/arrow/compute/kernels/vector_nested_test.cc

   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 #include <gtest/gtest.h>
  19
  20 #include "arrow/chunked_array.h"
  21 #include "arrow/compute/api.h"
  22 #include "arrow/compute/kernels/test_util.h"
  23 #include "arrow/result.h"
  24 #include "arrow/testing/gtest_util.h"
  25 #include "arrow/util/checked_cast.h"
  26
  27 namespace arrow {
  28 namespace compute {
  29
  30 using arrow::internal::checked_cast;
  31
  32 TEST(TestVectorNested, ListFlatten) {
  33   for (auto ty : {list(int16()), large_list(int16())}) {
  34     auto input = ArrayFromJSON(ty, "[[0, null, 1], null, [2, 3], []]");
  35     auto expected = ArrayFromJSON(int16(), "[0, null, 1, 2, 3]");
  36     CheckVectorUnary("list_flatten", input, expected);
  37
  38     // Construct a list with a non-empty null slot
  39     TweakValidityBit(input, 0, false);
  40     expected = ArrayFromJSON(int16(), "[2, 3]");
  41     CheckVectorUnary("list_flatten", input, expected);
  42   }
  43 }
  44
  45 TEST(TestVectorNested, ListFlattenChunkedArray) {
  46   for (auto ty : {list(int16()), large_list(int16())}) {
  47     auto input = ChunkedArrayFromJSON(ty, {"[[0, null, 1], null]", "[[2, 3], []]"});
  48     auto expected = ChunkedArrayFromJSON(int16(), {"[0, null, 1]", "[2, 3]"});
  49     CheckVectorUnary("list_flatten", input, expected);
  50
  51     input = ChunkedArrayFromJSON(ty, {});
  52     expected = ChunkedArrayFromJSON(int16(), {});
  53     CheckVectorUnary("list_flatten", input, expected);
  54   }
  55 }
  56
  57 TEST(TestVectorNested, ListFlattenFixedSizeList) {
  58   for (auto ty : {fixed_size_list(int16(), 2), fixed_size_list(uint32(), 2)}) {
  59     const auto& out_ty = checked_cast<const FixedSizeListType&>(*ty).value_type();
  60     {
  61       auto input = ArrayFromJSON(ty, "[[0, null], null, [2, 3], [0, 42]]");
  62       auto expected = ArrayFromJSON(out_ty, "[0, null, 2, 3, 0, 42]");
  63       CheckVectorUnary("list_flatten", input, expected);
  64     }
  65
  66     {
  67       // Test a chunked array
  68       auto input = ChunkedArrayFromJSON(ty, {"[[0, null], null]", "[[2, 3], [0, 42]]"});
  69       auto expected = ChunkedArrayFromJSON(out_ty, {"[0, null]", "[2, 3, 0, 42]"});
  70       CheckVectorUnary("list_flatten", input, expected);
  71
  72       input = ChunkedArrayFromJSON(ty, {});
  73       expected = ChunkedArrayFromJSON(out_ty, {});
  74       CheckVectorUnary("list_flatten", input, expected);
  75     }
  76   }
  77 }
  78
  79 TEST(TestVectorNested, ListParentIndices) {
  80   for (auto ty : {list(int16()), large_list(int16())}) {
  81     auto input = ArrayFromJSON(ty, "[[0, null, 1], null, [2, 3], [], [4, 5]]");
  82
  83     auto out_ty = ty->id() == Type::LIST ? int32() : int64();
  84     auto expected = ArrayFromJSON(out_ty, "[0, 0, 0, 2, 2, 4, 4]");
  85     CheckVectorUnary("list_parent_indices", input, expected);
  86   }
  87
  88   // Construct a list with a non-empty null slot
  89   auto input = ArrayFromJSON(list(int16()), "[[0, null, 1], [0, 0], [2, 3], [], [4, 5]]");
  90   TweakValidityBit(input, 1, false);
  91   auto expected = ArrayFromJSON(int32(), "[0, 0, 0, 1, 1, 2, 2, 4, 4]");
  92   CheckVectorUnary("list_parent_indices", input, expected);
  93 }
  94
  95 TEST(TestVectorNested, ListParentIndicesChunkedArray) {
  96   for (auto ty : {list(int16()), large_list(int16())}) {
  97     auto input =
  98         ChunkedArrayFromJSON(ty, {"[[0, null, 1], null]", "[[2, 3], [], [4, 5]]"});
  99
 100     auto out_ty = ty->id() == Type::LIST ? int32() : int64();
 101     auto expected = ChunkedArrayFromJSON(out_ty, {"[0, 0, 0]", "[2, 2, 4, 4]"});
 102     CheckVectorUnary("list_parent_indices", input, expected);
 103
 104     input = ChunkedArrayFromJSON(ty, {});
 105     expected = ChunkedArrayFromJSON(out_ty, {});
 106     CheckVectorUnary("list_parent_indices", input, expected);
 107   }
 108 }
 109
 110 TEST(TestVectorNested, ListParentIndicesFixedSizeList) {
 111   for (auto ty : {fixed_size_list(int16(), 2), fixed_size_list(uint32(), 2)}) {
 112     {
 113       auto input = ArrayFromJSON(ty, "[[0, null], null, [1, 2], [3, 4], [null, 5]]");
 114       auto expected = ArrayFromJSON(int32(), "[0, 0, 2, 2, 3, 3, 4, 4]");
 115       CheckVectorUnary("list_parent_indices", input, expected);
 116     }
 117     {
 118       // Test a chunked array
 119       auto input =
 120           ChunkedArrayFromJSON(ty, {"[[0, null], null, [1, 2]]", "[[3, 4], [null, 5]]"});
 121       auto expected = ChunkedArrayFromJSON(int32(), {"[0, 0, 2, 2]", "[3, 3, 4, 4]"});
 122       CheckVectorUnary("list_parent_indices", input, expected);
 123
 124       input = ChunkedArrayFromJSON(ty, {});
 125       expected = ChunkedArrayFromJSON(int32(), {});
 126       CheckVectorUnary("list_parent_indices", input, expected);
 127     }
 128   }
 129 }
 130
 131 }  // namespace compute
 132 }  // namespace arrow