This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 3a9249fa3a [core] Support count-limit for collect agg (#10540)
3a9249fa3a is described below
commit 3a9249fa3a0575576541500ca1b7a8aa339ee939
Author: tsreaper <[email protected]>
AuthorDate: Wed Oct 14 13:29:44 2026 +0800
[core] Support count-limit for collect agg (#10540)
---
.../primary-key-table/merge-engine/aggregation.mdx | 8 ++
.../main/java/org/apache/paimon/CoreOptions.java | 7 ++
.../compact/aggregate/FieldCollectAgg.java | 99 +++++++++++++---------
.../compact/aggregate/FieldIgnoreRetractAgg.java | 5 ++
.../aggregate/factory/FieldCollectAggFactory.java | 5 +-
.../compact/aggregate/FieldAggregatorTest.java | 45 ++++++++++
6 files changed, 126 insertions(+), 43 deletions(-)
diff --git a/docs/docs/primary-key-table/merge-engine/aggregation.mdx
b/docs/docs/primary-key-table/merge-engine/aggregation.mdx
index 3e81d1d52d..b0dbe03e00 100644
--- a/docs/docs/primary-key-table/merge-engine/aggregation.mdx
+++ b/docs/docs/primary-key-table/merge-engine/aggregation.mdx
@@ -181,6 +181,14 @@ count of distinct event IDs.
The collect function collects elements into an Array. You can set
`fields.<field-name>.distinct=true` to deduplicate elements.
It only supports ARRAY type.
+ Use `fields.<field-name>.count-limit=<Integer>` to limit the number of
collected elements. There is no limit by default.
+ A positive value `N` keeps the first `N` elements in aggregation order; a
negative value `-N` keeps the last `N` elements.
+
+ With `distinct=true`, the limit counts distinct elements: positive limits
retain the first occurrence of each element,
+ while negative limits retain the last occurrence. The retained elements keep
their original order.
+ For example, collecting `[1, 2, 1, 3]` with `distinct=true` produces `[1,
2]` for `count-limit=2`
+ and `[1, 3]` for `count-limit=-2`.
+
### merge_map
The merge_map function merge input maps. It only supports MAP type.
diff --git a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
index d44370d760..f9cb0914f6 100644
--- a/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
+++ b/paimon-api/src/main/java/org/apache/paimon/CoreOptions.java
@@ -3500,6 +3500,13 @@ public class CoreOptions implements Serializable {
.defaultValue(false));
}
+ public int fieldCollectAggCountLimit(String fieldName) {
+ return options.get(
+ key(FIELDS_PREFIX + "." + fieldName + "." + COUNT_LIMIT)
+ .intType()
+ .defaultValue(Integer.MAX_VALUE));
+ }
+
public String fieldListAggDelimiter(String fieldName) {
return options.get(
key(FIELDS_PREFIX + "." + fieldName + "." + LIST_AGG_DELIMITER)
diff --git
a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/FieldCollectAgg.java
b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/FieldCollectAgg.java
index 60f60e82ce..4559f07eb9 100644
---
a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/FieldCollectAgg.java
+++
b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/FieldCollectAgg.java
@@ -33,8 +33,8 @@ import javax.annotation.Nullable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
-import java.util.HashSet;
import java.util.Iterator;
+import java.util.LinkedHashSet;
import java.util.List;
import java.util.function.BiFunction;
@@ -46,12 +46,14 @@ public class FieldCollectAgg extends FieldAggregator {
private static final long serialVersionUID = 1L;
private final boolean distinct;
+ private final int countLimit;
private final InternalArray.ElementGetter elementGetter;
@Nullable private final BiFunction<Object, Object, Boolean> equaliser;
- public FieldCollectAgg(String name, ArrayType dataType, boolean distinct) {
+ public FieldCollectAgg(String name, ArrayType dataType, boolean distinct,
int countLimit) {
super(name, dataType);
this.distinct = distinct;
+ this.countLimit = countLimit;
this.elementGetter =
InternalArray.createElementGetter(dataType.getElementType());
// The equaliser is built from the element type alone: retract() needs
it whether or not
@@ -95,69 +97,82 @@ public class FieldCollectAgg extends FieldAggregator {
@Override
public Object aggReversed(Object accumulator, Object inputField) {
- // we don't need to actually do the reverse here for this agg
- // because accumulator has been distinct, just let accumulator be
accumulator will speed up
- // distinct process
- return agg(accumulator, inputField);
+ if (countLimit != Integer.MAX_VALUE) {
+ // order is important with limit
+ return agg(inputField, accumulator, true);
+ } else {
+ // we don't need to actually do the reverse here for this agg
+ // because accumulator has been distinct, just let accumulator be
accumulator will speed
+ // up
+ // distinct process
+ return agg(accumulator, inputField, false);
+ }
}
@Override
public Object agg(Object accumulator, Object inputField) {
- if (accumulator == null && inputField == null) {
- return null;
- }
+ return agg(accumulator, inputField, false);
+ }
- if ((accumulator == null || inputField == null) && !distinct) {
- return accumulator == null ? inputField : accumulator;
- }
+ private Object agg(Object accumulator, Object inputField, boolean
accumulatorHasDuplicate) {
+ InternalArray accumulatorArray = (InternalArray) accumulator;
+ InternalArray inputArray = (InternalArray) inputField;
- if (distinct && equaliser != null) {
- List<Object> collection = new ArrayList<>();
- // do not need to distinct accumulator, because the accumulator is
always distinct, no
- // need to distinct it every time
- collect(collection, accumulator);
- collectWithEqualiser(collection, inputField);
+ if (accumulatorArray == null && inputArray == null) {
+ return null;
+ } else if (inputArray == null
+ && !distinct
+ && accumulatorArray.size() <= Math.abs(countLimit)) {
+ return accumulatorArray;
+ } else if (accumulatorArray == null
+ && !distinct
+ && inputArray.size() <= Math.abs(countLimit)) {
+ return inputArray;
+ }
+
+ boolean removeDuplicate = distinct && equaliser != null;
+ Collection<Object> collection =
+ distinct && equaliser == null ? new LinkedHashSet<>() : new
ArrayList<>();
+ if (countLimit >= 0) {
+ collect(collection, accumulatorArray, removeDuplicate &&
accumulatorHasDuplicate);
+ collect(collection, inputArray, removeDuplicate);
return new GenericArray(collection.toArray());
} else {
- Collection<Object> collection = distinct ? new HashSet<>() : new
ArrayList<>();
- collect(collection, accumulator);
- collect(collection, inputField);
- return new GenericArray(collection.toArray());
+ collect(collection, inputArray, removeDuplicate);
+ collect(collection, accumulatorArray, removeDuplicate);
+ List<Object> result =
+ collection instanceof List
+ ? (List<Object>) collection
+ : new ArrayList<>(collection);
+ Collections.reverse(result);
+ return new GenericArray(result.toArray());
}
}
- private void collect(Collection<Object> collection, @Nullable Object data)
{
- if (data == null) {
+ private void collect(
+ Collection<Object> collection, @Nullable InternalArray array,
boolean removeDuplicate) {
+ if (array == null) {
return;
}
- InternalArray array = (InternalArray) data;
- for (int i = 0; i < array.size(); i++) {
- collection.add(elementGetter.getElementOrNull(array, i));
- }
- }
-
- private void collectWithEqualiser(List<Object> list, Object data) {
- if (data == null) {
- return;
- }
-
- InternalArray array = (InternalArray) data;
- for (int i = 0; i < array.size(); i++) {
- Object element = elementGetter.getElementOrNull(array, i);
- if (!contains(list, element)) {
- list.add(element);
+ long limit = Math.abs((long) countLimit);
+ for (int i = 0; i < array.size() && collection.size() < limit; i++) {
+ Object element =
+ elementGetter.getElementOrNull(
+ array, countLimit < 0 ? array.size() - 1 - i : i);
+ if (!(removeDuplicate && contains(collection, element))) {
+ collection.add(element);
}
}
}
- private boolean contains(List<Object> list, @Nullable Object element) {
+ private boolean contains(Collection<Object> list, @Nullable Object
element) {
if (element == null) {
return list.contains(null);
}
for (Object o : list) {
- if (equaliser.apply(o, element)) {
+ if (equals(o, element)) {
return true;
}
}
diff --git
a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/FieldIgnoreRetractAgg.java
b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/FieldIgnoreRetractAgg.java
index e98e64852b..341bb19d8e 100644
---
a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/FieldIgnoreRetractAgg.java
+++
b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/FieldIgnoreRetractAgg.java
@@ -35,6 +35,11 @@ public class FieldIgnoreRetractAgg extends FieldAggregator {
return aggregator.agg(accumulator, inputField);
}
+ @Override
+ public Object aggReversed(Object accumulator, Object inputField) {
+ return aggregator.aggReversed(accumulator, inputField);
+ }
+
@Override
public void reset() {
aggregator.reset();
diff --git
a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/factory/FieldCollectAggFactory.java
b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/factory/FieldCollectAggFactory.java
index a4325d165b..6b656073e6 100644
---
a/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/factory/FieldCollectAggFactory.java
+++
b/paimon-core/src/main/java/org/apache/paimon/mergetree/compact/aggregate/factory/FieldCollectAggFactory.java
@@ -37,7 +37,10 @@ public class FieldCollectAggFactory implements
FieldAggregatorFactory {
"Data type for collect column must be 'Array' but was '%s'.",
fieldType);
return new FieldCollectAgg(
- identifier(), (ArrayType) fieldType,
options.fieldCollectAggDistinct(field));
+ identifier(),
+ (ArrayType) fieldType,
+ options.fieldCollectAggDistinct(field),
+ options.fieldCollectAggCountLimit(field));
}
@Override
diff --git
a/paimon-core/src/test/java/org/apache/paimon/mergetree/compact/aggregate/FieldAggregatorTest.java
b/paimon-core/src/test/java/org/apache/paimon/mergetree/compact/aggregate/FieldAggregatorTest.java
index b915866e7f..4e9e2852af 100644
---
a/paimon-core/src/test/java/org/apache/paimon/mergetree/compact/aggregate/FieldAggregatorTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/mergetree/compact/aggregate/FieldAggregatorTest.java
@@ -2139,6 +2139,51 @@ public class FieldAggregatorTest {
return GenericRow.of(k0, k1, BinaryString.fromString(v), seq, ts);
}
+ @Test
+ public void testFieldCollectAggWithCountLimit() {
+ DataType[] elementTypes = {DataTypes.INT(),
RowType.of(DataTypes.INT())};
+ for (DataType elementType : elementTypes) {
+ assertCollectCountLimit(elementType, false, 3, 2, 1, 2);
+ assertCollectCountLimit(elementType, true, 3, 2, 1, 3);
+ assertCollectCountLimit(elementType, false, -3, 3, 3, 1);
+ assertCollectCountLimit(elementType, true, -3, 2, 3, 1);
+ assertCollectCountLimit(elementType, false, 0);
+ assertCollectCountLimit(elementType, true, 0);
+ }
+ }
+
+ private void assertCollectCountLimit(
+ DataType elementType, boolean distinct, int countLimit, int...
expected) {
+ Map<String, String> config = new HashMap<>();
+ config.put("fields.fieldName.distinct", Boolean.toString(distinct));
+ config.put("fields.fieldName.count-limit",
Integer.toString(countLimit));
+ CoreOptions options = CoreOptions.fromMap(config);
+ ArrayType arrayType = DataTypes.ARRAY(elementType);
+ FieldCollectAgg agg = new FieldCollectAggFactory().create(arrayType,
options, "fieldName");
+ InternalArray.ElementGetter elementGetter =
InternalArray.createElementGetter(elementType);
+ Object[] expectedElements = makeArray(elementType,
expected).toObjectArray();
+
+ InternalArray input = makeArray(elementType, 2, 1, 2, 3, 3, 1);
+ InternalArray single = (InternalArray) agg.agg(null, input);
+ assertThat(unnest(single,
elementGetter)).containsExactly(expectedElements);
+
+ InternalArray older = makeArray(elementType, 2, 1, 2);
+ InternalArray newer = makeArray(elementType, 3, 3, 1);
+ InternalArray forward = (InternalArray) agg.agg(agg.agg(null, older),
newer);
+ InternalArray reversed =
+ (InternalArray) agg.aggReversed(agg.aggReversed(null, newer),
older);
+ assertThat(unnest(forward,
elementGetter)).containsExactly(expectedElements);
+ assertThat(unnest(reversed,
elementGetter)).containsExactly(expectedElements);
+ }
+
+ private GenericArray makeArray(DataType elementType, int... values) {
+ Object[] elements = new Object[values.length];
+ for (int i = 0; i < values.length; i++) {
+ elements[i] = elementType instanceof RowType ?
GenericRow.of(values[i]) : values[i];
+ }
+ return new GenericArray(elements);
+ }
+
@Test
public void testFieldCollectAggWithDistinct() {
FieldCollectAgg agg =
