|
23 | 23 | #include "arrow/array.h" |
24 | 24 | #include "arrow/array/util.h" |
25 | 25 | #include "arrow/compute/api_vector.h" |
| 26 | +#include "arrow/compute/cast.h" |
26 | 27 | #include "arrow/compute/exec.h" |
27 | 28 | #include "arrow/compute/row/grouper.h" |
28 | 29 | #include "arrow/compute/row/grouper_internal.h" |
@@ -120,6 +121,8 @@ void TestGroupClassSupportedKeys( |
120 | 121 |
|
121 | 122 | ASSERT_OK(make_func({utf8(), binary(), large_utf8(), large_binary()})); |
122 | 123 |
|
| 124 | + ASSERT_OK(make_func({utf8_view(), binary_view()})); |
| 125 | + |
123 | 126 | ASSERT_OK(make_func({fixed_size_binary(16), fixed_size_binary(32)})); |
124 | 127 |
|
125 | 128 | ASSERT_OK(make_func({decimal128(32, 10), decimal256(76, 20)})); |
@@ -754,13 +757,24 @@ struct TestGrouper { |
754 | 757 | auto group_ids = id_batch.make_array(); |
755 | 758 | ValidateOutput(*group_ids); |
756 | 759 |
|
| 760 | + // Take has no view kernel yet (GH-43010); validate via the non-view cast. |
| 761 | + // Output type is checked separately by ExpectUniques. |
| 762 | + auto as_takeable = [](std::shared_ptr<Array> arr) -> std::shared_ptr<Array> { |
| 763 | + if (is_binary_view_like(*arr->type())) { |
| 764 | + auto target = arr->type_id() == Type::STRING_VIEW ? utf8() : binary(); |
| 765 | + arr = Cast(Datum(arr), target).ValueOrDie().make_array(); |
| 766 | + } |
| 767 | + return arr; |
| 768 | + }; |
| 769 | + |
757 | 770 | for (int i = 0; i < key_batch.num_values(); ++i) { |
758 | 771 | SCOPED_TRACE(ToChars(i) + "th key array"); |
759 | 772 | auto original = |
760 | 773 | key_batch[i].is_array() |
761 | 774 | ? key_batch[i].make_array() |
762 | 775 | : *MakeArrayFromScalar(*key_batch[i].scalar(), key_batch.length); |
763 | | - ASSERT_OK_AND_ASSIGN(auto encoded, Take(*uniques_[i].make_array(), *group_ids)); |
| 776 | + ASSERT_OK_AND_ASSIGN(auto encoded, |
| 777 | + Take(*as_takeable(uniques_[i].make_array()), *group_ids)); |
764 | 778 | std::shared_ptr<Array> expected = original; |
765 | 779 | if (can_be_null && original->type_id() != Type::NA) { |
766 | 780 | // To compute the expected output, mask out the original entries that |
@@ -791,7 +805,7 @@ struct TestGrouper { |
791 | 805 | expected_data->null_count = kUnknownNullCount; |
792 | 806 | expected = MakeArray(expected_data); |
793 | 807 | } |
794 | | - AssertArraysEqual(*expected, *encoded, /*verbose=*/true, |
| 808 | + AssertArraysEqual(*as_takeable(expected), *encoded, /*verbose=*/true, |
795 | 809 | EqualOptions().nans_equal(true)); |
796 | 810 | } |
797 | 811 | } |
@@ -913,6 +927,48 @@ TEST(Grouper, StringKey) { |
913 | 927 | } |
914 | 928 | } |
915 | 929 |
|
| 930 | +TEST(Grouper, StringViewKey) { |
| 931 | + // Mix inline (<=12 byte) and out-of-line view keys. |
| 932 | + for (auto ty : {utf8_view(), binary_view()}) { |
| 933 | + ARROW_SCOPED_TRACE("key type = ", *ty); |
| 934 | + { |
| 935 | + TestGrouper g({ty}); |
| 936 | + g.ExpectConsume(R"([["eh"], ["eh"]])", "[0, 0]"); |
| 937 | + g.ExpectConsume(R"([["eh"], ["eh"]])", "[0, 0]"); |
| 938 | + g.ExpectConsume(R"([["be"], [null]])", "[1, 2]"); |
| 939 | + // Several distinct out-of-line (>12 byte) views, including two that share |
| 940 | + // the same 4-byte inline prefix. |
| 941 | + g.ExpectConsume( |
| 942 | + R"([["a long out-of-line view"], ["a long out-of-line view"], |
| 943 | + ["a different long view"], ["a distinct long-ish view"]])", |
| 944 | + "[3, 3, 4, 5]"); |
| 945 | + g.ExpectUniques( |
| 946 | + R"([["eh"], ["be"], [null], ["a long out-of-line view"], |
| 947 | + ["a different long view"], ["a distinct long-ish view"]])"); |
| 948 | + } |
| 949 | + { |
| 950 | + TestGrouper g({ty}); |
| 951 | + g.ExpectPopulate(R"([["eh"], ["eh"]])"); |
| 952 | + g.ExpectPopulate(R"([["be"], [null]])"); |
| 953 | + g.ExpectLookup(R"([["be"], [null], ["da"]])", "[1, 2, null]"); |
| 954 | + } |
| 955 | + } |
| 956 | +} |
| 957 | + |
| 958 | +TEST(Grouper, StringViewInt64Key) { |
| 959 | + for (auto ty : {utf8_view(), binary_view()}) { |
| 960 | + ARROW_SCOPED_TRACE("key type = ", *ty); |
| 961 | + TestGrouper g({ty, int64()}); |
| 962 | + |
| 963 | + g.ExpectConsume(R"([["eh", 0], ["eh", 0]])", "[0, 0]"); |
| 964 | + g.ExpectConsume(R"([["eh", 0], ["eh", null]])", "[0, 1]"); |
| 965 | + g.ExpectConsume(R"([["eh", 1], ["bee", 1]])", "[2, 3]"); |
| 966 | + g.ExpectConsume(R"([["eh", null], ["bee", 1]])", "[1, 3]"); |
| 967 | + |
| 968 | + g.ExpectUniques(R"([["eh", 0], ["eh", null], ["eh", 1], ["bee", 1]])"); |
| 969 | + } |
| 970 | +} |
| 971 | + |
916 | 972 | TEST(Grouper, DictKey) { |
917 | 973 | // For dictionary keys, all batches must share a single dictionary. |
918 | 974 | // Eventually, differing dictionaries will be unified and indices transposed |
@@ -1074,6 +1130,10 @@ FieldVector AnnotateForRandomGeneration(FieldVector fields) { |
1074 | 1130 | } else if (is_binary_like(*field->type())) { |
1075 | 1131 | // (note this is unsupported for large binary types) |
1076 | 1132 | field = field->WithMergedMetadata(key_value_metadata({"unique"}, {"100"})); |
| 1133 | + } else if (is_binary_view_like(*field->type())) { |
| 1134 | + // Views don't support the "unique" knob; small lengths keep cardinality low. |
| 1135 | + field = field->WithMergedMetadata( |
| 1136 | + key_value_metadata({"min_length", "max_length"}, {"0", "5"})); |
1077 | 1137 | } |
1078 | 1138 | field = field->WithMergedMetadata(key_value_metadata({"null_probability"}, {"0.1"})); |
1079 | 1139 | } |
@@ -1130,6 +1190,22 @@ TEST(Grouper, RandomStringInt64DoubleInt32Keys) { |
1130 | 1190 | TestRandomLookup(TestGrouper({utf8(), int64(), float64(), int32()})); |
1131 | 1191 | } |
1132 | 1192 |
|
| 1193 | +TEST(Grouper, RandomStringViewKeys) { |
| 1194 | + for (auto ty : {utf8_view(), binary_view()}) { |
| 1195 | + ARROW_SCOPED_TRACE("key type = ", *ty); |
| 1196 | + TestRandomConsume(TestGrouper({ty})); |
| 1197 | + TestRandomLookup(TestGrouper({ty})); |
| 1198 | + } |
| 1199 | +} |
| 1200 | + |
| 1201 | +TEST(Grouper, RandomStringViewInt64Keys) { |
| 1202 | + for (auto ty : {utf8_view(), binary_view()}) { |
| 1203 | + ARROW_SCOPED_TRACE("key type = ", *ty); |
| 1204 | + TestRandomConsume(TestGrouper({ty, int64()})); |
| 1205 | + TestRandomLookup(TestGrouper({ty, int64()})); |
| 1206 | + } |
| 1207 | +} |
| 1208 | + |
1133 | 1209 | TEST(Grouper, NullKeys) { |
1134 | 1210 | { |
1135 | 1211 | TestGrouper g({null()}); |
|
0 commit comments