-
Notifications
You must be signed in to change notification settings - Fork 4.2k
GH-50478: [C++][Compute] Support string_view/binary_view in scalar string predicate kernels #50479
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
19583cc
764dff6
8082d31
098c9ee
26453db
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -889,6 +889,12 @@ void AddAsciiStringLength(FunctionRegistry* registry) { | |
| ty); | ||
| DCHECK_OK(func->AddKernel({ty}, int64(), std::move(exec))); | ||
| } | ||
| // View element length is int32-sized, so both view types emit int32. | ||
| for (const auto& ty : BinaryViewTypes()) { | ||
| auto exec = GenerateVarBinaryViewBase<applicator::ScalarUnaryNotNull, Int32Type, | ||
| BinaryLength>(*ty); | ||
| DCHECK_OK(func->AddKernel({ty}, int32(), std::move(exec))); | ||
| } | ||
| DCHECK_OK(func->AddKernel({InputType(Type::FIXED_SIZE_BINARY)}, int32(), | ||
| BinaryLength::FixedSizeExec)); | ||
| DCHECK_OK(registry->AddFunction(std::move(func))); | ||
|
|
@@ -1337,27 +1343,47 @@ struct RegexSubstringMatcher { | |
|
|
||
| template <typename Type, typename Matcher> | ||
| struct MatchSubstringImpl { | ||
| using offset_type = typename Type::offset_type; | ||
|
|
||
| static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out, | ||
| const Matcher* matcher) { | ||
| StringBoolTransform<Type>( | ||
| ctx, batch, | ||
| [&matcher](const void* raw_offsets, const uint8_t* data, int64_t length, | ||
| int64_t output_offset, uint8_t* output) { | ||
| const offset_type* offsets = reinterpret_cast<const offset_type*>(raw_offsets); | ||
| FirstTimeBitmapWriter bitmap_writer(output, output_offset, length); | ||
| for (int64_t i = 0; i < length; ++i) { | ||
| const char* current_data = reinterpret_cast<const char*>(data + offsets[i]); | ||
| int64_t current_length = offsets[i + 1] - offsets[i]; | ||
| if (matcher->Match(std::string_view(current_data, current_length))) { | ||
| if constexpr (is_binary_view_like_type<Type>::value) { | ||
| // Views have no packed offset buffer, so evaluate per element. Null slots are | ||
| // skipped: a view's null header is not validated and may carry a bogus | ||
| // buffer_index/offset that decoding would dereference. | ||
| const ArraySpan& input = batch[0].array; | ||
| ArraySpan* out_arr = out->array_span_mutable(); | ||
| FirstTimeBitmapWriter bitmap_writer(out_arr->buffers[1].data, out_arr->offset, | ||
| input.length); | ||
| VisitArrayValuesInline<Type>( | ||
| input, | ||
| [&](std::string_view val) { | ||
| if (matcher->Match(val)) { | ||
| bitmap_writer.Set(); | ||
| } | ||
| bitmap_writer.Next(); | ||
| } | ||
| bitmap_writer.Finish(); | ||
| }, | ||
| out); | ||
| }, | ||
| [&]() { bitmap_writer.Next(); }); | ||
| bitmap_writer.Finish(); | ||
| } else { | ||
| using offset_type = typename Type::offset_type; | ||
| StringBoolTransform<Type>( | ||
| ctx, batch, | ||
| [&matcher](const void* raw_offsets, const uint8_t* data, int64_t length, | ||
| int64_t output_offset, uint8_t* output) { | ||
| const offset_type* offsets = | ||
| reinterpret_cast<const offset_type*>(raw_offsets); | ||
| FirstTimeBitmapWriter bitmap_writer(output, output_offset, length); | ||
| for (int64_t i = 0; i < length; ++i) { | ||
| const char* current_data = reinterpret_cast<const char*>(data + offsets[i]); | ||
| int64_t current_length = offsets[i + 1] - offsets[i]; | ||
| if (matcher->Match(std::string_view(current_data, current_length))) { | ||
| bitmap_writer.Set(); | ||
| } | ||
| bitmap_writer.Next(); | ||
| } | ||
| bitmap_writer.Finish(); | ||
| }, | ||
| out); | ||
| } | ||
| return Status::OK(); | ||
| } | ||
| }; | ||
|
|
@@ -1611,6 +1637,19 @@ const FunctionDoc match_like_doc( | |
| {"strings"}, "MatchSubstringOptions", /*options_required=*/true); | ||
| #endif | ||
|
|
||
| // Register the view kernels for a match-substring-style predicate. Registered per | ||
| // view type (not via GenerateVarBinaryViewBase) so utf8_view -> StringViewType keeps | ||
| // the is_utf8 distinction used by ignore_case/regex folding. | ||
|
Comment on lines
+1640
to
+1642
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I've opened #50615 as it's frankly a pity to duplicate kernel code just for the is_utf8 flag. |
||
| template <template <typename...> class ExecTemplate, typename... Matcher> | ||
| void AddMatchSubstringViewKernels(ScalarFunction* func) { | ||
| DCHECK_OK(func->AddKernel({binary_view()}, boolean(), | ||
| ExecTemplate<BinaryViewType, Matcher...>::Exec, | ||
| MatchSubstringState::Init)); | ||
| DCHECK_OK(func->AddKernel({utf8_view()}, boolean(), | ||
| ExecTemplate<StringViewType, Matcher...>::Exec, | ||
| MatchSubstringState::Init)); | ||
| } | ||
|
|
||
| void AddAsciiStringMatchSubstring(FunctionRegistry* registry) { | ||
| { | ||
| auto func = std::make_shared<ScalarFunction>("match_substring", Arity::Unary(), | ||
|
|
@@ -1620,6 +1659,7 @@ void AddAsciiStringMatchSubstring(FunctionRegistry* registry) { | |
| DCHECK_OK( | ||
| func->AddKernel({ty}, boolean(), std::move(exec), MatchSubstringState::Init)); | ||
| } | ||
| AddMatchSubstringViewKernels<MatchSubstring, PlainSubstringMatcher>(func.get()); | ||
| DCHECK_OK(registry->AddFunction(std::move(func))); | ||
| } | ||
| { | ||
|
|
@@ -1631,6 +1671,7 @@ void AddAsciiStringMatchSubstring(FunctionRegistry* registry) { | |
| DCHECK_OK( | ||
| func->AddKernel({ty}, boolean(), std::move(exec), MatchSubstringState::Init)); | ||
| } | ||
| AddMatchSubstringViewKernels<MatchSubstring, PlainStartsWithMatcher>(func.get()); | ||
| DCHECK_OK(registry->AddFunction(std::move(func))); | ||
| } | ||
| { | ||
|
|
@@ -1641,6 +1682,7 @@ void AddAsciiStringMatchSubstring(FunctionRegistry* registry) { | |
| DCHECK_OK( | ||
| func->AddKernel({ty}, boolean(), std::move(exec), MatchSubstringState::Init)); | ||
| } | ||
| AddMatchSubstringViewKernels<MatchSubstring, PlainEndsWithMatcher>(func.get()); | ||
| DCHECK_OK(registry->AddFunction(std::move(func))); | ||
| } | ||
| #ifdef ARROW_WITH_RE2 | ||
|
|
@@ -1652,6 +1694,7 @@ void AddAsciiStringMatchSubstring(FunctionRegistry* registry) { | |
| DCHECK_OK( | ||
| func->AddKernel({ty}, boolean(), std::move(exec), MatchSubstringState::Init)); | ||
| } | ||
| AddMatchSubstringViewKernels<MatchSubstring, RegexSubstringMatcher>(func.get()); | ||
| DCHECK_OK(registry->AddFunction(std::move(func))); | ||
| } | ||
| { | ||
|
|
@@ -1662,6 +1705,7 @@ void AddAsciiStringMatchSubstring(FunctionRegistry* registry) { | |
| DCHECK_OK( | ||
| func->AddKernel({ty}, boolean(), std::move(exec), MatchSubstringState::Init)); | ||
| } | ||
| AddMatchSubstringViewKernels<MatchLike>(func.get()); | ||
| DCHECK_OK(registry->AddFunction(std::move(func))); | ||
| } | ||
| #endif | ||
|
|
@@ -1670,6 +1714,17 @@ void AddAsciiStringMatchSubstring(FunctionRegistry* registry) { | |
| // ---------------------------------------------------------------------- | ||
| // Substring find - lfind/index/etc. | ||
|
|
||
| // The output offset type for find/count kernels. Like TypeTraits<T>::OffsetType, | ||
| // but also defined for the view types, whose per-element length is int32-sized. | ||
| template <typename T, typename Enable = void> | ||
| struct StringOffsetType { | ||
| using type = typename TypeTraits<T>::OffsetType; | ||
| }; | ||
| template <typename T> | ||
| struct StringOffsetType<T, enable_if_binary_view_like<T>> { | ||
| using type = Int32Type; | ||
| }; | ||
|
|
||
| struct FindSubstring { | ||
| const PlainSubstringMatcher matcher_; | ||
|
|
||
|
|
@@ -1716,7 +1771,7 @@ struct FindSubstringRegex { | |
|
|
||
| template <typename InputType> | ||
| struct FindSubstringExec { | ||
| using OffsetType = typename TypeTraits<InputType>::OffsetType; | ||
| using OffsetType = typename StringOffsetType<InputType>::type; | ||
| static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { | ||
| const MatchSubstringOptions& options = MatchSubstringState::Get(ctx); | ||
| if (options.ignore_case) { | ||
|
|
@@ -1746,7 +1801,7 @@ const FunctionDoc find_substring_doc( | |
| #ifdef ARROW_WITH_RE2 | ||
| template <typename InputType> | ||
| struct FindSubstringRegexExec { | ||
| using OffsetType = typename TypeTraits<InputType>::OffsetType; | ||
| using OffsetType = typename StringOffsetType<InputType>::type; | ||
| static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { | ||
| const MatchSubstringOptions& options = MatchSubstringState::Get(ctx); | ||
| ARROW_ASSIGN_OR_RAISE(auto matcher, FindSubstringRegex::Make(options, false)); | ||
|
|
@@ -1774,6 +1829,13 @@ void AddAsciiStringFindSubstring(FunctionRegistry* registry) { | |
| GenerateVarBinaryToVarBinary<FindSubstringExec>(ty), | ||
| MatchSubstringState::Init)); | ||
| } | ||
| // Per view type to keep the is_utf8 distinction; view length fits int32. | ||
| DCHECK_OK(func->AddKernel({binary_view()}, int32(), | ||
| FindSubstringExec<BinaryViewType>::Exec, | ||
| MatchSubstringState::Init)); | ||
| DCHECK_OK(func->AddKernel({utf8_view()}, int32(), | ||
| FindSubstringExec<StringViewType>::Exec, | ||
| MatchSubstringState::Init)); | ||
| DCHECK_OK(func->AddKernel({InputType(Type::FIXED_SIZE_BINARY)}, int32(), | ||
| FindSubstringExec<FixedSizeBinaryType>::Exec, | ||
| MatchSubstringState::Init)); | ||
|
|
@@ -1789,6 +1851,12 @@ void AddAsciiStringFindSubstring(FunctionRegistry* registry) { | |
| GenerateVarBinaryToVarBinary<FindSubstringRegexExec>(ty), | ||
| MatchSubstringState::Init)); | ||
| } | ||
| DCHECK_OK(func->AddKernel({binary_view()}, int32(), | ||
| FindSubstringRegexExec<BinaryViewType>::Exec, | ||
| MatchSubstringState::Init)); | ||
| DCHECK_OK(func->AddKernel({utf8_view()}, int32(), | ||
| FindSubstringRegexExec<StringViewType>::Exec, | ||
| MatchSubstringState::Init)); | ||
| DCHECK_OK(func->AddKernel({InputType(Type::FIXED_SIZE_BINARY)}, int32(), | ||
| FindSubstringRegexExec<FixedSizeBinaryType>::Exec, | ||
| MatchSubstringState::Init)); | ||
|
|
@@ -1862,7 +1930,7 @@ struct CountSubstringRegex { | |
|
|
||
| template <typename InputType> | ||
| struct CountSubstringRegexExec { | ||
| using OffsetType = typename TypeTraits<InputType>::OffsetType; | ||
| using OffsetType = typename StringOffsetType<InputType>::type; | ||
| static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { | ||
| const MatchSubstringOptions& options = MatchSubstringState::Get(ctx); | ||
| ARROW_ASSIGN_OR_RAISE( | ||
|
|
@@ -1876,7 +1944,7 @@ struct CountSubstringRegexExec { | |
|
|
||
| template <typename InputType> | ||
| struct CountSubstringExec { | ||
| using OffsetType = typename TypeTraits<InputType>::OffsetType; | ||
| using OffsetType = typename StringOffsetType<InputType>::type; | ||
| static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { | ||
| const MatchSubstringOptions& options = MatchSubstringState::Get(ctx); | ||
| if (options.ignore_case) { | ||
|
|
@@ -1923,6 +1991,12 @@ void AddAsciiStringCountSubstring(FunctionRegistry* registry) { | |
| GenerateVarBinaryToVarBinary<CountSubstringExec>(ty), | ||
| MatchSubstringState::Init)); | ||
| } | ||
| DCHECK_OK(func->AddKernel({binary_view()}, int32(), | ||
| CountSubstringExec<BinaryViewType>::Exec, | ||
| MatchSubstringState::Init)); | ||
| DCHECK_OK(func->AddKernel({utf8_view()}, int32(), | ||
| CountSubstringExec<StringViewType>::Exec, | ||
| MatchSubstringState::Init)); | ||
| DCHECK_OK(func->AddKernel({InputType(Type::FIXED_SIZE_BINARY)}, int32(), | ||
| CountSubstringExec<FixedSizeBinaryType>::Exec, | ||
| MatchSubstringState::Init)); | ||
|
|
@@ -1938,6 +2012,12 @@ void AddAsciiStringCountSubstring(FunctionRegistry* registry) { | |
| GenerateVarBinaryToVarBinary<CountSubstringRegexExec>(ty), | ||
| MatchSubstringState::Init)); | ||
| } | ||
| DCHECK_OK(func->AddKernel({binary_view()}, int32(), | ||
| CountSubstringRegexExec<BinaryViewType>::Exec, | ||
| MatchSubstringState::Init)); | ||
| DCHECK_OK(func->AddKernel({utf8_view()}, int32(), | ||
| CountSubstringRegexExec<StringViewType>::Exec, | ||
| MatchSubstringState::Init)); | ||
| DCHECK_OK(func->AddKernel({InputType(Type::FIXED_SIZE_BINARY)}, int32(), | ||
| CountSubstringRegexExec<FixedSizeBinaryType>::Exec, | ||
| MatchSubstringState::Init)); | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -22,6 +22,7 @@ | |
| #include "arrow/array.h" | ||
| #include "arrow/chunked_array.h" | ||
| #include "arrow/compute/api_scalar.h" | ||
| #include "arrow/compute/cast.h" | ||
| #include "arrow/compute/exec.h" | ||
| #include "arrow/datum.h" | ||
| #include "arrow/testing/gtest_util.h" | ||
|
|
@@ -39,25 +40,34 @@ namespace compute { | |
|
|
||
| constexpr auto kSeed = 0x94378165; | ||
|
|
||
| static void UnaryStringBenchmark(benchmark::State& state, const std::string& func_name, | ||
| const FunctionOptions* options = nullptr) { | ||
| // Shared dataset for the string benchmarks: ~1M only-ASCII values. | ||
| static std::shared_ptr<Array> MakeBenchmarkStrings() { | ||
| const int64_t array_length = 1 << 20; | ||
| const int64_t value_min_size = 0; | ||
| const int64_t value_max_size = 32; | ||
| const double null_probability = 0.01; | ||
| random::RandomArrayGenerator rng(kSeed); | ||
| return rng.String(array_length, value_min_size, value_max_size, null_probability); | ||
| } | ||
|
|
||
| // NOTE: this produces only-Ascii data | ||
| auto values = | ||
| rng.String(array_length, value_min_size, value_max_size, null_probability); | ||
| static void UnaryStringBenchmark(benchmark::State& state, const std::string& func_name, | ||
| const FunctionOptions* options = nullptr, | ||
| const std::shared_ptr<DataType>& input_type = utf8()) { | ||
| std::shared_ptr<Array> values = MakeBenchmarkStrings(); | ||
| const int64_t array_length = values->length(); | ||
| // Report throughput based on the raw string bytes, before any conversion. | ||
| const int64_t data_nbytes = values->data()->buffers[2]->size(); | ||
| if (input_type->id() != Type::STRING) { | ||
| values = Cast(*values, input_type).ValueOrDie(); | ||
| } | ||
| // Make sure lookup tables are initialized before measuring | ||
| ABORT_NOT_OK(CallFunction(func_name, {values}, options)); | ||
|
|
||
| for (auto _ : state) { | ||
| ABORT_NOT_OK(CallFunction(func_name, {values}, options)); | ||
| } | ||
| state.SetItemsProcessed(state.iterations() * array_length); | ||
| state.SetBytesProcessed(state.iterations() * values->data()->buffers[2]->size()); | ||
| state.SetBytesProcessed(state.iterations() * data_nbytes); | ||
| } | ||
|
|
||
| static void AsciiLower(benchmark::State& state) { | ||
|
|
@@ -77,6 +87,28 @@ static void MatchSubstring(benchmark::State& state) { | |
| UnaryStringBenchmark(state, "match_substring", &options); | ||
| } | ||
|
|
||
| // Predicate run directly on a string_view array (the path added by this PR). | ||
| static void MatchSubstringView(benchmark::State& state) { | ||
| MatchSubstringOptions options("abac"); | ||
| UnaryStringBenchmark(state, "match_substring", &options, utf8_view()); | ||
| } | ||
|
|
||
| // Baseline: the pre-PR workaround of casting the view column to utf8 first. The | ||
| // gap vs. MatchSubstringView is the cast that direct view support avoids. | ||
| static void MatchSubstringViewCast(benchmark::State& state) { | ||
|
Comment on lines
+96
to
+98
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We don't care about benchmarking this IMHO. It was useful to judge this PR's usefulness, but it's not useful for measuring Arrow C++ performance. |
||
| std::shared_ptr<Array> utf8_values = MakeBenchmarkStrings(); | ||
| const int64_t array_length = utf8_values->length(); | ||
| const int64_t data_nbytes = utf8_values->data()->buffers[2]->size(); | ||
| std::shared_ptr<Array> view_values = Cast(*utf8_values, utf8_view()).ValueOrDie(); | ||
| MatchSubstringOptions options("abac"); | ||
| for (auto _ : state) { | ||
| ASSIGN_OR_ABORT(auto casted, Cast(*view_values, utf8())); | ||
| ABORT_NOT_OK(CallFunction("match_substring", {casted}, &options)); | ||
| } | ||
| state.SetItemsProcessed(state.iterations() * array_length); | ||
| state.SetBytesProcessed(state.iterations() * data_nbytes); | ||
| } | ||
|
|
||
| static void SplitPattern(benchmark::State& state) { | ||
| SplitPatternOptions options("a"); | ||
| UnaryStringBenchmark(state, "split_pattern", &options); | ||
|
|
@@ -242,6 +274,8 @@ BENCHMARK(AsciiLower); | |
| BENCHMARK(AsciiUpper); | ||
| BENCHMARK(IsAlphaNumericAscii); | ||
| BENCHMARK(MatchSubstring); | ||
| BENCHMARK(MatchSubstringView); | ||
| BENCHMARK(MatchSubstringViewCast); | ||
| BENCHMARK(SplitPattern); | ||
| BENCHMARK(TrimSingleAscii); | ||
| BENCHMARK(TrimManyAscii); | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Note: I think ideally we would also use
VisitArrayValuesInlinefor other binary types, because the matcher may be expensive and it's better to skip it on null entries.(for example, if the array has 99% nulls, it would be wasteful to run a regexp matcher on each underlying empty slot)