lidavidm commented on a change in pull request #12162:
URL: https://github.com/apache/arrow/pull/12162#discussion_r788838333
##########
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) {
Review comment:
I would say these are the "groups" you need:
`PhysicalIntegralArrowTypes`, `DecimalArrowTypes`, `BaseBinaryArrowTypes`, and
then miscellaneous types (Boolean, FixedSizeString, MonthDayNanoInterval)
See here for all of the categories (note they overlap/not all are relevant):
https://github.com/apache/arrow/blob/c39caedc2e9296fcb53d1c587708836ec6628ddb/cpp/src/arrow/testing/gtest_util.h#L153-L185
##########
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:
Hmm, some cases to consider: empty maps, null maps (both already
covered); null keys, null values (do we want to support looking up null keys?
should consider that), and null maps where the values are not actually empty
(this is harder to test, there's a helper for it, let me find it)
--
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]
