Skip to content

Commit ddda17f

Browse files
yujun777yujun777
committed
[refactor](function) Return the state struct from TOPK_MERGE_MIN/MAX
The IVM apply expression writes the merge result straight back to the hidden top-K state columns, whose type comes from the aggregate. The scalar therefore has to return the state type itself; a 4-field result cannot be reassembled into the required (values, bound) struct because FE casts cannot turn a nullable field into a required one. Key changes: - topk_merge_min/max now return STRUCT<values ARRAY<STRUCT<value, count>>, bound>, identical to topk_min_count/max; n' and exhausted are derived by the caller (n, dn stay as inputs) - bound is max(S') only when pos > k; pos == k keeps bE, because nothing known was dropped - the merge stops scanning after the (k+1)-th positive candidate, which is enough to decide the bound - FE resultType delegates to the aggregate's result type; tests and the regression .out are updated, with a new pos == k case Unit Test: - FE: TopKFunctionSignatureTest - BE: run-be-ut.sh --filter=TopKCountTest.* - Regression: query_p0/aggregate/topk_min_max_count
1 parent 5e96382 commit ddda17f

5 files changed

Lines changed: 60 additions & 45 deletions

File tree

‎be/src/exprs/function/function_topk_merge.cpp‎

Lines changed: 22 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,11 @@ DataTypePtr find_value_type(const DataTypes& arguments) {
215215
* MAX): outside it at least one side may have dropped values. The post count of a value is
216216
* S + I - D; non-positive counts are removed (net-zero cancellations), then the first k values are
217217
* kept.
218+
*
219+
* The result is the state itself, i.e. the same type topk_min_count/topk_max_count return:
220+
* STRUCT<values ARRAY<STRUCT<value T, count BIGINT>>, bound T NULL>. n' = n + dn and
221+
* exhausted = (size(values) == 0 && n' > 0) are derived by the caller, which writes n' to the
222+
* COUNT state column and uses exhausted to trigger the fallback.
218223
*/
219224
template <bool IsMin>
220225
class FunctionTopKMerge : public IFunction {
@@ -237,14 +242,15 @@ class FunctionTopKMerge : public IFunction {
237242
throw Exception(ErrorCode::INVALID_ARGUMENT,
238243
"{} can not infer the value type from arguments", name);
239244
}
245+
// The result is the state type, identical to the aggregate's return type, so that the
246+
// apply expression can write it straight back to the hidden MIN_TOPK/MAX_TOPK and BOUND
247+
// columns.
240248
auto pair_type = std::make_shared<DataTypeStruct>(
241249
DataTypes {value_type, std::make_shared<DataTypeInt64>()},
242250
Strings {"value", "count"});
243251
auto values_type = std::make_shared<DataTypeArray>(make_nullable(pair_type));
244-
return std::make_shared<DataTypeStruct>(
245-
DataTypes {values_type, make_nullable(value_type),
246-
std::make_shared<DataTypeInt64>(), std::make_shared<DataTypeBool>()},
247-
Strings {"values", "bound", "n", "exhausted"});
252+
return std::make_shared<DataTypeStruct>(DataTypes {values_type, make_nullable(value_type)},
253+
Strings {"values", "bound"});
248254
}
249255

250256
Status execute_impl(FunctionContext* /*context*/, Block& block, const ColumnNumbers& arguments,
@@ -269,8 +275,6 @@ class FunctionTopKMerge : public IFunction {
269275
auto& out_value_column = out_pair.get_column(0);
270276
auto& out_count_column = assert_cast<ColumnInt64&>(out_pair.get_column(1));
271277
auto& out_bound_column = assert_cast<ColumnNullable&>(result_struct.get_column(1));
272-
auto& out_n_column = assert_cast<ColumnInt64&>(result_struct.get_column(2));
273-
auto& out_exhausted_column = assert_cast<ColumnUInt8&>(result_struct.get_column(3));
274278

275279
for (size_t row = 0; row != input_rows_count; ++row) {
276280
const int64_t k = read_int64(*columns[8], row, "k");
@@ -297,14 +301,16 @@ class FunctionTopKMerge : public IFunction {
297301
out_count_column, out_pair_nullable, &positive));
298302
out_values.get_offsets().push_back(out_pair.size());
299303

300-
// bound: NULL while the state is complete, the k-th kept value when the state is
301-
// truncated, or the effective bound when fewer than k values survived.
304+
// bound: NULL while the state is complete, the k-th kept value only when values were
305+
// actually dropped (pos > k), or the effective bound otherwise.
302306
// b IS NULL means the pre-image covers every distinct value; a NULL b_i/b_d means
303307
// that side of the delta was not truncated.
308+
// pos == k is NOT truncation: every known value fit, so keep bE instead of shrinking
309+
// the trusted range to max(S').
304310
const bool sides_complete = b_view.is_null && bi_view.is_null && bd_view.is_null;
305311
if (new_n == 0 || (sides_complete && std::cmp_less_equal(positive, k))) {
306312
out_bound_column.insert_default();
307-
} else if (std::cmp_greater_equal(positive, k)) {
313+
} else if (std::cmp_greater(positive, k)) {
308314
out_bound_column.get_nested_column().insert_from(out_value_column,
309315
out_pair.size() - 1);
310316
out_bound_column.get_null_map_data().push_back(0);
@@ -313,9 +319,6 @@ class FunctionTopKMerge : public IFunction {
313319
out_bound_column.get_nested_column().insert_from(*effective_bound, 0);
314320
out_bound_column.get_null_map_data().push_back(0);
315321
}
316-
317-
out_n_column.get_data().push_back(new_n);
318-
out_exhausted_column.get_data().push_back(positive == 0 && new_n > 0 ? 1 : 0);
319322
}
320323

321324
block.replace_by_position(result, std::move(result_column));
@@ -387,7 +390,8 @@ class FunctionTopKMerge : public IFunction {
387390

388391
/// Merges the trusted part (v <= bE for MIN) of S, I and D into the output list, keeping the
389392
/// first k values with a positive post count S + I - D. `positive` receives the number of
390-
/// positive candidates in the trusted range, which may exceed k.
393+
/// positive candidates, capped at k + 1: once more than k are found the scan stops, which is
394+
/// enough to decide the bound.
391395
Status append_trusted_values(const TopKArrayView& s_view, const TopKArrayView& i_view,
392396
const TopKArrayView& d_view, const IColumn* effective_bound,
393397
int64_t k, PrimitiveType type, IColumn& out_value_column,
@@ -434,6 +438,11 @@ class FunctionTopKMerge : public IFunction {
434438
out_pair_nullable.get_null_map_data().push_back(0);
435439
++kept;
436440
}
441+
if (std::cmp_greater(*positive, k)) {
442+
// pos > k is decided (a value was dropped); the rest only affects counts
443+
// outside the kept prefix, so stop scanning.
444+
break;
445+
}
437446
}
438447
}
439448
return Status::OK();

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/TopKMergeFunction.java‎

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,10 @@
2525
import org.apache.doris.nereids.trees.expressions.functions.agg.TopKCountFunction;
2626
import org.apache.doris.nereids.types.ArrayType;
2727
import org.apache.doris.nereids.types.BigIntType;
28-
import org.apache.doris.nereids.types.BooleanType;
2928
import org.apache.doris.nereids.types.DataType;
3029
import org.apache.doris.nereids.types.StructField;
3130
import org.apache.doris.nereids.types.StructType;
3231

33-
import com.google.common.collect.ImmutableList;
34-
3532
import java.util.List;
3633

3734
/**
@@ -40,16 +37,17 @@
4037
*
4138
* <pre>
4239
* topk_merge_min(S, b, I, b_i, D, b_d, n, dn, k)
43-
* -> STRUCT&lt;values ARRAY&lt;STRUCT&lt;value T, count BIGINT&gt;&gt;, bound T, n BIGINT,
44-
* exhausted BOOLEAN&gt;
40+
* -> STRUCT&lt;values ARRAY&lt;STRUCT&lt;value T, count BIGINT&gt;&gt;, bound T&gt;
4541
* </pre>
4642
*
43+
* <p>The result is the state itself and has exactly the type the aggregate returns, so apply can
44+
* write it back to the hidden top-K state columns without rebuilding a struct. The caller derives
45+
* {@code n' = n + dn} and {@code exhausted = (size(values) == 0 && n' > 0)} itself.
46+
*
4747
* <p>A NULL list is an empty list; a NULL {@code b}/{@code b_i}/{@code b_d} means that side is
4848
* complete/untruncated. The result struct itself is never NULL, only its {@code bound} field can
4949
* be: the merge only trusts values in {@code v <= min(b, b_i, b_d)}, so {@code bound} is the
5050
* effective upper bound (for MIN; the lower bound for MAX) of the resulting state.
51-
* {@code exhausted} is true when the value range is used up: the merge found no value to keep but
52-
* {@code n + dn > 0}, which requires a fallback to a complete refresh.
5351
*/
5452
public abstract class TopKMergeFunction extends ScalarFunction
5553
implements CustomSignature, AlwaysNotNullable {
@@ -134,15 +132,8 @@ private DataType inferValueType() {
134132
return null;
135133
}
136134

137-
/**
138-
* Builds {@code STRUCT<values ARRAY<STRUCT<value T, count BIGINT>>, bound T NULL, n BIGINT,
139-
* exhausted BOOLEAN>}.
140-
*/
135+
/** The merge result is the aggregate state type: {@code STRUCT<values, bound>}. */
141136
private static StructType resultType(DataType valueType) {
142-
return new StructType(ImmutableList.of(
143-
new StructField("values", ArrayType.of(TopKCountFunction.pairType(valueType)), false, ""),
144-
new StructField("bound", valueType, true, ""),
145-
new StructField("n", BigIntType.INSTANCE, false, ""),
146-
new StructField("exhausted", BooleanType.INSTANCE, false, "")));
137+
return TopKCountFunction.resultType(valueType);
147138
}
148139
}

‎fe/fe-core/src/test/java/org/apache/doris/nereids/trees/expressions/functions/agg/TopKFunctionSignatureTest.java‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
import org.apache.doris.nereids.trees.expressions.literal.NullLiteral;
2929
import org.apache.doris.nereids.types.ArrayType;
3030
import org.apache.doris.nereids.types.BigIntType;
31-
import org.apache.doris.nereids.types.BooleanType;
3231
import org.apache.doris.nereids.types.DataType;
3332
import org.apache.doris.nereids.types.IntegerType;
3433
import org.apache.doris.nereids.types.StructType;
@@ -116,15 +115,13 @@ void testMergeSignature() {
116115
Assertions.assertEquals(valueType, signature.getArgType(1));
117116
Assertions.assertEquals(BigIntType.INSTANCE, signature.getArgType(8));
118117

118+
// The merge result is the state itself, identical to the aggregate's return type.
119119
StructType resultType = (StructType) signature.returnType;
120-
Assertions.assertEquals(4, resultType.getFields().size());
120+
Assertions.assertEquals(TopKCountFunction.resultType(valueType), resultType);
121+
Assertions.assertEquals(2, resultType.getFields().size());
121122
assertPairType(valueType, (StructType) ((ArrayType) resultType.getField("values")
122123
.getDataType()).getItemType());
123124
Assertions.assertTrue(resultType.getField("bound").isNullable());
124-
Assertions.assertFalse(resultType.getField("n").isNullable());
125-
Assertions.assertEquals(BigIntType.INSTANCE, resultType.getField("n").getDataType());
126-
Assertions.assertFalse(resultType.getField("exhausted").isNullable());
127-
Assertions.assertEquals(BooleanType.INSTANCE, resultType.getField("exhausted").getDataType());
128125
Assertions.assertFalse(function.nullable());
129126
}
130127
}

‎regression-test/data/query_p0/aggregate/topk_min_max_count.out‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -109,32 +109,35 @@ e 0 0 {"values":[{"value":0, "count":5}], "bound":null} {"values":[{"value":0, "
109109
{"values":[{"value":1, "count":1}], "bound":null} {"values":[{"value":1, "count":1}], "bound":null}
110110

111111
-- !merge_min_complete --
112-
{"values":[{"value":1, "count":2}, {"value":2, "count":1}, {"value":4, "count":1}], "bound":null, "n":4, "exhausted":0}
112+
{"values":[{"value":1, "count":2}, {"value":2, "count":1}, {"value":4, "count":1}], "bound":null}
113113

114114
-- !merge_max_complete --
115-
{"values":[{"value":5, "count":1}, {"value":4, "count":1}, {"value":3, "count":1}, {"value":2, "count":1}], "bound":null, "n":4, "exhausted":0}
115+
{"values":[{"value":5, "count":1}, {"value":4, "count":1}, {"value":3, "count":1}, {"value":2, "count":1}], "bound":null}
116116

117117
-- !merge_min_pre_image_bound --
118-
{"values":[{"value":1, "count":1}, {"value":2, "count":1}], "bound":2, "n":3, "exhausted":0}
118+
{"values":[{"value":1, "count":1}, {"value":2, "count":1}], "bound":2}
119119

120120
-- !merge_min_insert_bound --
121-
{"values":[{"value":1, "count":1}, {"value":2, "count":1}], "bound":2, "n":3, "exhausted":0}
121+
{"values":[{"value":1, "count":1}, {"value":2, "count":1}], "bound":2}
122122

123123
-- !merge_min_delete_bound --
124-
{"values":[{"value":1, "count":1}], "bound":2, "n":2, "exhausted":0}
124+
{"values":[{"value":1, "count":1}], "bound":2}
125125

126126
-- !merge_min_exhausted --
127-
{"values":[], "bound":2, "n":1, "exhausted":1}
127+
{"values":[], "bound":2}
128128

129129
-- !merge_min_empty_after_delete --
130-
{"values":[], "bound":null, "n":0, "exhausted":0}
130+
{"values":[], "bound":null}
131131

132132
-- !merge_min_truncate_to_k --
133-
{"values":[{"value":1, "count":1}, {"value":2, "count":1}], "bound":2, "n":4, "exhausted":0}
133+
{"values":[{"value":1, "count":1}, {"value":2, "count":1}], "bound":2}
134+
135+
-- !merge_min_pos_equals_k --
136+
{"values":[{"value":1, "count":1}, {"value":2, "count":1}], "bound":5}
134137

135138
-- !merge_max_truncate_to_k --
136-
{"values":[{"value":4, "count":1}, {"value":3, "count":1}], "bound":3, "n":4, "exhausted":0}
139+
{"values":[{"value":4, "count":1}, {"value":3, "count":1}], "bound":3}
137140

138141
-- !merge_all_null --
139-
{"values":[], "bound":null, "n":0, "exhausted":0}
142+
{"values":[], "bound":null}
140143

‎regression-test/suites/query_p0/aggregate/topk_min_max_count.groovy‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,6 +308,21 @@ suite("topk_min_max_count") {
308308
CAST(4 AS BIGINT), CAST(0 AS BIGINT), CAST(2 AS BIGINT))
309309
"""
310310

311+
// pos == k is not truncation: every known value fits, so the state keeps bE (5) instead of
312+
// shrinking the trusted range to max(S') (2).
313+
order_qt_merge_min_pos_equals_k """
314+
SELECT topk_merge_min(
315+
CAST(array(named_struct('value', 1, 'count', 1))
316+
AS ARRAY<STRUCT<value:INT, count:BIGINT>>),
317+
CAST(NULL AS INT),
318+
CAST(array(named_struct('value', 2, 'count', 1))
319+
AS ARRAY<STRUCT<value:INT, count:BIGINT>>),
320+
CAST(5 AS INT),
321+
CAST(array() AS ARRAY<STRUCT<value:INT, count:BIGINT>>),
322+
CAST(NULL AS INT),
323+
CAST(1 AS BIGINT), CAST(1 AS BIGINT), CAST(2 AS BIGINT))
324+
"""
325+
311326
// MAX is symmetric.
312327
order_qt_merge_max_truncate_to_k """
313328
SELECT topk_merge_max(

0 commit comments

Comments
 (0)