Skip to content

Commit 7239379

Browse files
authored
GH-50223: [C++][Compute] Support string_view/binary_view keys in the hash-aggregate Grouper (#50224)
### Rationale for this change Grouper rejected `string_view` and `binary_view` group keys with `NotImplemented: Keys of type string_view`, so `Table.group_by` on a view-typed key failed. ### What changes are included in this PR? - Add `BinaryViewKeyEncoder` in `row_encoder_internal.{h,cc}`. It encodes view keys using the existing variable-length row format and decodes them into a fresh view array of the original type. - `GrouperImpl::Make` dispatches view keys to the new encoder. - `GrouperFastImpl::CanUse` rejects view keys, so they fall back to GrouperImpl. ### Are these changes tested? Yes, unit tests added. ### Are there any user-facing changes? Yes. `Table.group_by` now accepts `string_view` and `binary_view` group keys * GitHub Issue: #50223 Authored-by: Fangchen Li <fangchen.li@outlook.com> Signed-off-by: Antoine Pitrou <antoine@python.org>
1 parent dd69699 commit 7239379

6 files changed

Lines changed: 252 additions & 16 deletions

File tree

cpp/src/arrow/compute/row/grouper.cc

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -373,6 +373,12 @@ struct GrouperImpl : public Grouper {
373373
continue;
374374
}
375375

376+
if (is_binary_view_like(key_type.id())) {
377+
impl->encoders_[i] =
378+
std::make_unique<internal::BinaryViewKeyEncoder>(key_type.GetSharedPtr());
379+
continue;
380+
}
381+
376382
if (key_type.id() == Type::NA) {
377383
impl->encoders_[i] = std::make_unique<internal::NullKeyEncoder>();
378384
continue;
@@ -556,7 +562,10 @@ struct GrouperFastImpl : public Grouper {
556562
}
557563
#if ARROW_LITTLE_ENDIAN
558564
for (size_t i = 0; i < key_types.size(); ++i) {
559-
if (is_large_binary_like(key_types[i].id())) {
565+
// View keys fall back to GrouperImpl; the fast row format can't hold
566+
// their variadic buffers.
567+
if (is_large_binary_like(key_types[i].id()) ||
568+
is_binary_view_like(key_types[i].id())) {
560569
return false;
561570
}
562571
}

cpp/src/arrow/compute/row/grouper_benchmark.cc

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,11 @@
2020
#include "arrow/util/key_value_metadata.h"
2121
#include "arrow/util/string.h"
2222

23+
#include "arrow/compute/cast.h"
2324
#include "arrow/compute/row/grouper.h"
2425
#include "arrow/testing/gtest_util.h"
2526
#include "arrow/testing/random.h"
27+
#include "arrow/type_traits.h"
2628
#include "arrow/util/benchmark_util.h"
2729

2830
namespace arrow {
@@ -62,8 +64,20 @@ static ExecBatch MakeRandomExecBatch(const DataTypeVector& types, int64_t num_ro
6264
std::vector<Datum> values;
6365
values.resize(num_types);
6466
for (int i = 0; i < num_types; ++i) {
65-
auto field = ::arrow::field("", types[i], metadata);
66-
values[i] = rng.ArrayOf(*field, num_rows, alignment, memory_pool);
67+
// The "unique" cardinality knob isn't honored for view types, so generate the
68+
// plain equivalent and cast it (outside the timed loop) to match {utf8}.
69+
const auto& type = types[i];
70+
const bool is_view = is_binary_view_like(*type);
71+
std::shared_ptr<DataType> gen_type = type;
72+
if (is_view) {
73+
gen_type = type->id() == Type::STRING_VIEW ? utf8() : binary();
74+
}
75+
auto field = ::arrow::field("", gen_type, metadata);
76+
Datum value = rng.ArrayOf(*field, num_rows, alignment, memory_pool);
77+
if (is_view) {
78+
ASSIGN_OR_ABORT(value, compute::Cast(value, type));
79+
}
80+
values[i] = std::move(value);
6781
}
6882

6983
return ExecBatch(std::move(values), num_rows);
@@ -125,6 +139,8 @@ BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{boolean}", {boolean()})->Apply(SetArg
125139
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{int32}", {int32()})->Apply(SetArgs);
126140
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{int64}", {int64()})->Apply(SetArgs);
127141
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{utf8}", {utf8()})->Apply(SetArgs);
142+
// View keys use the generic GrouperImpl path (the {utf8} fast path rejects them).
143+
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{utf8_view}", {utf8_view()})->Apply(SetArgs);
128144
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{fixed_size_binary(32)}",
129145
{fixed_size_binary(32)})
130146
->Apply(SetArgs);
@@ -147,6 +163,9 @@ BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{int32, boolean, utf8}",
147163
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{int32, int64, boolean, utf8}",
148164
{int32(), int64(), boolean(), utf8()})
149165
->Apply(SetArgs);
166+
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{int32, int64, boolean, utf8_view}",
167+
{int32(), int64(), boolean(), utf8_view()})
168+
->Apply(SetArgs);
150169
BENCHMARK_CAPTURE(GrouperWithMultiTypes,
151170
"{utf8, int32, int64, fixed_size_binary(32), boolean}",
152171
{utf8(), int32(), int64(), fixed_size_binary(32), boolean()})

cpp/src/arrow/compute/row/grouper_test.cc

Lines changed: 78 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
#include "arrow/array.h"
2424
#include "arrow/array/util.h"
2525
#include "arrow/compute/api_vector.h"
26+
#include "arrow/compute/cast.h"
2627
#include "arrow/compute/exec.h"
2728
#include "arrow/compute/row/grouper.h"
2829
#include "arrow/compute/row/grouper_internal.h"
@@ -120,6 +121,8 @@ void TestGroupClassSupportedKeys(
120121

121122
ASSERT_OK(make_func({utf8(), binary(), large_utf8(), large_binary()}));
122123

124+
ASSERT_OK(make_func({utf8_view(), binary_view()}));
125+
123126
ASSERT_OK(make_func({fixed_size_binary(16), fixed_size_binary(32)}));
124127

125128
ASSERT_OK(make_func({decimal128(32, 10), decimal256(76, 20)}));
@@ -754,13 +757,24 @@ struct TestGrouper {
754757
auto group_ids = id_batch.make_array();
755758
ValidateOutput(*group_ids);
756759

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+
757770
for (int i = 0; i < key_batch.num_values(); ++i) {
758771
SCOPED_TRACE(ToChars(i) + "th key array");
759772
auto original =
760773
key_batch[i].is_array()
761774
? key_batch[i].make_array()
762775
: *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));
764778
std::shared_ptr<Array> expected = original;
765779
if (can_be_null && original->type_id() != Type::NA) {
766780
// To compute the expected output, mask out the original entries that
@@ -791,7 +805,7 @@ struct TestGrouper {
791805
expected_data->null_count = kUnknownNullCount;
792806
expected = MakeArray(expected_data);
793807
}
794-
AssertArraysEqual(*expected, *encoded, /*verbose=*/true,
808+
AssertArraysEqual(*as_takeable(expected), *encoded, /*verbose=*/true,
795809
EqualOptions().nans_equal(true));
796810
}
797811
}
@@ -913,6 +927,48 @@ TEST(Grouper, StringKey) {
913927
}
914928
}
915929

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+
916972
TEST(Grouper, DictKey) {
917973
// For dictionary keys, all batches must share a single dictionary.
918974
// Eventually, differing dictionaries will be unified and indices transposed
@@ -1074,6 +1130,10 @@ FieldVector AnnotateForRandomGeneration(FieldVector fields) {
10741130
} else if (is_binary_like(*field->type())) {
10751131
// (note this is unsupported for large binary types)
10761132
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"}));
10771137
}
10781138
field = field->WithMergedMetadata(key_value_metadata({"null_probability"}, {"0.1"}));
10791139
}
@@ -1130,6 +1190,22 @@ TEST(Grouper, RandomStringInt64DoubleInt32Keys) {
11301190
TestRandomLookup(TestGrouper({utf8(), int64(), float64(), int32()}));
11311191
}
11321192

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+
11331209
TEST(Grouper, NullKeys) {
11341210
{
11351211
TestGrouper g({null()});

cpp/src/arrow/compute/row/row_encoder_internal.cc

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
#include "arrow/compute/row/row_encoder_internal.h"
1919

20+
#include "arrow/array/builder_binary.h"
2021
#include "arrow/util/bitmap_writer.h"
2122
#include "arrow/util/logging_internal.h"
2223

@@ -256,6 +257,34 @@ Result<std::shared_ptr<ArrayData>> DictionaryKeyEncoder::Decode(uint8_t** encode
256257
return data;
257258
}
258259

260+
Result<std::shared_ptr<ArrayData>> BinaryViewKeyEncoder::Decode(uint8_t** encoded_bytes,
261+
int32_t length,
262+
MemoryPool* pool) {
263+
// Build a fresh view array; MakeBuilder gives the type-faithful builder and
264+
// Append handles inline vs. out-of-line storage.
265+
std::unique_ptr<ArrayBuilder> builder;
266+
RETURN_NOT_OK(MakeBuilder(pool, type_, &builder));
267+
auto& view_builder = checked_cast<BinaryViewBuilder&>(*builder);
268+
RETURN_NOT_OK(view_builder.Reserve(length));
269+
270+
for (int32_t i = 0; i < length; ++i) {
271+
uint8_t*& encoded_ptr = encoded_bytes[i];
272+
const bool is_valid = (*encoded_ptr++ == kValidByte);
273+
const auto key_length = util::SafeLoadAs<Offset>(encoded_ptr);
274+
encoded_ptr += sizeof(Offset);
275+
if (is_valid) {
276+
RETURN_NOT_OK(view_builder.Append(encoded_ptr, key_length));
277+
} else {
278+
RETURN_NOT_OK(view_builder.AppendNull());
279+
}
280+
encoded_ptr += key_length; // zero for null rows
281+
}
282+
283+
std::shared_ptr<Array> out;
284+
RETURN_NOT_OK(view_builder.Finish(&out));
285+
return out->data();
286+
}
287+
259288
void RowEncoder::Init(const std::vector<TypeHolder>& column_types, ExecContext* ctx) {
260289
ctx_ = ctx;
261290
encoders_.resize(column_types.size());
@@ -301,6 +330,11 @@ void RowEncoder::Init(const std::vector<TypeHolder>& column_types, ExecContext*
301330
continue;
302331
}
303332

333+
if (is_binary_view_like(type.id())) {
334+
encoders_[i] = std::make_shared<BinaryViewKeyEncoder>(type.GetSharedPtr());
335+
continue;
336+
}
337+
304338
// We should not get here
305339
ARROW_DCHECK(false);
306340
}

cpp/src/arrow/compute/row/row_encoder_internal.h

Lines changed: 39 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -137,15 +137,22 @@ struct ARROW_COMPUTE_EXPORT DictionaryKeyEncoder : FixedWidthKeyEncoder {
137137
std::shared_ptr<Array> dictionary_;
138138
};
139139

140-
template <typename T>
141-
struct VarLengthKeyEncoder : KeyEncoder {
142-
using Offset = typename T::offset_type;
140+
// Shared implementation of the variable-length row encoding used by the plain
141+
// binary/string encoder and the view encoder. The on-row format is
142+
// [null byte][length : Offset][raw key bytes]; only Decode differs between the
143+
// two, so subclasses override just that.
144+
template <typename VisitType, typename OffsetType>
145+
struct BaseVarLengthKeyEncoder : KeyEncoder {
146+
using Offset = OffsetType;
147+
148+
explicit BaseVarLengthKeyEncoder(std::shared_ptr<DataType> type)
149+
: type_(std::move(type)) {}
143150

144151
void AddLength(const ExecValue& data, int64_t batch_length, int32_t* lengths) override {
145152
if (data.is_array()) {
146153
int64_t i = 0;
147154
ARROW_DCHECK_EQ(data.array.length, batch_length);
148-
VisitArraySpanInline<T>(
155+
VisitArraySpanInline<VisitType>(
149156
data.array,
150157
[&](std::string_view bytes) {
151158
lengths[i++] +=
@@ -155,8 +162,9 @@ struct VarLengthKeyEncoder : KeyEncoder {
155162
} else {
156163
const Scalar& scalar = *data.scalar;
157164
const int32_t buffer_size =
158-
scalar.is_valid ? static_cast<int32_t>(UnboxScalar<T>::Unbox(scalar).size())
159-
: 0;
165+
scalar.is_valid
166+
? static_cast<int32_t>(UnboxScalar<VisitType>::Unbox(scalar).size())
167+
: 0;
160168
for (int64_t i = 0; i < batch_length; i++) {
161169
lengths[i] += kExtraByteForNull + sizeof(Offset) + buffer_size;
162170
}
@@ -185,8 +193,8 @@ struct VarLengthKeyEncoder : KeyEncoder {
185193
};
186194
if (data.is_array()) {
187195
ARROW_DCHECK_EQ(data.length(), batch_length);
188-
VisitArraySpanInline<T>(data.array, handle_next_valid_value,
189-
handle_next_null_value);
196+
VisitArraySpanInline<VisitType>(data.array, handle_next_valid_value,
197+
handle_next_null_value);
190198
} else {
191199
const auto& scalar = data.scalar_as<BaseBinaryScalar>();
192200
if (scalar.is_valid) {
@@ -210,11 +218,24 @@ struct VarLengthKeyEncoder : KeyEncoder {
210218
encoded_ptr += sizeof(Offset);
211219
}
212220

221+
std::shared_ptr<DataType> type_;
222+
};
223+
224+
// Encodes plain binary/string keys, decoding back into the aliasable
225+
// offsets + data buffer layout.
226+
template <typename T>
227+
struct VarLengthKeyEncoder : BaseVarLengthKeyEncoder<T, typename T::offset_type> {
228+
using Base = BaseVarLengthKeyEncoder<T, typename T::offset_type>;
229+
using Base::Base;
230+
using Base::type_;
231+
using typename Base::Offset;
232+
213233
Result<std::shared_ptr<ArrayData>> Decode(uint8_t** encoded_bytes, int32_t length,
214234
MemoryPool* pool) override {
215235
std::shared_ptr<Buffer> null_buf;
216236
int32_t null_count;
217-
ARROW_RETURN_NOT_OK(DecodeNulls(pool, length, encoded_bytes, &null_buf, &null_count));
237+
ARROW_RETURN_NOT_OK(
238+
this->DecodeNulls(pool, length, encoded_bytes, &null_buf, &null_count));
218239

219240
Offset length_sum = 0;
220241
for (int32_t i = 0; i < length; ++i) {
@@ -246,10 +267,17 @@ struct VarLengthKeyEncoder : KeyEncoder {
246267
type_, length, {std::move(null_buf), std::move(offset_buf), std::move(key_buf)},
247268
null_count);
248269
}
270+
};
249271

250-
explicit VarLengthKeyEncoder(std::shared_ptr<DataType> type) : type_(std::move(type)) {}
272+
// Encodes BinaryView/StringView keys in the same row format as the plain
273+
// encoder above. The row copies the key bytes, so Decode builds a fresh view
274+
// array rather than aliasing the input's variadic buffers.
275+
struct ARROW_COMPUTE_EXPORT BinaryViewKeyEncoder
276+
: BaseVarLengthKeyEncoder<BinaryViewType, int32_t> {
277+
using BaseVarLengthKeyEncoder::BaseVarLengthKeyEncoder;
251278

252-
std::shared_ptr<DataType> type_;
279+
Result<std::shared_ptr<ArrayData>> Decode(uint8_t** encoded_bytes, int32_t length,
280+
MemoryPool* pool) override;
253281
};
254282

255283
struct ARROW_COMPUTE_EXPORT NullKeyEncoder : KeyEncoder {

0 commit comments

Comments
 (0)