lidavidm commented on a change in pull request #12162:
URL: https://github.com/apache/arrow/pull/12162#discussion_r788841782
##########
File path: cpp/src/arrow/compute/kernels/scalar_nested_test.cc
##########
@@ -225,6 +225,167 @@ TEST(TestScalarNested, StructField) {
}
}
+TEST(TestScalarNested, MapArrayLookup) {
+ MapArrayLookupOptions foo_all(MakeScalar("foo"), MapArrayLookupOptions::ALL);
+ MapArrayLookupOptions foo_first(MakeScalar("foo"),
MapArrayLookupOptions::FIRST);
+ MapArrayLookupOptions foo_last(MakeScalar("foo"),
MapArrayLookupOptions::LAST);
+
+ auto type = map(utf8(), int32());
+ const char* input = R"(
+ [
+ [["foo", 99], ["bar", 1], ["hello", 2], ["foo", 3], ["lesgo", 5],
["whatnow", 8]],
+ null,
+ [["nothing", null], ["hat", null], ["foo", 101], ["sorry", 1], ["dip",
null],
+ ["foo", 22]],
+ []
+ ]
+ )";
+ auto map_array = ArrayFromJSON(type, input);
+
+ CheckScalar("map_array_lookup", {map_array},
+ ArrayFromJSON(list(int32()), "[[99, 3], null, [101, 22],
null]"), &foo_all);
+ CheckScalar("map_array_lookup", {map_array},
+ ArrayFromJSON(int32(), "[99, null, 101, null]"), &foo_first);
+ CheckScalar("map_array_lookup", {map_array},
+ ArrayFromJSON(int32(), "[3, null, 22, null]"), &foo_last);
+}
+
+TEST(TestScalarNested, MapArrayLookupNested) {
+ auto type = map(utf8(), map(int16(), int16()));
+ const char* input = R"(
+ [
+ [
+ [
+ "just",
+ [[0, 0], [1, 1]]
+ ],
+ [
+ "random",
+ [[2, 2], [3, 3]]
+ ],
+ [
+ "foo",
+ [[4, 4], [5, 5]]
+ ],
+ [
+ "values",
+ [[6, 6], [7, 7]]
+ ],
+ [
+ "foo",
+ [[8, 8], [9, 9]]
+ ],
+ [
+ "point",
+ [[10, 10], [11, 11]]
+ ],
+ [
+ "foo",
+ [[12, 12], [13, 13]]
+ ]
+ ],
+ null,
+ [
+ [
+ "yet",
+ [[0, 1], [1, 2]]
+ ],
+ [
+ "more",
+ [[2, 3], [3, 4]]
+ ],
+ [
+ "foo",
+ [[4, 5], [5, 6]]
+ ],
+ [
+ "random",
+ [[6, 7], [7, 8]]
+ ],
+ [
+ "foo",
+ [[8, 9], [9, 10]]
+ ],
+ [
+ "values",
+ [[10, 11], [11, 12]]
+ ],
+ [
+ "foo",
+ [[12, 13], [13, 14]]
+ ]
+ ],
+ []
+ ]
+ )";
+ auto map_array = ArrayFromJSON(type, input);
+
+ MapArrayLookupOptions foo_all(MakeScalar("foo"), MapArrayLookupOptions::ALL);
+ MapArrayLookupOptions foo_first(MakeScalar("foo"),
MapArrayLookupOptions::FIRST);
+ MapArrayLookupOptions foo_last(MakeScalar("foo"),
MapArrayLookupOptions::LAST);
+
+ auto foo_all_output = ArrayFromJSON(
+ list(map(int16(), int16())),
+ "[ [[[4, 4], [5, 5]], [[8, 8], [9, 9]], [[12, 12], [13, 13]]], null,
[[[4, 5], [5, "
+ "6]], [[8, 9], [9, 10]], [[12, 13], [13, 14]]], null ]");
+
+ CheckScalar("map_array_lookup", {map_array}, foo_all_output, &foo_all);
+ CheckScalar("map_array_lookup", {map_array},
+ ArrayFromJSON(map(int16(), int16()),
+ "[ [[4, 4], [5, 5]], null, [[4, 5], [5, 6]], null
]"),
+ &foo_first);
+ CheckScalar("map_array_lookup", {map_array},
+ ArrayFromJSON(map(int16(), int16()),
+ "[ [[12, 12], [13, 13]], null, [[12, 13], [13,
14]], null ]"),
+ &foo_last);
+}
+
+template <typename Type>
+class TestMapArrayLookupIntegralKeys : public ::testing ::Test {};
+
+TYPED_TEST_SUITE(TestMapArrayLookupIntegralKeys, IntegralArrowTypes);
+
+TYPED_TEST(TestMapArrayLookupIntegralKeys, StringItems) {
+ auto type = default_type_instance<TypeParam>();
+
+ auto one_scalar = MakeScalar(type, 1).ValueOrDie();
+ MapArrayLookupOptions one_all(one_scalar, MapArrayLookupOptions::ALL);
+ MapArrayLookupOptions one_first(one_scalar, MapArrayLookupOptions::FIRST);
+ MapArrayLookupOptions one_last(one_scalar, MapArrayLookupOptions::LAST);
+
+ auto map_type = map(type, utf8());
+ const char* input = R"(
+ [
+ [
+ [0, "zero"], [1, "first_one"], [2, "two"], [3, "three"], [1,
"second_one"],
+ [1, "last_one"]
+ ],
+ null,
+ [
+ [0, "zero_hero"], [9, "almost_six"], [1, "the_dumb_one"], [7,
"eleven"],
+ [1, "the_chosen_one"], [42, "meaning of life?"], [1, "just_one"],
+ [1, "no more ones!"]
+ ],
+ []
+ ]
+ )";
+ auto map_array = ArrayFromJSON(map_type, input);
+ CheckScalar("map_array_lookup", {map_array},
+ ArrayFromJSON(utf8(), R"(["first_one", null, "the_dumb_one",
null])"),
+ &one_first);
+ CheckScalar("map_array_lookup", {map_array},
+ ArrayFromJSON(utf8(), R"(["last_one", null, "no more ones!",
null])"),
+ &one_last);
+ CheckScalar("map_array_lookup", {map_array}, ArrayFromJSON(list(utf8()), R"([
Review comment:
For the last case, use TweakValidityBit:
https://github.com/apache/arrow/blob/c39caedc2e9296fcb53d1c587708836ec6628ddb/cpp/src/arrow/testing/gtest_util.h#L347-L353
```cpp
auto arr = ArrayFromJSON(ty, "[["foo", 1]]");
auto tweaked = TweakValidityBit(arr, 0, false); // Now arr[0] is null, but
the underlying data is still there
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
