unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to OSSpedia
duckdb/duckdb
mid7/10
data-structurestype-systemstesting-strategy

Fix reservoir_quantile for DECIMAL types (preserve precision and support lists)

Project: duckdb/duckdb | Language: C++ | Concepts: data-structures, type-systems, testing-strategy | Difficulty: mid
Link: https://github.com/duckdb/duckdb/pull/24488

The problem

The reservoir_quantile aggregate function did not work correctly for DECIMAL columns. DuckDB represents decimals using underlying integer types (like INT16, INT32, INT64, INT128), but the aggregate function resolution was missing these physical types. As a result, calling reservoir_quantile on a DECIMAL column would fail, and there was no test coverage for this scenario.

Why it's non-trivial

A naive fix might try to just add support for DECIMAL types by treating them as numeric types, but this would lose precision or fail to handle the internal integer representation correctly. Additionally, handling lists of quantiles and serialization/deserialization of aggregate state with decimal types requires careful type tracking and implementation selection.

The approach

1. Track and Serialize Decimal Type Information

The bind data for the aggregate function was extended to remember the decimal logical type, so it can be used for correct serialization and deserialization:

struct ReservoirQuantileBindData : public FunctionData {
-	ReservoirQuantileBindData() {
+	ReservoirQuantileBindData() : decimal_type(LogicalType::INVALID) {
 	}
...
+	LogicalType decimal_type;
};

Annotation:
This ensures that the decimal type information is always available in the bind data, which is needed when (de)serializing the function state.

Serialization and deserialization logic was updated to include this new field:

static void Serialize(Serializer &serializer, const optional_ptr<FunctionData> bind_data_p,
                      const BoundAggregateFunction &function) {
...
+	serializer.WriteProperty(102, "decimal_type", bind_data.decimal_type);
}
...
static unique_ptr<FunctionData> Deserialize(Deserializer &deserializer, BoundAggregateFunction &function) {
	auto result = make_uniq<ReservoirQuantileBindData>();
	deserializer.ReadProperty(100, "quantiles", result->quantiles);
	deserializer.ReadProperty(101, "sample_size", result->sample_size);
+	result->decimal_type =
+	    deserializer.ReadPropertyWithExplicitDefault<LogicalType>(102, "decimal_type", LogicalType::INVALID);
	return std::move(result);
}

Annotation:
This ensures that when the aggregate function state is serialized (for things like distributed execution or plan caching), the decimal type is preserved and restored.

2. Add Physical Type Mappings for Decimals

The code now explicitly maps decimal logical types to their correct underlying integer physical types, and sets up the correct aggregate function implementation for each:

static void SetDecimalImplementation(BoundAggregateFunction &function, const LogicalType &decimal_type, bool is_list) {
...
+	if (is_list) {
+		function.ReplaceImplementation(GetReservoirQuantileListAggregateFunction(decimal_type));
+	} else {
+		switch (decimal_type.InternalType()) {
+		case PhysicalType::INT16:
+			function.ReplaceImplementation(
+			    AggregateFunction::UnaryAggregate<ReservoirQuantileState<int16_t>, int16_t, int16_t,
+			                                      ReservoirQuantileScalarOperation>(decimal_type, decimal_type));
+			break;
+		case PhysicalType::INT32:
+			function.ReplaceImplementation(
+			    AggregateFunction::UnaryAggregate<ReservoirQuantileState<int32_t>, int32_t, int32_t,
+			                                      ReservoirQuantileScalarOperation>(decimal_type, decimal_type));
+			break;
+		case PhysicalType::INT64:
+			function.ReplaceImplementation(
+			    AggregateFunction::UnaryAggregate<ReservoirQuantileState<int64_t>, int64_t, int64_t,
+			                                      ReservoirQuantileScalarOperation>(decimal_type, decimal_type));
+			break;
+		case PhysicalType::INT128:
+			function.ReplaceImplementation(
+			    AggregateFunction::UnaryAggregate<ReservoirQuantileState<hugeint_t>, hugeint_t, hugeint_t,
+			                                      ReservoirQuantileScalarOperation>(decimal_type, decimal_type));
+			break;
+		default:
+			throw InternalException("Invalid physical type for decimal reservoir quantile");
+		}
+	}
...
}

Annotation:
This code ensures that the correct template instantiation is used for the underlying integer type of the decimal, preserving precision and correctness.

3. Add Regression Tests for Decimal Quantiles

A new test file was added to verify that reservoir_quantile works on decimal columns, both for scalars and lists:

+statement ok
+CREATE TABLE decimals(d DECIMAL(10,2));
+
+statement ok
+INSERT INTO decimals VALUES (1.10), (2.20), (3.30), (4.40), (5.50), (6.60), (7.70), (8.80), (9.90), (10.00);
+
+query TT
+SELECT reservoir_quantile(d, 0.5), typeof(reservoir_quantile(d, 0.5)) FROM decimals;
+----
+5.50	DECIMAL(10,2)
+
+query TT
+SELECT reservoir_quantile(d, [0.1, 0.5, 0.9]), typeof(reservoir_quantile(d, [0.1, 0.5, 0.9])) FROM decimals;
+----
+[1.10, 5.50, 9.90]	DECIMAL(10,2)[]

Annotation:
This test ensures that the function now works as expected for decimal types, and that the output type is correct.

4. Update Plan Serialization Test Coverage

The new test is included in the plan serialization verification config:

@@ -50,6 +50,7 @@
         "test/sql/aggregate/aggregates/test_histogram_3529.test",
         "test/sql/aggregate/aggregates/test_mode.test",
         "test/sql/aggregate/aggregates/test_ordered_aggregates.test",
+        "test/sql/aggregate/aggregates/test_reservoir_quantile_decimal.test",
         ...

Annotation:
This ensures that serialization/deserialization of aggregate state with decimals is exercised in the test suite.

From the review thread

lnkuiper: Thanks for the PR! This breaks:

SET debug_verify_serializer=true;
SELECT reservoir_quantile(d, 0.5)
FROM (VALUES (1.55::DECIMAL(10,2))) t(d);

Could you have a look?

ptquocle: Thanks for reviewing the PR. I'll take a look tomorrow and push a fix then

Takeaways

  • When supporting new logical types (like DECIMAL) in aggregate functions, always map them to their correct physical representation to avoid precision loss or runtime errors.
  • Always serialize and deserialize all relevant type information for aggregate function state, especially for user-defined or parameterized types.
  • Add regression tests that cover both the new type support and serialization/deserialization paths to catch subtle bugs and ensure correctness.