Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 100 additions & 20 deletions cpp/src/arrow/compute/kernels/scalar_string_ascii.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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)));
Expand Down Expand Up @@ -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();
}
};
Expand Down Expand Up @@ -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.
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(),
Expand All @@ -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)));
}
{
Expand All @@ -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)));
}
{
Expand All @@ -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
Expand All @@ -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)));
}
{
Expand All @@ -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
Expand All @@ -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_;

Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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));
Expand All @@ -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));
Expand Down Expand Up @@ -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(
Expand All @@ -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) {
Expand Down Expand Up @@ -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));
Expand All @@ -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));
Expand Down
46 changes: 40 additions & 6 deletions cpp/src/arrow/compute/kernels/scalar_string_benchmark.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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) {
Expand All @@ -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) {
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);
Expand Down Expand Up @@ -242,6 +274,8 @@ BENCHMARK(AsciiLower);
BENCHMARK(AsciiUpper);
BENCHMARK(IsAlphaNumericAscii);
BENCHMARK(MatchSubstring);
BENCHMARK(MatchSubstringView);
BENCHMARK(MatchSubstringViewCast);
BENCHMARK(SplitPattern);
BENCHMARK(TrimSingleAscii);
BENCHMARK(TrimManyAscii);
Expand Down
35 changes: 28 additions & 7 deletions cpp/src/arrow/compute/kernels/scalar_string_internal.h
Original file line number Diff line number Diff line change
Expand Up @@ -220,14 +220,33 @@ struct StringPredicateFunctor {
Status st = Status::OK();
EnsureUtf8LookupTablesFilled();
const ArraySpan& input = batch[0].array;
ArrayIterator<Type> input_it(input);
ArraySpan* out_arr = out->array_span_mutable();
::arrow::internal::GenerateBitsUnrolled(
out_arr->buffers[1].data, out_arr->offset, input.length, [&]() -> bool {
std::string_view val = input_it();
return Predicate::Call(ctx, reinterpret_cast<const uint8_t*>(val.data()),
val.size(), &st);
});
if constexpr (is_binary_view_like_type<Type>::value) {
// Skip null slots: a view's null header is not validated and may carry a
// bogus buffer_index/offset that decoding would dereference. (The base
// binary/string path below can safely read null slots.)
FirstTimeBitmapWriter writer(out_arr->buffers[1].data, out_arr->offset,
input.length);
VisitArrayValuesInline<Type>(
input,
[&](std::string_view val) {
if (Predicate::Call(ctx, reinterpret_cast<const uint8_t*>(val.data()),
val.size(), &st)) {
writer.Set();
}
writer.Next();
},
[&]() { writer.Next(); });
writer.Finish();
} else {
ArrayIterator<Type> input_it(input);
::arrow::internal::GenerateBitsUnrolled(
out_arr->buffers[1].data, out_arr->offset, input.length, [&]() -> bool {
std::string_view val = input_it();
return Predicate::Call(ctx, reinterpret_cast<const uint8_t*>(val.data()),
val.size(), &st);
});
}
return st;
}
};
Expand All @@ -240,6 +259,8 @@ void AddUnaryStringPredicate(std::string name, FunctionRegistry* registry,
auto exec = GenerateVarBinaryToVarBinary<StringPredicateFunctor, Predicate>(ty);
ARROW_DCHECK_OK(func->AddKernel({ty}, boolean(), std::move(exec)));
}
ARROW_DCHECK_OK(func->AddKernel(
{utf8_view()}, boolean(), StringPredicateFunctor<StringViewType, Predicate>::Exec));
ARROW_DCHECK_OK(registry->AddFunction(std::move(func)));
}

Expand Down
Loading
Loading