Skip to content

GH-53159: [C++][Compute] Cache substring and regex matchers across invocations in scalar string kernels #53160

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.

Sign up for GitHub

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

Open
ishaaaan17 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

GH-53159: [C++][Compute] Cache substring and regex matchers across invocations in scalar string kernels #53160

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
240 changes: 156 additions & 84 deletions cpp/src/arrow/compute/kernels/scalar_string_ascii.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1216,8 +1216,6 @@ void StringBoolTransform(KernelContext* ctx, const ExecSpan& batch,
}
}

using MatchSubstringState = OptionsWrapper<MatchSubstringOptions>;

// This is an implementation of the Knuth-Morris-Pratt algorithm
struct PlainSubstringMatcher {
const MatchSubstringOptions& options_;
Expand Down Expand Up @@ -1308,7 +1306,7 @@ struct PlainEndsWithMatcher {

#ifdef ARROW_WITH_RE2
struct RegexSubstringMatcher {
const MatchSubstringOptions& options_;
MatchSubstringOptions options_;
const RE2 regex_match_;

static Result<std::unique_ptr<RegexSubstringMatcher>> Make(
Expand All @@ -1318,11 +1316,11 @@ struct RegexSubstringMatcher {
return matcher;
}

explicit RegexSubstringMatcher(const MatchSubstringOptions& options,
bool is_utf8 = true, bool literal = false)
: options_(options),
explicit RegexSubstringMatcher(MatchSubstringOptions options, bool is_utf8 = true,
bool literal = false)
: options_(std::move(options)),
regex_match_(options_.pattern,
MakeRE2Options(is_utf8, options.ignore_case, literal)) {}
MakeRE2Options(is_utf8, options_.ignore_case, literal)) {}

bool Match(std::string_view current) const {
auto piece = re2::StringPiece(current.data(), current.length());
Expand All @@ -1331,6 +1329,46 @@ struct RegexSubstringMatcher {
};
#endif

struct MatchSubstringState : public KernelState {
MatchSubstringOptions options;

std::unique_ptr<PlainSubstringMatcher> plain_substring_matcher;
std::unique_ptr<PlainStartsWithMatcher> plain_starts_with_matcher;
std::unique_ptr<PlainEndsWithMatcher> plain_ends_with_matcher;
#ifdef ARROW_WITH_RE2
std::unique_ptr<RegexSubstringMatcher> regex_substring_matcher;

enum class LikeKind { kNone, kSubstring, kStartsWith, kEndsWith, kRegex };
LikeKind like_kind = LikeKind::kNone;
std::unique_ptr<MatchSubstringState> like_state;
#endif

explicit MatchSubstringState(MatchSubstringOptions options)
: options(std::move(options)) {}

static Result<std::unique_ptr<KernelState>> Init(KernelContext* ctx,
const KernelInitArgs& args) {
if (auto options = static_cast<const MatchSubstringOptions*>(args.options)) {
return std::make_unique<MatchSubstringState>(*options);
}

return Status::Invalid(
"Attempted to initialize KernelState from null FunctionOptions");
}

static const MatchSubstringOptions& Get(const KernelState& state) {
return ::arrow::internal::checked_cast<const MatchSubstringState&>(state).options;
}

static const MatchSubstringOptions& Get(KernelContext* ctx) {
return Get(*ctx->state());
}

static MatchSubstringState* GetState(KernelContext* ctx) {
return ::arrow::internal::checked_cast<MatchSubstringState*>(ctx->state());
}
};

template <typename Type, typename Matcher>
struct MatchSubstringImpl {
using offset_type = typename Type::offset_type;
Expand Down Expand Up @@ -1361,7 +1399,6 @@ struct MatchSubstringImpl {
template <typename Type, typename Matcher>
struct MatchSubstring {
static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
// TODO Cache matcher across invocations (for regex compilation)
ARROW_ASSIGN_OR_RAISE(auto matcher, Matcher::Make(MatchSubstringState::Get(ctx)));
return MatchSubstringImpl<Type, Matcher>::Exec(ctx, batch, out, matcher.get());
}
Expand All @@ -1371,80 +1408,98 @@ struct MatchSubstring {
template <typename Type>
struct MatchSubstring<Type, RegexSubstringMatcher> {
static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
// TODO Cache matcher across invocations (for regex compilation)
const bool is_utf8 = is_string_or_string_view(batch[0].type()->id());
ARROW_ASSIGN_OR_RAISE(auto matcher, RegexSubstringMatcher::Make(
MatchSubstringState::Get(ctx), is_utf8));
return MatchSubstringImpl<Type, RegexSubstringMatcher>::Exec(ctx, batch, out,
matcher.get());
auto* state = MatchSubstringState::GetState(ctx);
if (!state->regex_substring_matcher) {
const bool is_utf8 = is_string_or_string_view(batch[0].type()->id());
ARROW_ASSIGN_OR_RAISE(state->regex_substring_matcher,
RegexSubstringMatcher::Make(state->options, is_utf8));
}
return MatchSubstringImpl<Type, RegexSubstringMatcher>::Exec(
ctx, batch, out, state->regex_substring_matcher.get());
}
};
#endif

template <typename Type>
struct MatchSubstring<Type, PlainSubstringMatcher> {
static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
auto options = MatchSubstringState::Get(ctx);
if (options.ignore_case) {
auto* state = MatchSubstringState::GetState(ctx);
if (state->options.ignore_case) {
#ifdef ARROW_WITH_RE2
const bool is_utf8 = is_string_or_string_view(batch[0].type()->id());
ARROW_ASSIGN_OR_RAISE(auto matcher, RegexSubstringMatcher::Make(options, is_utf8,
/*literal=*/true));
return MatchSubstringImpl<Type, RegexSubstringMatcher>::Exec(ctx, batch, out,
matcher.get());
if (!state->regex_substring_matcher) {
const bool is_utf8 = is_string_or_string_view(batch[0].type()->id());
ARROW_ASSIGN_OR_RAISE(
state->regex_substring_matcher,
RegexSubstringMatcher::Make(state->options, is_utf8, /*literal=*/true));
}
return MatchSubstringImpl<Type, RegexSubstringMatcher>::Exec(
ctx, batch, out, state->regex_substring_matcher.get());
#else
return Status::NotImplemented("ignore_case requires RE2");
#endif
}
ARROW_ASSIGN_OR_RAISE(auto matcher, PlainSubstringMatcher::Make(options));
return MatchSubstringImpl<Type, PlainSubstringMatcher>::Exec(ctx, batch, out,
matcher.get());
if (!state->plain_substring_matcher) {
ARROW_ASSIGN_OR_RAISE(state->plain_substring_matcher,
PlainSubstringMatcher::Make(state->options));
}
return MatchSubstringImpl<Type, PlainSubstringMatcher>::Exec(
ctx, batch, out, state->plain_substring_matcher.get());
}
};

template <typename Type>
struct MatchSubstring<Type, PlainStartsWithMatcher> {
static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
auto options = MatchSubstringState::Get(ctx);
if (options.ignore_case) {
auto* state = MatchSubstringState::GetState(ctx);
if (state->options.ignore_case) {
#ifdef ARROW_WITH_RE2
MatchSubstringOptions converted_options = options;
converted_options.pattern = "^" + RE2::QuoteMeta(options.pattern);
const bool is_utf8 = is_string_or_string_view(batch[0].type()->id());
ARROW_ASSIGN_OR_RAISE(auto matcher,
RegexSubstringMatcher::Make(converted_options, is_utf8));
return MatchSubstringImpl<Type, RegexSubstringMatcher>::Exec(ctx, batch, out,
matcher.get());
if (!state->regex_substring_matcher) {
MatchSubstringOptions converted_options = state->options;
converted_options.pattern = "^" + RE2::QuoteMeta(state->options.pattern);
const bool is_utf8 = is_string_or_string_view(batch[0].type()->id());
ARROW_ASSIGN_OR_RAISE(state->regex_substring_matcher,
RegexSubstringMatcher::Make(converted_options, is_utf8));
}
return MatchSubstringImpl<Type, RegexSubstringMatcher>::Exec(
ctx, batch, out, state->regex_substring_matcher.get());
#else
return Status::NotImplemented("ignore_case requires RE2");
#endif
}
ARROW_ASSIGN_OR_RAISE(auto matcher, PlainStartsWithMatcher::Make(options));
return MatchSubstringImpl<Type, PlainStartsWithMatcher>::Exec(ctx, batch, out,
matcher.get());
if (!state->plain_starts_with_matcher) {
ARROW_ASSIGN_OR_RAISE(state->plain_starts_with_matcher,
PlainStartsWithMatcher::Make(state->options));
}
return MatchSubstringImpl<Type, PlainStartsWithMatcher>::Exec(
ctx, batch, out, state->plain_starts_with_matcher.get());
}
};

template <typename Type>
struct MatchSubstring<Type, PlainEndsWithMatcher> {
static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
auto options = MatchSubstringState::Get(ctx);
if (options.ignore_case) {
auto* state = MatchSubstringState::GetState(ctx);
if (state->options.ignore_case) {
#ifdef ARROW_WITH_RE2
MatchSubstringOptions converted_options = options;
converted_options.pattern = RE2::QuoteMeta(options.pattern) + "$";
const bool is_utf8 = is_string_or_string_view(batch[0].type()->id());
ARROW_ASSIGN_OR_RAISE(auto matcher,
RegexSubstringMatcher::Make(converted_options, is_utf8));
return MatchSubstringImpl<Type, RegexSubstringMatcher>::Exec(ctx, batch, out,
matcher.get());
if (!state->regex_substring_matcher) {
MatchSubstringOptions converted_options = state->options;
converted_options.pattern = RE2::QuoteMeta(state->options.pattern) + "$";
const bool is_utf8 = is_string_or_string_view(batch[0].type()->id());
ARROW_ASSIGN_OR_RAISE(state->regex_substring_matcher,
RegexSubstringMatcher::Make(converted_options, is_utf8));
}
return MatchSubstringImpl<Type, RegexSubstringMatcher>::Exec(
ctx, batch, out, state->regex_substring_matcher.get());
#else
return Status::NotImplemented("ignore_case requires RE2");
#endif
}
ARROW_ASSIGN_OR_RAISE(auto matcher, PlainEndsWithMatcher::Make(options));
return MatchSubstringImpl<Type, PlainEndsWithMatcher>::Exec(ctx, batch, out,
matcher.get());
if (!state->plain_ends_with_matcher) {
ARROW_ASSIGN_OR_RAISE(state->plain_ends_with_matcher,
PlainEndsWithMatcher::Make(state->options));
}
return MatchSubstringImpl<Type, PlainEndsWithMatcher>::Exec(
ctx, batch, out, state->plain_ends_with_matcher.get());
}
};

Expand Down Expand Up @@ -1536,51 +1591,68 @@ struct MatchLike {
"should only codegen on physical types");

static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
ARROW_ASSIGN_OR_RAISE(
const auto like_constants,
MatchLikeConstants::Instance(
/*is_utf8=*/is_string_or_string_view(batch[0].type()->id())));
auto* original_state = MatchSubstringState::GetState(ctx);
if (!original_state->like_state) {
ARROW_ASSIGN_OR_RAISE(
const auto like_constants,
MatchLikeConstants::Instance(
/*is_utf8=*/is_string_or_string_view(batch[0].type()->id())));

const auto& original_options = original_state->options;
std::string pattern;
bool matched = false;
if (!original_options.ignore_case) {
if ((matched = RE2::FullMatch(original_options.pattern,
like_constants->like_pattern_is_substring_match,
&pattern))) {
original_state->like_kind = MatchSubstringState::LikeKind::kSubstring;
original_state->like_state = std::make_unique<MatchSubstringState>(
MatchSubstringOptions{pattern, false});
} else if ((matched = RE2::FullMatch(original_options.pattern,
like_constants->like_pattern_is_starts_with,
&pattern))) {
original_state->like_kind = MatchSubstringState::LikeKind::kStartsWith;
original_state->like_state = std::make_unique<MatchSubstringState>(
MatchSubstringOptions{pattern, false});
} else if ((matched = RE2::FullMatch(original_options.pattern,
like_constants->like_pattern_is_ends_with,
&pattern))) {
original_state->like_kind = MatchSubstringState::LikeKind::kEndsWith;
original_state->like_state = std::make_unique<MatchSubstringState>(
MatchSubstringOptions{pattern, false});
}
}

auto original_options = MatchSubstringState::Get(ctx);
auto original_state = ctx->state();
if (!matched) {
original_state->like_kind = MatchSubstringState::LikeKind::kRegex;
original_state->like_state =
std::make_unique<MatchSubstringState>(MatchSubstringOptions{
MakeLikeRegex(original_options), original_options.ignore_case});
}
}

ctx->SetState(original_state->like_state.get());
Status status;
std::string pattern;
bool matched = false;
if (!original_options.ignore_case) {
if ((matched = RE2::FullMatch(original_options.pattern,
like_constants->like_pattern_is_substring_match,
&pattern))) {
MatchSubstringOptions converted_options{pattern, original_options.ignore_case};
MatchSubstringState converted_state(converted_options);
ctx->SetState(&converted_state);
switch (original_state->like_kind) {
case MatchSubstringState::LikeKind::kSubstring:
status =
MatchSubstring<PhysicalType, PlainSubstringMatcher>::Exec(ctx, batch, out);
} else if ((matched = RE2::FullMatch(original_options.pattern,
like_constants->like_pattern_is_starts_with,
&pattern))) {
MatchSubstringOptions converted_options{pattern, original_options.ignore_case};
MatchSubstringState converted_state(converted_options);
ctx->SetState(&converted_state);
break;
case MatchSubstringState::LikeKind::kStartsWith:
status =
MatchSubstring<PhysicalType, PlainStartsWithMatcher>::Exec(ctx, batch, out);
} else if ((matched = RE2::FullMatch(original_options.pattern,
like_constants->like_pattern_is_ends_with,
&pattern))) {
MatchSubstringOptions converted_options{pattern, original_options.ignore_case};
MatchSubstringState converted_state(converted_options);
ctx->SetState(&converted_state);
break;
case MatchSubstringState::LikeKind::kEndsWith:
status =
MatchSubstring<PhysicalType, PlainEndsWithMatcher>::Exec(ctx, batch, out);
}
}

if (!matched) {
MatchSubstringOptions converted_options{MakeLikeRegex(original_options),
original_options.ignore_case};
MatchSubstringState converted_state(converted_options);
ctx->SetState(&converted_state);
status = MatchSubstring<PhysicalType, RegexSubstringMatcher>::Exec(ctx, batch, out);
break;
case MatchSubstringState::LikeKind::kRegex:
status =
MatchSubstring<PhysicalType, RegexSubstringMatcher>::Exec(ctx, batch, out);
break;
default:
status = Status::Invalid("Invalid like pattern kind");
break;
}
ctx->SetState(original_state);
return status;
Expand Down
Loading