unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to OSSpedia
duckdb/duckdb
senior9/10
algorithmic-optimizationdata-structuresquery-optimizationtesting-strategy

Add multiplicity-aware combine support and fix LIST state combination

Project: duckdb/duckdb | Language: C++ | Concepts: algorithmic-optimization, data-structures, query-optimization, testing-strategy | Difficulty: senior
Link: https://github.com/duckdb/duckdb/pull/24817

The problem

DuckDB's partial aggregate pushdown optimization rewrites aggregates over joins into partial aggregate states, which are then combined using a multiplicity count (i.e., how many times each partial state should be counted due to join duplicates). However, list-backed aggregates (LIST, MEDIAN, QUANTILE_CONT, MAD) ignored this multiplicity when combining states, causing them to undercount duplicates and produce incorrect results.

Why it's non-trivial

A naive fix would simply append each aggregate state once, which is what the existing code did. This ignores the actual number of times each state should contribute (its multiplicity), leading to incorrect results, especially for aggregates sensitive to value counts like median or list. Correctly handling multiplicity requires careful changes to the combine logic, including efficient handling of edge cases (zero, negative, or NULL multiplicity), and ensuring the fix integrates with DuckDB's aggregate framework and optimizer.

The approach

1. Forwarding multiplicity in combine calls

The core combine function for aggregate states is updated to forward the multiplicity information so that downstream combine logic can use it:

 void CombineAggrStateCombine(Vector &source, Vector &target, AggregateInputData
-	auto &bind_data = aggr_input_data.bind_data->Cast<ExportAggregateBindData>();
-	AggregateInputData combine_input(bind_data.aggr, bind_data.bind_data.get(), aggr_input_data.allocator,
-	                                 aggr_input_data.combine_type);
+	auto &bind_data = aggr_input_data.bind_data->Cast<ExportAggregateBindData>();
+	AggregateInputData combine_input(bind_data.aggr, bind_data.bind_data.get(), aggr_input_data.allocator,
+	                                 aggr_input_data.combine_type);
+	combine_input.combine_multiplicities = aggr_input_data.combine_multiplicities;
 	bind_data.aggr.GetStateCombineCallback()(source, target, combine_input, count);
 }

Annotation:
This ensures that when combining aggregate states, the multiplicity vector (if present) is forwarded to the actual combine logic, allowing it to handle each state's intended count.

2. Refactoring list state absorption

A helper is introduced to encapsulate the logic for merging one list state into another, handling empty/null cases:

inline void ListAbsorbState(ListAggState &source, ListAggState &target) {
	if (source.linked_list.total_capacity == 0) {
		// NULL, no need to append. This can happen with a filtered aggregate.
		return;
	}
	if (target.linked_list.total_capacity == 0) {
		target.linked_list = source.linked_list;
		return;
	}
	// Append the linked list.
	target.linked_list.last_segment->next = source.linked_list.first_segment;
	target.linked_list.last_segment = source.linked_list.last_segment;
	target.linked_list.total_capacity += source.linked_list.total_capacity;
}

Annotation:
This function is used to efficiently merge list states, handling cases where either side is empty or null.

3. Making ListCombineFunction multiplicity-aware

The main fix: the custom combine function for list-backed aggregates is updated to respect multiplicity, including efficient paths for multiplicity 0, 1, and general N:

-	if (aggr_input_data.combine_type == AggregateCombineType::ALLOW_DESTRUCTIVE) {
+	if (aggr_input_data.combine_type == AggregateCombineType::ALLOW_DESTRUCTIVE &&
+	    !aggr_input_data.combine_multiplicities) {
 		ListAbsorbFunction(states_vector, combined, aggr_input_data, count);
 		return;
 	}

Annotation:
Only use the destructive (efficient) path if no multiplicities are present.

+	UnifiedVectorFormat multiplicities;
+	const int64_t *multiplicity_data = nullptr;
+	if (aggr_input_data.combine_multiplicities) {
+		aggr_input_data.combine_multiplicities->ToUnifiedFormat(multiplicities);
+		multiplicity_data = UnifiedVectorFormat::GetData<int64_t>(multiplicities);
+	}

Annotation:
Extract multiplicity data if provided.

+		idx_t multiplicity = 1;
+		if (multiplicity_data) {
+			const auto multiplicity_idx = multiplicities.sel->get_index(i);
+			if (!multiplicities.validity.RowIsValid(multiplicity_idx)) {
+				continue;
+			}
+			const auto signed_multiplicity = multiplicity_data[multiplicity_idx];
+			if (signed_multiplicity < 0) {
+				throw InvalidInputException("combine_aggr multiplicity must be non-negative");
+			}
+			multiplicity = static_cast<idx_t>(signed_multiplicity);
+		}
+		if (multiplicity == 0 || source.linked_list.total_capacity == 0) {
+			continue;
+		}
+		if (multiplicity == 1 && aggr_input_data.combine_type == AggregateCombineType::ALLOW_DESTRUCTIVE) {
+			ListAbsorbState(source, target);
+			continue;
+		}

Annotation:

  • If multiplicity is NULL or zero, skip.
  • If multiplicity is negative, throw.
  • If multiplicity is one and destructive combine is allowed, use the efficient merge.
  • Otherwise, repeat the append as many times as multiplicity.
+		for (idx_t repeat_idx = 0; repeat_idx < multiplicity; repeat_idx++) {
+			functions.AppendListEntry(aggr_input_data.allocator, target.linked_list, input_data,
+			                          list_entry_t(0, entry_count));
+		}

Annotation:
This loop appends the source's list values to the target the correct number of times, ensuring join multiplicity is respected.

4. Comprehensive testing

New and updated tests ensure correctness for all affected aggregates and edge cases:

  • Test that join multiplicity is respected for LIST, MEDIAN, QUANTILE_CONT, and MAD.
  • Test that zero/NULL multiplicity is a no-op, and negative multiplicity raises an error.
  • Test both with and without the partial aggregate pushdown optimizer enabled.

Example (from new test file):

SELECT
    m.g,
    median(v.v),
    quantile_cont(v.v, 0.5),
    mad(v.v),
    list_sort(list(v.v))
FROM issue_24806_multiplicity m
INNER JOIN issue_24806_values v USING (k)
GROUP BY m.g;
----
1	1.0	1.0	0.0	[1, 1, 1, 1, 1, 1, 42, 42, 42, 42]

Annotation:
This checks that the aggregates see the correct number of duplicates after a join.

From the review thread

Mytherin: Thanks but I think we shouldn't expose another callback just for this - can you instead try to fix the callback for aggregates that don't implement it correctly?

Alignyx: Hi @Mytherin, thank you for your guidance. I’ve fixed the callback for the affected aggregates.

The problem occurred during partial aggregate pushdown. DuckDB first computed an intermediate aggregate state on one side of the join and then called combine_aggr(state, multiplicity) to reproduce the effect of the joined rows. The multiplicity tells the aggregate how many times each intermediate state should contribute to the final result.

The standard aggregate combine callback already supports this through AggregateInputData::combine_multiplicities. However, list-backed aggregates such as LIST, MEDIAN, QUANTILE, and MAD use a custom ListCombineFunction. That function ignored the multiplicity and appended each intermediate state only once. As a result, duplicate rows introduced by the join were lost, producing incorrect aggregate results.

I fixed ListCombineFunction so that it now reads and applies the multiplicity passed to the existing combine callback:

  • A NULL or zero multiplicity contributes nothing.
  • A multiplicity of one keeps the existing efficient destructive-merge path.
  • A multiplicity greater than one appends the state’s values the required number of times.
  • A negative multiplicity raises an error.

I also fixed the combine_aggr wrapper so that it forwards combine_multiplicities when it calls the underlying aggregate’s existing combine callback.

Takeaways

  • Aggregate combine logic must explicitly account for multiplicity when combining partial results, especially in the presence of joins or optimizations that alter row counts.
  • Edge cases (zero, null, negative multiplicity) should be handled explicitly and tested, as they can silently lead to incorrect results or crashes.
  • When fixing a bug in a shared callback or framework, prefer to update the existing implementation to handle all cases correctly, rather than adding special-case hooks.
  • Comprehensive, scenario-based tests (including optimizer toggles) are essential to verify correctness and prevent regressions in complex query optimizations.