Skip to content

GH-50615: [C++] Reduce code generation for string kernels #50616

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

Merged
pitrou merged 3 commits into from
Jul 27, 2026
Merged

GH-50615: [C++] Reduce code generation for string kernels #50616

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9d2e900
GH-50615: [C++] Reduce code generation for string kernels
pitrou Jul 23, 2026
5930e92
Address review comments
pitrou Jul 27, 2026
7aed06a
Fix typo :)
pitrou Jul 27, 2026
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
48 changes: 8 additions & 40 deletions cpp/src/arrow/compute/kernels/codegen_internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1314,10 +1314,14 @@ KernelType GenerateTypeAgnosticPrimitive(detail::GetTypeId get_id) {
}
}

// similar to GenerateTypeAgnosticPrimitive, but for base variable binary types
template <template <typename...> class Generator, typename KernelType = ArrayKernelExec,
typename... Args>
KernelType GenerateTypeAgnosticVarBinaryBase(detail::GetTypeId get_id) {
// Similar to GenerateTypeAgnosticPrimitive, but for base variable binary types
//
// Note that we don't offer to generate separate code for String types, because
// the utf8-ness of a type can be retrieved and handled efficiently at runtime.
// This helps cut down on code generation (see GH-50615).
template <template <typename...> class Generator, typename... Args>
auto GenerateTypeAgnosticVarBinaryBase(detail::GetTypeId get_id) {
using KernelType = decltype(&Generator<BinaryType, Args...>::Exec);
switch (get_id.id) {
case Type::BINARY:
case Type::STRING:
Expand All @@ -1331,24 +1335,6 @@ KernelType GenerateTypeAgnosticVarBinaryBase(detail::GetTypeId get_id) {
}
}

// Generate a kernel given a templated functor for binary and string types
template <template <typename...> class Generator, typename... Args>
ArrayKernelExec GenerateVarBinaryToVarBinary(detail::GetTypeId get_id) {
switch (get_id.id) {
case Type::BINARY:
return Generator<BinaryType, Args...>::Exec;
case Type::STRING:
return Generator<StringType, Args...>::Exec;
case Type::LARGE_BINARY:
return Generator<LargeBinaryType, Args...>::Exec;
case Type::LARGE_STRING:
return Generator<LargeStringType, Args...>::Exec;
default:
ARROW_DCHECK(false);
return nullptr;
}
}

// Generate a kernel given a templated functor for base binary types. Generates
// a single kernel for binary/string and large binary/large string. If your kernel
// implementation needs access to the specific type at compile time, please use
Expand All @@ -1370,24 +1356,6 @@ ArrayKernelExec GenerateVarBinaryBase(detail::GetTypeId get_id) {
}
}

// See BaseBinary documentation
template <template <typename...> class Generator, typename Type0, typename... Args>
ArrayKernelExec GenerateVarBinary(detail::GetTypeId get_id) {
switch (get_id.id) {
case Type::BINARY:
return Generator<Type0, BinaryType, Args...>::Exec;
case Type::STRING:
return Generator<Type0, StringType, Args...>::Exec;
case Type::LARGE_BINARY:
return Generator<Type0, LargeBinaryType, Args...>::Exec;
case Type::LARGE_STRING:
return Generator<Type0, LargeStringType, Args...>::Exec;
default:
ARROW_DCHECK(false);
return nullptr;
}
}

// Generate a kernel given a templated functor for binary-view types. Generates a
// single kernel for binary/string-view.
//
Expand Down
3 changes: 1 addition & 2 deletions cpp/src/arrow/compute/kernels/scalar_compare.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -825,8 +825,7 @@ std::shared_ptr<ScalarFunction> MakeScalarMinMax(std::string name, FunctionDoc d
DCHECK_OK(func->AddKernel(std::move(kernel)));
}
for (const auto& ty : BaseBinaryTypes()) {
auto exec =
GenerateTypeAgnosticVarBinaryBase<BinaryScalarMinMax, ArrayKernelExec, Op>(ty);
auto exec = GenerateTypeAgnosticVarBinaryBase<BinaryScalarMinMax, Op>(ty);
ScalarKernel kernel{KernelSignature::Make({ty}, ty, /*is_varargs=*/true), exec,
MinMaxState::Init};
kernel.null_handling = NullHandling::COMPUTED_NO_PREALLOCATE;
Expand Down
2 changes: 1 addition & 1 deletion cpp/src/arrow/compute/kernels/scalar_if_else.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1328,7 +1328,7 @@ void AddBinaryIfElseKernels(const std::shared_ptr<IfElseFunction>& scalar_functi
const std::vector<std::shared_ptr<DataType>>& types) {
for (auto&& type : types) {
auto exec =
internal::GenerateTypeAgnosticVarBinaryBase<ResolveIfElseExec, ArrayKernelExec,
internal::GenerateTypeAgnosticVarBinaryBase<ResolveIfElseExec,
/*AllocateMem=*/std::true_type>(
*type);
// cond array needs to be boolean always
Expand Down
248 changes: 147 additions & 101 deletions cpp/src/arrow/compute/kernels/scalar_string_ascii.cc
View file
Open in desktop

Large diffs are not rendered by default.

45 changes: 29 additions & 16 deletions cpp/src/arrow/compute/kernels/scalar_string_internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

#include "arrow/compute/api_scalar.h"
#include "arrow/compute/kernels/common_internal.h"
#include "arrow/type_traits.h"

namespace arrow {
namespace compute {
Expand Down Expand Up @@ -68,10 +69,13 @@ static int64_t GetVarBinaryValuesLength(const ArraySpan& span) {
///
/// and returns the number of codeunits of the `output` sequence or a negative
/// value if an invalid input sequence is detected.
template <typename Type, typename StringTransform>
template <typename PhysicalType, typename StringTransform>
struct StringTransformExecBase {
using offset_type = typename Type::offset_type;
using ArrayType = typename TypeTraits<Type>::ArrayType;
using offset_type = typename PhysicalType::offset_type;
using ArrayType = typename TypeTraits<PhysicalType>::ArrayType;

static_assert(!is_string_or_string_view(PhysicalType::type_id),
"should only codegen on physical types");

static Status Execute(KernelContext* ctx, StringTransform* transform,
const ExecSpan& batch, ExecResult* out) {
Expand Down Expand Up @@ -121,9 +125,11 @@ struct StringTransformExecBase {
}
};

template <typename Type, typename StringTransform>
struct StringTransformExec : public StringTransformExecBase<Type, StringTransform> {
using StringTransformExecBase<Type, StringTransform>::Execute;
template <typename Type, typename StringTransform,
typename PhysicalType = typename Type::PhysicalType>
struct StringTransformExec
: public StringTransformExecBase<PhysicalType, StringTransform> {
using StringTransformExecBase<PhysicalType, StringTransform>::Execute;

static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
StringTransform transform;
Expand All @@ -132,11 +138,12 @@ struct StringTransformExec : public StringTransformExecBase<Type, StringTransfor
}
};

template <typename Type, typename StringTransform>
template <typename Type, typename StringTransform,
typename PhysicalType = typename Type::PhysicalType>
struct StringTransformExecWithState
: public StringTransformExecBase<Type, StringTransform> {
: public StringTransformExecBase<PhysicalType, StringTransform> {
using State = typename StringTransform::State;
using StringTransformExecBase<Type, StringTransform>::Execute;
using StringTransformExecBase<PhysicalType, StringTransform>::Execute;

static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
StringTransform transform(State::Get(ctx));
Expand All @@ -151,7 +158,7 @@ void MakeUnaryStringBatchKernel(
MemAllocation::type mem_allocation = MemAllocation::PREALLOCATE) {
auto func = std::make_shared<ScalarFunction>(name, Arity::Unary(), std::move(doc));
for (const auto& ty : StringTypes()) {
auto exec = GenerateVarBinaryToVarBinary<ExecFunctor>(ty);
auto exec = GenerateTypeAgnosticVarBinaryBase<ExecFunctor>(ty);
ScalarKernel kernel{{ty}, ty, std::move(exec)};
kernel.mem_allocation = mem_allocation;
ARROW_DCHECK_OK(func->AddKernel(std::move(kernel)));
Expand Down Expand Up @@ -216,6 +223,9 @@ static inline FunctionDoc StringClassifyDoc(std::string class_summary,

template <typename Type, typename Predicate>
struct StringPredicateFunctor {
static_assert(!is_string_or_string_view(Type::type_id),
"should only codegen on physical types");

static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
Status st = Status::OK();
EnsureUtf8LookupTablesFilled();
Expand All @@ -237,7 +247,7 @@ void AddUnaryStringPredicate(std::string name, FunctionRegistry* registry,
FunctionDoc doc) {
auto func = std::make_shared<ScalarFunction>(name, Arity::Unary(), std::move(doc));
for (const auto& ty : StringTypes()) {
auto exec = GenerateVarBinaryToVarBinary<StringPredicateFunctor, Predicate>(ty);
auto exec = GenerateTypeAgnosticVarBinaryBase<StringPredicateFunctor, Predicate>(ty);
ARROW_DCHECK_OK(func->AddKernel({ty}, boolean(), std::move(exec)));
}
ARROW_DCHECK_OK(registry->AddFunction(std::move(func)));
Expand Down Expand Up @@ -281,7 +291,7 @@ struct ReplaceStringSliceTransformBase : public StringTransformBase {
template <typename Options>
struct StringSplitFinderBase {
virtual ~StringSplitFinderBase() = default;
virtual Status PreExec(const Options& options) { return Status::OK(); }
virtual Status PreExec(const Options& options, bool is_utf8) { return Status::OK(); }

// Derived classes should also define these methods:
// static bool Find(const uint8_t* begin, const uint8_t* end,
Expand Down Expand Up @@ -319,8 +329,9 @@ struct StringSplitExec {
}

Status Execute(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
const bool is_utf8 = is_string_or_string_view(batch[0].type()->id());
SplitFinder finder;
RETURN_NOT_OK(finder.PreExec(options));
RETURN_NOT_OK(finder.PreExec(options, is_utf8));
// TODO(wesm): refactor to not require creating ArrayData
const ArrayType input(batch[0].array.ToArrayData());

Expand All @@ -347,9 +358,11 @@ struct StringSplitExec {
*list_offsets++ = static_cast<list_offset_type>(builder.length());
}
// Assign string array to list child data
std::shared_ptr<Array> string_array;
RETURN_NOT_OK(builder.Finish(&string_array));
output_list->child_data.push_back(string_array->data());
ARROW_ASSIGN_OR_RAISE(auto physical_array, builder.Finish());
// We got the physical type (e.g. binary instead of utf8), need to patch it
auto child_data = physical_array->data()->Copy();
child_data->type = batch[0].type()->GetSharedPtr();
output_list->child_data.push_back(std::move(child_data));
return Status::OK();
}

Expand Down
Loading
Loading