Skip to content

GH-40207: [C++] TakeCC: Concatenate only once and delegate to TakeAA instead of TakeCA#40206

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 8 commits into from
Feb 28, 2024
Merged

GH-40207: [C++] TakeCC: Concatenate only once and delegate to TakeAA instead of TakeCA #40206

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
c7ca222
Also encode the return type of Take*() specializations
felipecrv Feb 22, 2024
367ff9c
TakeChunked benchmarks from #13857
felipecrv Feb 23, 2024
a9a8304
TakeCCC: Concatenate only once and delegate to TakeAAA instead of Tak…
felipecrv Feb 23, 2024
140e255
benchmarks: Rename TakeChunked to TakeChunkedChunked
felipecrv Feb 23, 2024
40d585b
Improve the TakeChunked benchmarks
felipecrv Feb 23, 2024
954919c
Fix benchmarks after code review
felipecrv Feb 27, 2024
2486e99
Consistently order the take benchmarks definitions and registrations
felipecrv Feb 27, 2024
8643da4
Add more benchmarks for Take()
felipecrv Feb 28, 2024
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
179 changes: 167 additions & 12 deletions cpp/src/arrow/compute/kernels/vector_selection_benchmark.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -115,6 +115,24 @@ struct TakeBenchmark {
indices_have_nulls(indices_have_nulls),
monotonic_indices(monotonic_indices) {}

static constexpr int kStringMinLength = 0;
static constexpr int kStringMaxLength = 32;
static constexpr int kByteWidthRange = 2;

template <typename GenChunk>
std::shared_ptr<ChunkedArray> GenChunkedArray(int64_t num_chunks,
GenChunk&& gen_chunk) {
const int64_t chunk_length =
std::llround(args.size / static_cast<double>(num_chunks));
ArrayVector chunks;
for (int64_t i = 0; i < num_chunks; ++i) {
const int64_t fitting_chunk_length =
std::min(chunk_length, args.size - i * chunk_length);
chunks.push_back(gen_chunk(fitting_chunk_length));
}
return std::make_shared<ChunkedArray>(std::move(chunks));
}

void Int64() {
auto values = rand.Int64(args.size, -100, 100, args.null_proportion);
Bench(values);
Expand All@@ -129,19 +147,43 @@ struct TakeBenchmark {
}

void FixedSizeBinary() {
const int32_t byte_width = static_cast<int32_t>(state.range(2));
const auto byte_width = static_cast<int32_t>(state.range(kByteWidthRange));
auto values = rand.FixedSizeBinary(args.size, byte_width, args.null_proportion);
Bench(values);
state.counters["byte_width"] = byte_width;
}

void String() {
int32_t string_min_length = 0, string_max_length = 32;
auto values = std::static_pointer_cast<StringArray>(rand.String(
args.size, string_min_length, string_max_length, args.null_proportion));
auto values = std::static_pointer_cast<StringArray>(
rand.String(args.size, kStringMinLength, kStringMaxLength, args.null_proportion));
Bench(values);
}

void ChunkedInt64(int64_t num_chunks, bool chunk_indices_too) {
auto chunked_array = GenChunkedArray(num_chunks, [this](int64_t chunk_length) {
return rand.Int64(chunk_length, -100, 100, args.null_proportion);
});
BenchChunked(chunked_array, chunk_indices_too);
}

void ChunkedFSB(int64_t num_chunks, bool chunk_indices_too) {
const auto byte_width = static_cast<int32_t>(state.range(kByteWidthRange));
auto chunked_array =
GenChunkedArray(num_chunks, [this, byte_width](int64_t chunk_length) {
return rand.FixedSizeBinary(chunk_length, byte_width, args.null_proportion);
});
BenchChunked(chunked_array, chunk_indices_too);
state.counters["byte_width"] = byte_width;
}

void ChunkedString(int64_t num_chunks, bool chunk_indices_too) {
auto chunked_array = GenChunkedArray(num_chunks, [this](int64_t chunk_length) {
return std::static_pointer_cast<StringArray>(rand.String(
chunk_length, kStringMinLength, kStringMaxLength, args.null_proportion));
});
BenchChunked(chunked_array, chunk_indices_too);
}

void Bench(const std::shared_ptr<Array>& values) {
double indices_null_proportion = indices_have_nulls ? args.null_proportion : 0;
auto indices =
Expand All@@ -158,6 +200,40 @@ struct TakeBenchmark {
}
state.SetItemsProcessed(state.iterations() * values->length());
}

void BenchChunked(const std::shared_ptr<ChunkedArray>& values, bool chunk_indices_too) {
double indices_null_proportion = indices_have_nulls ? args.null_proportion : 0;
auto indices =
rand.Int32(values->length(), 0, static_cast<int32_t>(values->length() - 1),
indices_null_proportion);

if (monotonic_indices) {
auto arg_sorter = *SortIndices(*indices);
indices = *Take(*indices, *arg_sorter);
}
std::shared_ptr<ChunkedArray> chunked_indices;
if (chunk_indices_too) {
std::vector<std::shared_ptr<Array>> indices_chunks;
int64_t offset = 0;
for (int i = 0; i < values->num_chunks(); ++i) {
auto chunk = indices->Slice(offset, values->chunk(i)->length());
indices_chunks.push_back(std::move(chunk));
offset += values->chunk(i)->length();
}
chunked_indices = std::make_shared<ChunkedArray>(std::move(indices_chunks));
}

if (chunk_indices_too) {
for (auto _ : state) {
ABORT_NOT_OK(Take(values, chunked_indices).status());
}
} else {
for (auto _ : state) {
ABORT_NOT_OK(Take(values, indices).status());
}
}
Comment thread
pitrou marked this conversation as resolved.
Outdated
state.SetItemsProcessed(state.iterations() * values->length());
}
};

struct FilterBenchmark {
Expand DownExpand Up@@ -298,23 +374,23 @@ static void FilterRecordBatchWithNulls(benchmark::State& state) {
}

static void TakeInt64RandomIndicesNoNulls(benchmark::State& state) {
TakeBenchmark(state, false).Int64();
TakeBenchmark(state, /*indices_with_nulls=*/false).Int64();
}

static void TakeInt64RandomIndicesWithNulls(benchmark::State& state) {
TakeBenchmark(state, true).Int64();
TakeBenchmark(state, /*indices_with_nulls=*/true).Int64();
}

static void TakeInt64MonotonicIndices(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/false, /*monotonic=*/true).Int64();
}

static void TakeFixedSizeBinaryRandomIndicesNoNulls(benchmark::State& state) {
TakeBenchmark(state, false).FixedSizeBinary();
TakeBenchmark(state, /*indices_with_nulls=*/false).FixedSizeBinary();
}

static void TakeFixedSizeBinaryRandomIndicesWithNulls(benchmark::State& state) {
TakeBenchmark(state, true).FixedSizeBinary();
TakeBenchmark(state, /*indices_with_nulls=*/true).FixedSizeBinary();
}

static void TakeFixedSizeBinaryMonotonicIndices(benchmark::State& state) {
Expand All@@ -323,29 +399,91 @@ static void TakeFixedSizeBinaryMonotonicIndices(benchmark::State& state) {
}

static void TakeFSLInt64RandomIndicesNoNulls(benchmark::State& state) {
TakeBenchmark(state, false).FSLInt64();
TakeBenchmark(state, /*indices_with_nulls=*/false).FSLInt64();
}

static void TakeFSLInt64RandomIndicesWithNulls(benchmark::State& state) {
TakeBenchmark(state, true).FSLInt64();
TakeBenchmark(state, /*indices_with_nulls=*/true).FSLInt64();
}

static void TakeFSLInt64MonotonicIndices(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/false, /*monotonic=*/true).FSLInt64();
}

static void TakeStringRandomIndicesNoNulls(benchmark::State& state) {
TakeBenchmark(state, false).String();
TakeBenchmark(state, /*indices_with_nulls=*/false).String();
}

static void TakeStringRandomIndicesWithNulls(benchmark::State& state) {
TakeBenchmark(state, true).String();
TakeBenchmark(state, /*indices_with_nulls=*/true).String();
}

static void TakeStringMonotonicIndices(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/false, /*monotonic=*/true).FSLInt64();
}

static void TakeChunkedChunkedInt64RandomIndicesNoNulls(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/false)
.ChunkedInt64(/*num_chunks=*/100, /*chunk_indices_too=*/true);
}

static void TakeChunkedChunkedInt64RandomIndicesWithNulls(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/true)
.ChunkedInt64(/*num_chunks=*/100, /*chunk_indices_too=*/true);
}

static void TakeChunkedChunkedInt64MonotonicIndices(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/false, /*monotonic=*/true)
.ChunkedInt64(
/*num_chunks=*/100, /*chunk_indices_too=*/true);
}

static void TakeChunkedChunkedFSBRandomIndicesNoNulls(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/false)
.ChunkedFSB(/*num_chunks=*/100, /*chunk_indices_too=*/true);
}

static void TakeChunkedChunkedFSBRandomIndicesWithNulls(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/true)
.ChunkedFSB(/*num_chunks=*/100, /*chunk_indices_too=*/true);
}

static void TakeChunkedChunkedFSBMonotonicIndices(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/false, /*monotonic=*/true)
.ChunkedFSB(/*num_chunks=*/100, /*chunk_indices_too=*/true);
}

static void TakeChunkedChunkedStringRandomIndicesNoNulls(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/false)
.ChunkedString(/*num_chunks=*/100, /*chunk_indices_too=*/true);
}

static void TakeChunkedChunkedStringRandomIndicesWithNulls(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/true)
.ChunkedString(/*num_chunks=*/100, /*chunk_indices_too=*/true);
}

static void TakeChunkedChunkedStringMonotonicIndices(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/false, /*monotonic=*/true)
.ChunkedString(/*num_chunks=*/100, /*chunk_indices_too=*/true);
}

static void TakeChunkedFlatInt64RandomIndicesNoNulls(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/false)
.ChunkedInt64(/*num_chunks=*/100, /*chunk_indices_too=*/false);
}

static void TakeChunkedFlatInt64RandomIndicesWithNulls(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/true)
.ChunkedInt64(/*num_chunks=*/100, /*chunk_indices_too=*/false);
}

static void TakeChunkedFlatInt64MonotonicIndices(benchmark::State& state) {
TakeBenchmark(state, /*indices_with_nulls=*/false, /*monotonic=*/true)
.ChunkedInt64(
/*num_chunks=*/100, /*chunk_indices_too=*/false);
}

void FilterSetArgs(benchmark::internal::Benchmark* bench) {
for (int64_t size : g_data_sizes) {
for (int i = 0; i < static_cast<int>(g_filter_params.size()); ++i) {
Expand DownExpand Up@@ -405,6 +543,7 @@ void TakeFSBSetArgs(benchmark::internal::Benchmark* bench) {
}
}

// Flat values x Flat indices
BENCHMARK(TakeInt64RandomIndicesNoNulls)->Apply(TakeSetArgs);
BENCHMARK(TakeInt64RandomIndicesWithNulls)->Apply(TakeSetArgs);
BENCHMARK(TakeInt64MonotonicIndices)->Apply(TakeSetArgs);
Expand All@@ -418,5 +557,21 @@ BENCHMARK(TakeStringRandomIndicesNoNulls)->Apply(TakeSetArgs);
BENCHMARK(TakeStringRandomIndicesWithNulls)->Apply(TakeSetArgs);
BENCHMARK(TakeStringMonotonicIndices)->Apply(TakeSetArgs);

// Chunked values x Chunked indices
BENCHMARK(TakeChunkedChunkedInt64RandomIndicesNoNulls)->Apply(TakeSetArgs);
BENCHMARK(TakeChunkedChunkedInt64RandomIndicesWithNulls)->Apply(TakeSetArgs);
BENCHMARK(TakeChunkedChunkedInt64MonotonicIndices)->Apply(TakeSetArgs);
BENCHMARK(TakeChunkedChunkedFSBRandomIndicesNoNulls)->Apply(TakeFSBSetArgs);
BENCHMARK(TakeChunkedChunkedFSBRandomIndicesWithNulls)->Apply(TakeFSBSetArgs);
BENCHMARK(TakeChunkedChunkedFSBMonotonicIndices)->Apply(TakeFSBSetArgs);
BENCHMARK(TakeChunkedChunkedStringRandomIndicesNoNulls)->Apply(TakeSetArgs);
BENCHMARK(TakeChunkedChunkedStringRandomIndicesWithNulls)->Apply(TakeSetArgs);
BENCHMARK(TakeChunkedChunkedStringMonotonicIndices)->Apply(TakeSetArgs);

// Chunked values x Flat indices
BENCHMARK(TakeChunkedFlatInt64RandomIndicesNoNulls)->Apply(TakeSetArgs);
BENCHMARK(TakeChunkedFlatInt64RandomIndicesWithNulls)->Apply(TakeSetArgs);
BENCHMARK(TakeChunkedFlatInt64MonotonicIndices)->Apply(TakeSetArgs);

} // namespace compute
} // namespace arrow
Loading