Skip to content

perf: kll_sketch allocates its items on first use - #558

Draft
rustyconover wants to merge 3 commits into
apache:masterfrom
rustyconover:feat/kll-lazy-items
Draft

rustyconover wants to merge 3 commits into
apache:masterfrom
rustyconover:feat/kll-lazy-items

Conversation

@rustyconover

@rustyconover rustyconover commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor

kll_sketch's constructor allocates storage for k items straight away. For the default k = 200 with double items, that's 1.6 KB per sketch before anything has been added.

That's costly when an application creates many sketches that stay empty or small, for example one per group in a grouped aggregation. Query engines also create short-lived partial states that are only ever merged into, or that see a single row each before being combined.

This PR allocates items_ on the first update() or non-empty merge() instead:

  • internal_update(), which both update() and the level-zero part of merge() go through, allocates on demand.
  • merge() also allocates before merging the higher levels, because the other sketch's level zero can be empty while its higher levels aren't.
  • The copy and converting constructors allocate only if the source has items.
  • The destructor and const_iterator already handle items_ == nullptr.
  • get_sorted_view() returns an empty view before sorting or constructing item ranges when the sketch is empty. This single check protects both paths without allocating item storage; sort_level_zero() is private and is called only from this guarded entry point.

Results and the serialized format are unchanged.

Tests

Regression coverage uses test_allocator to check:

  • Empty sketches, their copies, and empty merges allocate only the levels; the first update allocates item storage, and destruction releases it.
  • Copy and move merges into sketches without item storage preserve counts and extrema through estimation mode.
  • Empty sorted views are safe and allocate nothing after construction, copying, type conversion, moving, assignment over an existing cached view, and both serialization round trips. First updates after those empty views still allocate and sort correctly.
  • Numeric sketches and sketches receiving only NaN updates also return empty sorted views safely.

The new custom-type empty-view test reproduces applying non-zero offset 1600 to null pointer under UBSan before the guard is added.

Validation:

  • All 17 CTest suites pass on Linux aarch64 with GCC 14.
  • KLL, integration, and deserialization-hardening suites pass with GCC AddressSanitizer, UndefinedBehaviorSanitizer, and libstdc++ assertions enabled.
  • The full KLL suite passes on macOS with Apple Clang 17, libc++, and UBSan.

Context

We found this in the DuckDB DataSketches extension. For GROUP BY queries with many groups, DuckDB creates a separate aggregate state for each input row and merges them later, and the eager allocation made a KLL aggregate over 2M groups use tens of GB. DuckDB made the same change to its own t-digest for the same reason. req_sketch (its first compactor) and tdigest (centroids_ and buffer_) allocate eagerly too; if this approach is welcome, I can send the same change for them.

🤖 Generated with Claude Code

@proost proost left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Idea is good, But null checking isn't enough.

For example, sort_level_zero methods should check null for items_.

And we should guard all null cases in one place not each call sites.

rustyconover and others added 2 commits October 29, 2026 12:13
The constructor allocated storage for k items up front, so every empty
sketch held k * sizeof(T) bytes (1.6 KB for the default k with doubles). Code
that creates many sketches that stay empty or small (one per group in a
grouped aggregation) paid for all of it.

Allocate the items on the first update or non-empty merge instead. The copy
and converting constructors only allocate if the source has items; the
destructor and the iterator already handle an unallocated buffer. Results and
the serialized format are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Check for an empty sketch in get_sorted_view before sorting level zero or forming item ranges. This shared entry point is the only caller of sort_level_zero, so no per-call-site null checks or eager allocation are needed.

Exercise empty views after construction, copies, moves, assignments and deserialization, then verify that first updates still allocate and sort correctly. The new custom-type regression reproduces the null-pointer offset under UBSan before the fix.
Review feedback: rather than each caller guarding against items_ being
nullptr, keep that handling inside the helpers that own item storage:

- allocate_items(): the only place items_ is allocated (update, merge)
- copy_items(): the only place items are copied (copy and converting
  constructors), replacing two duplicated null checks
- item_ptr(): the only place an item index becomes a pointer for readers;
  without item storage every range is empty, so sort_level_zero(),
  check_sorting() and get_sorted_view() need no guards of their own

get_sorted_view() is back to its original form, without the is_empty()
early return. UBSan flags the unguarded pointer arithmetic on an empty
sketch ("applying non-zero offset to null pointer"), and passes with
item_ptr().

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@rustyconover

Copy link
Copy Markdown
Contributor Author

@proost thanks for the review, agreed on both points. I've pushed 1fcf432, which moves all handling of an unallocated items_ into the helpers that own the item storage, so callers no longer guard anything:

  • allocate_items() is the only place items_ gets allocated (update and merge).
  • copy_items() is the only place items are copied. The copy and converting constructors both use it, which replaces their duplicated null checks.
  • item_ptr(index) is the only place an item index becomes a pointer for readers. Without item storage it returns nullptr, so every level range is empty. sort_level_zero(), check_sorting() and get_sorted_view() go through it and no longer need guards of their own.

That also lets me drop the is_empty() early return I had added to get_sorted_view(), so it's back to its original form. Serialization and the quantile and rank queries already check is_empty() before touching items. The iterator already treats a null items_ as begin() == end().

To confirm the guard is load-bearing: if I put the raw items_ + levels_[0] back in sort_level_zero(), the UBSan build fails the new empty-sketch tests with applying non-zero offset 800 to null pointer. With item_ptr(), kll_test passes in Release, UBSan (undefined,float-cast-overflow) and ASan builds.

The CI workflows on this PR are still waiting for maintainer approval, so they haven't run yet.

@leerho

leerho commented Oct 30, 2026

Copy link
Copy Markdown
Member

Thanks, Rusty. The DuckDB numbers make a strong case, and the reworked storage helpers from proost's review are a clean way to handle an unallocated items_.

Because the same choice affects several sketches (KLL, t-digest and REQ here, and the same sketches in Java), we'd like to agree on one convention across the project before merging any of them. I've posted a proposal on dev@ as a lazy consensus ending Nov 7, 01:00 UTC: https://lists.apache.org/thread/rv730qo61qqkd99dcy4wnglf3hz1s0fz

In short, it proposes allocating item storage on first use by default, with an explicit reserve() for users who want it up front, and keeping the basic exception guarantee on the first update()/merge() when allocation fails.

I'm marking this PR as a draft until that resolves. If it passes, we'd like to build on this PR, possibly with a reserve() added. Comments on the proposal are very welcome, on the list or here.

@leerho
leerho marked this pull request as draft October 30, 2026 23:26

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants