Extend FUNNEL_COUNT to support multiple CORRELATE_BY columns by tarun11Mavani · Pull Request #18760 · apache/pinot

tarun11Mavani · 2026-06-15T06:06:10Z

Summary

Extends FUNNEL_COUNT to accept multiple columns in CORRELATE_BY(col1, col2, ...),
enabling funnel analysis that tracks users through steps within a composite key
(e.g., per user per device category), not just a single dimension.

Design

Doc with example: https://docs.google.com/document/d/1gWQ7XBbJdQcUdZvBevFnGTVbCVJ3fN49biIsSOtRdhM/edit?tab=t.0

The single-key aggregation path is preserved as a zero-overhead fast path — structurally
identical to the original single-column implementation — so existing queries see no
regression. Multi-key support is added as a separate code path selected once per block.

AggregationStrategy: Split into two abstract methods (addSingleKey / addMultiKey)
with separate aggregation loops for single-key and multi-key, eliminating per-row branching
on the dominant single-key path.
DictIdsWrapper: Added composite-key mapping for multi-column CORRELATE_BY. Uses
stride-based arithmetic when the product of dictionary sizes fits in int, falling back
to a HashMap<IntArrayList, Integer> for large key spaces. Also adds toCompositeString
for length-prefix encoded composite string keys used during result extraction.
SortedAggregationResult: Updated to handle multi-key by tracking secondary keys via
a HashMap within each primary-key group (data is sorted on the primary column only).
BitmapAggregationStrategy, SortedAggregationStrategy,
ThetaSketchAggregationStrategy: Implement both addSingleKey and addMultiKey.
SetResultExtractionStrategy, BitmapResultExtractionStrategy: Updated to
reverse-map composite IDs back to per-column dictionary values during result extraction.
FunnelCountSortedAggregationFunction: Propagates multi-dictionary context through
the sorted aggregation result extraction pipeline.

Example Query

SELECT FUNNEL_COUNT(
  STEPS(step1_col, step2_col, step3_col),
  CORRELATE_BY(user_id, device_category),
  SETTINGS('theta_sketch')
) FROM myTable

HashMap Fallback Usage

The composite-key mapping in DictIdsWrapper uses stride-based arithmetic when the product of per-segment dictionary sizes fits in int, falling back to a HashMap<IntArrayList, Integer> only when that product exceeds Integer.MAX_VALUE. In practice the fallback is rarely exercised:

The limit is per-segment, not global. The check multiplies segment-local dictionary cardinalities, which are a fraction of a column's global cardinality. A column with hundreds of millions of distinct values globally typically has far fewer per segment.
It's a product against a low-cardinality second key. The intended use of multi-key CORRELATE_BY is one high-cardinality key (e.g., user_id) correlated with a low-cardinality dimension (e.g., device_category, country, platform). With a second key in the typical low-cardinality range, the primary key can hold a large number of distinct values per segment before overflow:
- second key ~100 → primary key up to ~21M / segment
- second key ~1,000 → primary key up to ~2.1M / segment
- second key ~10,000 → primary key up to ~215K / segment

These ceilings sit well above realistic per-segment cardinalities for the high-cardinality key, so the stride path covers the common case. The HashMap fallback only kicks in for unusual queries that correlate two genuinely high-cardinality columns in the same segment, and exists purely to keep those queries correct rather than fast.

Test Plan

Existing single-key funnel integration tests pass unchanged
New multi-key integration tests: testMultiKeyOverall, testMultiKeyGroupBy, testMultiKeyWithFilter, testMultiKeyEmptyResult
All strategies tested: BITMAP, SORTED, THETA_SKETCH, SET
New unit tests: DictIdsWrapperTest (8 tests), SortedAggregationResultTest (2 tests)
JMH benchmarks verify zero regression on single-key path
Multi-key path benchmarked for throughput baseline

tarun11Mavani · 2026-06-15T06:07:59Z

Performance Validation (JMH)

Ran BenchmarkFunnelCountAggregation locally (JDK 21, 1 fork, 3 warmup / 5 measurement iterations of 5s each) on a 10,000-row block with 4 funnel steps.

Single-key path — Before (baseline) vs After (this PR):

Strategy	Baseline (ops/s)	After (ops/s)	Delta
bitmap	275.4 ± 9.7	274.6 ± 10.2	-0.3%
set	270.6 ± 23.9	275.4 ± 14.4	+1.8%
theta_sketch	1685.2 ± 47.2	1579.4 ± 503.8	-6.3%*
partitioned	276.6 ± 10.5	273.5 ± 14.8	-1.1%
partitioned_sorted	3688.5 ± 45.4	3352.2 ± 1022	-9.1%*

*theta_sketch and partitioned_sorted show large error bars indicating JVM warmup variance, not a real regression. Scores overlap within error margins.

Multi-key path (new feature, this PR only):

Strategy	Throughput (ops/s)
bitmap	369.0 ± 20.0
set	362.0 ± 18.1
theta_sketch	287.2 ± 7.1
partitioned	371.5 ± 21.1
partitioned_sorted	912.1 ± 15.6

Single-key path shows NO statistically significant regression. All deltas are within error margins. The bitmap/set/partitioned strategies (which dominate real workloads) are within ±2% of baseline — effectively identical.

tarun11Mavani · 2026-06-15T08:14:48Z

cc @darioliberman @rohityadav1993

codecov-commenter · 2026-06-15T08:26:29Z

Codecov Report

❌ Patch coverage is 46.36015% with 140 lines in your changes missing coverage. Please review.
✅ Project coverage is 64.74%. Comparing base (302783f) to head (e72173d).
⚠️ Report is 10 commits behind head on master.

Files with missing lines	Patch %	Lines
...gregation/function/funnel/AggregationStrategy.java	18.29%	64 Missing and 3 partials ⚠️
...unction/funnel/BitmapResultExtractionStrategy.java	11.53%	22 Missing and 1 partial ⚠️
...n/function/funnel/SetResultExtractionStrategy.java	20.00%	11 Missing and 1 partial ⚠️
...ation/function/funnel/SortedAggregationResult.java	78.00%	9 Missing and 2 partials ⚠️
...ry/aggregation/function/funnel/DictIdsWrapper.java	85.50%	9 Missing and 1 partial ⚠️
...unction/funnel/ThetaSketchAggregationStrategy.java	0.00%	8 Missing ⚠️
...ion/function/funnel/BitmapAggregationStrategy.java	0.00%	3 Missing ⚠️
...n/funnel/FunnelCountSortedAggregationFunction.java	40.00%	1 Missing and 2 partials ⚠️
...ion/function/funnel/SortedAggregationStrategy.java	0.00%	3 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #18760      +/-   ##
============================================
- Coverage     64.77%   64.74%   -0.03%     
  Complexity     1322     1322              
============================================
  Files          3393     3393              
  Lines        211022   211316     +294     
  Branches      33135    33216      +81     
============================================
+ Hits         136687   136816     +129     
- Misses        63322    63459     +137     
- Partials      11013    11041      +28

Flag	Coverage Δ
custom-integration1	`100.00% <ø> (ø)`
integration	`100.00% <ø> (ø)`
integration1	`100.00% <ø> (ø)`
integration2	`0.00% <ø> (ø)`
java-21	`64.74% <46.36%> (-0.03%)`	⬇️
temurin	`64.74% <46.36%> (-0.03%)`	⬇️
unittests	`64.74% <46.36%> (-0.03%)`	⬇️
unittests1	`56.94% <46.36%> (-0.05%)`	⬇️
unittests2	`37.13% <0.00%> (-0.05%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

rohityadav1993 · 2026-06-18T18:24:40Z

+    int[] dictIds = new int[numKeys];
+    while (iterator.hasNext()) {
+      wrapper.reverseCompositeId(iterator.next(), dictIds);
+      valueBitmap.add(DictIdsWrapper.toCompositeString(wrapper._dictionaries, dictIds).hashCode());


DictIdsWrapper.toCompositeString(wrapper._dictionaries, dictIds).hashCode())
This now generates a theoretical hash collision so bitmap strategy for multikey correlation is not accurate anymore.
We will have to either call this out as limitation in docs or find an alternative.

This is actually an existing limitation of the bitmap strategy, not something new from multi-key. The single-key path in convertToValueBitmap already uses .hashCode() for LONG, FLOAT, DOUBLE, and STRING types — only INT gets exact values stored directly. The multi-key path is consistent: toCompositeString itself is collision-free (length-prefix encoding is injective), but the .hashCode() mapping to 32-bit int has the same collision properties as single-key STRING at line 109.

I've updated the method Javadoc on convertCompositeToValueBitmap to call this out more explicitly, linking it to the existing single-key non-INT approximation.

Let's document this for multi-key int values it can be approximate.

yes. already covered in the class/method Javadoc.

I don't understand the point about all the meticulous process of keeping maps, reversing, building composite strings and what have you, we are anyways getting a hash code here.

Make sense. I'll replace the toCompositeString().hashCode() with a direct hash-combining loop over the dictionary values.

rohityadav1993 · 2026-06-18T18:44:12Z

+
+  @Override
+  void addMultiKey(UpdateSketch[] stepsSketches, int step, Dictionary[] dictionaries, int[] correlationDictIds) {
+    stepsSketches[step].update(DictIdsWrapper.toCompositeString(dictionaries, correlationDictIds));


I think there will be a lot of new string creation cost and subsequent GC pressure with toCompositeString for each row.

Similarly at other places if cardinality of distinct correlation multi-keys is high in a query.

Fair point — toCompositeString does allocate a new StringBuilder + String per row. A couple of things to note though:

Theta sketch's update() only accepts primitives (int, long, double) or String/byte[]. Since a multi-key tuple has no single primitive representation, some form of serialization is unavoidable here.

The single-key STRING path (line 75) already allocates a string per row via dictionary.getStringValue(), so the cost pattern is structurally similar — just slightly more overhead from the length-prefix encoding.

JMH baseline for multi-key theta_sketch is 287 ops/s, which we can measure future optimizations against.

One option would be switching to update(byte[]) with a reusable ByteBuffer to avoid the String intermediate, but wanted to keep it simple for the initial implementation. Do you have other optimization ideas in mind?

rohityadav1993 · 2026-06-23T15:59:27Z

+
+  @Override
+  void addMultiKey(DictIdsWrapper dictIdsWrapper, int step, Dictionary[] dictionaries, int[] correlationDictIds) {
+    dictIdsWrapper._stepsBitmaps[step].add(dictIdsWrapper.getCompositeCorrelationId(correlationDictIds));


Is it possible to cache/memoize getCompositeCorrelationId, this gets computed for every step that matches for a multi-key

On the stride path (common case), getCompositeCorrelationId is just multiply-add arithmetic — caching costs more than it saves (array comparison overhead > a few multiplies).

On the HashMap path (rare overflow case), repeated map lookups per step do cost more, but this path is already exceptional. And the number of repeated calls per row is bounded by the number of matching steps (typically 2-5).

Hence, I feel we are better off without any cache here.

rohityadav1993 · 2026-06-23T16:00:07Z

Please address the comments, overall LGTM.

Enable funnel analysis that tracks users through steps within a composite key (e.g., per user per device category) by accepting multiple columns in CORRELATE_BY(col1, col2, ...). The single-key path is preserved as a zero-overhead fast path with separate addSingleKey/addMultiKey abstract methods and dedicated aggregation loops, ensuring no regression for existing single-column queries. Multi-key composite ID mapping uses stride-based arithmetic when the product of dictionary sizes fits in int, with a HashMap fallback for large key spaces. Co-authored-by: Cursor <cursoragent@cursor.com>

Benchmark was used for local validation only; not needed in the PR. Co-authored-by: Cursor <cursoragent@cursor.com>

Keep the original `add(Dictionary, A, int, int)` abstract method unchanged. The new multi-key method is added as `addMultiKey(A, int, Dictionary[], int[])`. Co-authored-by: Cursor <cursoragent@cursor.com>

…egationResult double-count - Add DictIdsWrapperTest covering the HashMap fallback path (large-cardinality composite keys where product of dict sizes exceeds Integer.MAX_VALUE): path selection, sequential ID assignment, same-key idempotency, key-order sensitivity, and round-trip for 2- and 3-column keys. Also covers stride-path reverseCompositeId round-trip. Add isHashMapPath() predicate to DictIdsWrapper for test introspection (avoids widening _strides visibility). - Add SortedAggregationResultTest with multi-key extraction scenarios. - Fix SortedAggregationResult.extractResult(): clear _secondaryKeySteps after flushMultiKeyGroup() so a second call (defensive) returns zeros rather than double-counting the last open primary group.

Add method-level doc on convertCompositeToValueBitmap linking the multi-key .hashCode() usage to the existing single-key non-INT approximation in convertToValueBitmap.

tarun11Mavani · 2026-06-24T05:47:49Z

cc @shauryachats @Jackie-Jiang @xiangfu0 to review.

dario-liberman · 2026-06-24T08:18:08Z

+  /**
+   * Creates an aggregation result for multi-key correlation.
+   */
+  abstract A createAggregationResultMultiKey(Dictionary[] dictionaries);


I think we can start with this for now, but it would be better in my opinion if we do not have a physical column requirement on secondary correlation keys.
I think for the primary correlation key it makes sense that we require it to be a simple dictionary encoded column, but for secondary correlation I believe we could allow for any expression including projections.
Again, I think it is fine to start with this restriction for this pull request, but would be great to follow-up with a generalization.

make sense. We'll keep the physical column restriction for this PR and follow up with expression support for secondary keys.

dario-liberman · 2026-06-24T08:25:05Z

+   * Maps a tuple of per-column dictionary IDs to a single composite int suitable for RoaringBitmap.
+   * Only used for multi-key; for single-key, callers should add the dictId directly.
+   */
+  int getCompositeCorrelationId(int[] dictIds) {


In my opinion for composite case we should use 64 bits (i.e. long). The first 32 bits for the primary correlation key, the other 32 bits for the secondary correlation keys.

To make sure I understand: with ((long) primaryDictId << 32) | hash32(secondaryDictIds), the result is a 64-bit value, but RoaringBitmap only stores 32-bit ints. Are you suggesting we:
(a) use Roaring64Bitmap instead,
(b) hash/cast the long down to 32 bits for the bitmap,
or (c) something else?
Also — this makes the set strategy and partitioned-bitmap strategy approximate for multi-key (they're exact today since composite IDs are collision-free). We need to agree on this and document it.

not sure if we need 64 bits to be honest, but the main point is that the only scenarios that are collision free today are the partitioned strategies and the set for ints, all others convert the value (e.g. a uuid string) to a 32 bits hash.
Even for the strategies where we use dict IDs today, we still need to explain the collision semantics in the doc for multi key case. I was thinking that the contract could be that you would never have smaller counts due to collisions when adding a secondary correlation key compared to having only a primary correlation key. Say counting by user should never be larger than counting by user+device.

dario-liberman · 2026-06-24T08:47:15Z

+    if (existingId != null) {
+      return existingId;
+    }
+    IntArrayList insertKey = new IntArrayList(dictIds);


All of this memory allocation and subsequent garbage collection seems to me completely unnecessary.
We do not have a collision free counting strategy anyways.

agreed.
If we adopt the approach from above comment, this entire HashMap path goes away. If we keep it, I can at least eliminate the double-store (IntArrayList + dictIds.clone()). Will address based on how we resolve comment 2.

dario-liberman · 2026-06-24T09:04:24Z

+      _secondaryKeySteps.clear();
+    }
+
+    _lookupKey.clear();


Consider simply _lookupKey = IntArrayList.wrap(correlationIds)

correlationIds includes the primary key at index 0 (we only need secondary keys from index 1+), and wrap() aliases the underlying array which gets mutated on the next row

All I am trying to find is a way to avoid allocations. There must be ways to reuse fast util maps or perhaps have array lookups for scenarios up to say 5 entries. Teams mostly use secondary correlation for things like order id, mobile session if, device id, etc, but it is very rare to have multiple simultaneous trips per user, or multiple sessions or devices per user, and even when it happens we usually talk about a small number. I think we need as much as possible to optimise for this in the sorted case.

dario-liberman · 2026-06-24T09:06:46Z

+    boolean[] steps = _secondaryKeySteps.get(_lookupKey);
+    if (steps == null) {
+      steps = new boolean[_numSteps];
+      _secondaryKeySteps.put(new IntArrayList(_lookupKey), steps);


Why do we need to create a clone of lookupKey here? Are we mutating it somewhere?

The clone is needed because _lookupKey is a reusable buffer that gets clear()+add() on every row. Without the clone, the HashMap key mutates in-place on the next row.

I see you try to remove the sorting key from the array, but it is innocuous, it's the same for every entry here. The cost of all this moving around and specially the cost of memory allocation and garbage collection is much higher than adding one more int to a hash code calculation or equality calculation which is all that you aim to achieve here.

dario-liberman · 2026-06-25T01:25:23Z

my main concern is to avoid memory allocations for partitioned sorted strategy.

…map extraction SortedAggregationResult: replace HashMap<IntArrayList, boolean[]> with pre-allocated flat arrays and linear scan. Zero allocations in the hot loop for typical workloads (1-5 secondary key combos per primary group). BitmapResultExtractionStrategy: replace toCompositeString().hashCode() with direct type-aware hash combining, avoiding StringBuilder/String allocation per composite ID during extraction.

tarun11Mavani force-pushed the funnel-multi-correlate-keys branch from e1d2196 to d6bb092 Compare June 15, 2026 06:41

tarun11Mavani marked this pull request as ready for review June 15, 2026 07:10

rohityadav1993 reviewed Jun 18, 2026

View reviewed changes

rohityadav1993 reviewed Jun 23, 2026

View reviewed changes

tarun11Mavani and others added 5 commits June 24, 2026 04:41

Remove benchmark file from PR

bccbbd2

Benchmark was used for local validation only; not needed in the PR. Co-authored-by: Cursor <cursoragent@cursor.com>

Preserve original add() signature for backward compatibility

a3fcf9b

Keep the original `add(Dictionary, A, int, int)` abstract method unchanged. The new multi-key method is added as `addMultiKey(A, int, Dictionary[], int[])`. Co-authored-by: Cursor <cursoragent@cursor.com>

Clarify hash approximation in BitmapResultExtractionStrategy Javadoc

b21be3a

Add method-level doc on convertCompositeToValueBitmap linking the multi-key .hashCode() usage to the existing single-key non-INT approximation in convertToValueBitmap.

tarun11Mavani force-pushed the funnel-multi-correlate-keys branch from e30af27 to b21be3a Compare June 24, 2026 04:41

rohityadav1993 approved these changes Jun 24, 2026

View reviewed changes

dario-liberman reviewed Jun 24, 2026

View reviewed changes

dario-liberman approved these changes Jun 24, 2026

View reviewed changes

dario-liberman approved these changes Jun 25, 2026

View reviewed changes

tarun11Mavani force-pushed the funnel-multi-correlate-keys branch from f1adfa8 to e72173d Compare June 25, 2026 13:25

Uh oh!

Conversation

tarun11Mavani commented Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Design

Example Query

HashMap Fallback Usage

Test Plan

Uh oh!

tarun11Mavani commented Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Performance Validation (JMH)

Uh oh!

tarun11Mavani commented Jun 15, 2026

Uh oh!

codecov-commenter commented Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rohityadav1993 commented Jun 23, 2026

Uh oh!

tarun11Mavani commented Jun 24, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dario-liberman commented Jun 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

tarun11Mavani commented Jun 15, 2026 •

edited

Loading

tarun11Mavani commented Jun 15, 2026 •

edited

Loading

codecov-commenter commented Jun 15, 2026 •

edited

Loading