Skip to content

Commit 64675aa

Browse files
author
hongli.wwj
committed
[core] Apply LIMIT during merge-file split read when safe.
1 parent 94b468a commit 64675aa

8 files changed

Lines changed: 654 additions & 47 deletions

File tree

‎paimon-common/src/main/java/org/apache/paimon/reader/LimitRecordReader.java‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package org.apache.paimon.reader;
20+
21+
import javax.annotation.Nullable;
22+
23+
import java.io.IOException;
24+
import java.util.concurrent.atomic.AtomicLong;
25+
26+
import static org.apache.paimon.utils.Preconditions.checkArgument;
27+
28+
/** A {@link RecordReader} that stops after reading a given number of records. */
29+
public final class LimitRecordReader<T> implements RecordReader<T> {
30+
31+
private final RecordReader<T> reader;
32+
private final long limit;
33+
private final AtomicLong recordCount = new AtomicLong(0);
34+
35+
public LimitRecordReader(RecordReader<T> reader, long limit) {
36+
checkArgument(limit > 0, "Limit must be positive.");
37+
this.reader = reader;
38+
this.limit = limit;
39+
}
40+
41+
public static <T> RecordReader<T> limit(RecordReader<T> reader, @Nullable Integer limit) {
42+
if (limit == null || limit <= 0) {
43+
return reader;
44+
}
45+
return new LimitRecordReader<>(reader, limit);
46+
}
47+
48+
@Override
49+
@Nullable
50+
public RecordIterator<T> readBatch() throws IOException {
51+
if (recordCount.get() >= limit) {
52+
return null;
53+
}
54+
RecordIterator<T> iterator = reader.readBatch();
55+
if (iterator == null) {
56+
return null;
57+
}
58+
return new LimitRecordIterator<>(iterator);
59+
}
60+
61+
@Override
62+
public void close() throws IOException {
63+
reader.close();
64+
}
65+
66+
private class LimitRecordIterator<T> implements RecordIterator<T> {
67+
68+
private final RecordIterator<T> iterator;
69+
70+
private LimitRecordIterator(RecordIterator<T> iterator) {
71+
this.iterator = iterator;
72+
}
73+
74+
@Override
75+
@Nullable
76+
public T next() throws IOException {
77+
if (recordCount.get() >= limit) {
78+
return null;
79+
}
80+
T next = iterator.next();
81+
if (next != null) {
82+
recordCount.incrementAndGet();
83+
}
84+
return next;
85+
}
86+
87+
@Override
88+
public void releaseBatch() {
89+
iterator.releaseBatch();
90+
}
91+
}
92+
}

‎paimon-core/src/main/java/org/apache/paimon/operation/MergeFileSplitRead.java‎

Lines changed: 65 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
package org.apache.paimon.operation;
2020

2121
import org.apache.paimon.CoreOptions;
22+
import org.apache.paimon.CoreOptions.MergeEngine;
2223
import org.apache.paimon.KeyValue;
2324
import org.apache.paimon.KeyValueFileStore;
2425
import org.apache.paimon.data.BinaryRow;
@@ -41,6 +42,7 @@
4142
import org.apache.paimon.mergetree.compact.MergeFunctionWrapper;
4243
import org.apache.paimon.mergetree.compact.ReducerMergeFunctionWrapper;
4344
import org.apache.paimon.predicate.Predicate;
45+
import org.apache.paimon.reader.LimitRecordReader;
4446
import org.apache.paimon.reader.ReaderSupplier;
4547
import org.apache.paimon.reader.RecordReader;
4648
import org.apache.paimon.schema.TableSchema;
@@ -82,13 +84,16 @@ public class MergeFileSplitRead implements SplitRead<KeyValue> {
8284
private final MergeSorter mergeSorter;
8385
private final List<String> sequenceFields;
8486
private final boolean sequenceOrder;
87+
private final CoreOptions coreOptions;
8588

8689
@Nullable private RowType readKeyType;
8790
@Nullable private RowType outerReadType;
8891

8992
@Nullable private List<Predicate> filtersForKeys;
9093
@Nullable private List<Predicate> filtersForAll;
9194

95+
@Nullable private Integer limit;
96+
9297
private boolean forceKeepDelete = false;
9398

9499
public MergeFileSplitRead(
@@ -100,6 +105,7 @@ public MergeFileSplitRead(
100105
MergeFunctionFactory<KeyValue> mfFactory,
101106
KeyValueFileReaderFactory.Builder readerFactoryBuilder) {
102107
this.tableSchema = schema;
108+
this.coreOptions = options;
103109
this.readerFactoryBuilder = readerFactoryBuilder;
104110
this.fileIO = readerFactoryBuilder.fileIO();
105111
this.keyComparator = keyComparator;
@@ -177,6 +183,12 @@ public MergeFileSplitRead forceKeepDelete() {
177183
return this;
178184
}
179185

186+
@Override
187+
public MergeFileSplitRead withLimit(@Nullable Integer limit) {
188+
this.limit = limit;
189+
return this;
190+
}
191+
180192
@Override
181193
public MergeFileSplitRead withFilter(Predicate predicate) {
182194
if (predicate == null) {
@@ -312,6 +324,7 @@ public RecordReader<KeyValue> createMergeReader(
312324
reader = new DropDeleteReader(reader);
313325
}
314326

327+
reader = LimitRecordReader.limit(reader, effectiveReadLimit());
315328
return projectOuter(projectKey(reader));
316329
}
317330

@@ -334,7 +347,58 @@ public RecordReader<KeyValue> createNoMergeReader(
334347
suppliers.add(() -> readerFactory.createRecordReader(file));
335348
}
336349

337-
return projectOuter(ConcatRecordReader.create(suppliers));
350+
return LimitRecordReader.limit(
351+
projectOuter(ConcatRecordReader.create(suppliers)), effectiveReadLimit());
352+
}
353+
354+
/**
355+
* Limit on merge read is only safe when filters are fully applied before truncation. This
356+
* aligns with {@link KeyValueFileStoreScan#limitPushdownEnabled()} and additionally rejects
357+
* non-primary-key filters, which are not pushed down to overlapping L0 sections (see {@link
358+
* #withFilter(Predicate)}).
359+
*/
360+
@Nullable
361+
private Integer effectiveReadLimit() {
362+
if (limit == null || limit <= 0) {
363+
return null;
364+
}
365+
366+
if (forceKeepDelete) {
367+
return null;
368+
}
369+
370+
MergeEngine mergeEngine = coreOptions.mergeEngine();
371+
if (mergeEngine == MergeEngine.PARTIAL_UPDATE || mergeEngine == MergeEngine.AGGREGATE) {
372+
return null;
373+
}
374+
375+
if (coreOptions.deletionVectorsEnabled()) {
376+
return null;
377+
}
378+
379+
if (hasNonPrimaryKeyFilter()) {
380+
return null;
381+
}
382+
383+
return limit;
384+
}
385+
386+
private boolean hasNonPrimaryKeyFilter() {
387+
if (filtersForAll == null || filtersForAll.isEmpty()) {
388+
return false;
389+
}
390+
391+
List<String> primaryKeys = tableSchema.trimmedPrimaryKeys();
392+
Set<String> nonPrimaryKeys =
393+
tableSchema.fieldNames().stream()
394+
.filter(name -> !primaryKeys.contains(name))
395+
.collect(Collectors.toSet());
396+
for (Predicate filter : filtersForAll) {
397+
if (containsFields(filter, nonPrimaryKeys)) {
398+
return true;
399+
}
400+
}
401+
return false;
338402
}
339403

340404
/**

‎paimon-core/src/main/java/org/apache/paimon/operation/SplitRead.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,18 @@ public SplitRead<R> withFilter(@Nullable Predicate predicate) {
8383
return this;
8484
}
8585

86+
@Override
87+
public SplitRead<R> withTopN(@Nullable TopN topN) {
88+
read.withTopN(topN);
89+
return this;
90+
}
91+
92+
@Override
93+
public SplitRead<R> withLimit(@Nullable Integer limit) {
94+
read.withLimit(limit);
95+
return this;
96+
}
97+
8698
@Override
8799
public RecordReader<R> createReader(Split split) throws IOException {
88100
return splitConvert.apply(split);

‎paimon-core/src/main/java/org/apache/paimon/table/format/FormatTableRead.java‎

Lines changed: 2 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import org.apache.paimon.metrics.MetricRegistry;
2424
import org.apache.paimon.predicate.Predicate;
2525
import org.apache.paimon.predicate.PredicateProjectionConverter;
26+
import org.apache.paimon.reader.LimitRecordReader;
2627
import org.apache.paimon.reader.RecordReader;
2728
import org.apache.paimon.table.FormatTable;
2829
import org.apache.paimon.table.source.Split;
@@ -31,7 +32,6 @@
3132

3233
import java.io.IOException;
3334
import java.util.Optional;
34-
import java.util.concurrent.atomic.AtomicLong;
3535

3636
/** A {@link TableRead} implementation for {@link FormatTable}. */
3737
public class FormatTableRead implements TableRead {
@@ -80,11 +80,7 @@ public RecordReader<InternalRow> createReader(Split split) throws IOException {
8080
if (executeFilter) {
8181
reader = executeFilter(reader);
8282
}
83-
if (limit != null && limit > 0) {
84-
reader = applyLimit(reader, limit);
85-
}
86-
87-
return reader;
83+
return LimitRecordReader.limit(reader, limit);
8884
}
8985

9086
private RecordReader<InternalRow> executeFilter(RecordReader<InternalRow> reader) {
@@ -106,44 +102,4 @@ private RecordReader<InternalRow> executeFilter(RecordReader<InternalRow> reader
106102
Predicate finalFilter = predicate;
107103
return reader.filter(finalFilter::test);
108104
}
109-
110-
private RecordReader<InternalRow> applyLimit(RecordReader<InternalRow> reader, int limit) {
111-
return new RecordReader<InternalRow>() {
112-
private final AtomicLong recordCount = new AtomicLong(0);
113-
114-
@Override
115-
public RecordIterator<InternalRow> readBatch() throws IOException {
116-
if (recordCount.get() >= limit) {
117-
return null;
118-
}
119-
RecordIterator<InternalRow> iterator = reader.readBatch();
120-
if (iterator == null) {
121-
return null;
122-
}
123-
return new RecordIterator<InternalRow>() {
124-
@Override
125-
public InternalRow next() throws IOException {
126-
if (recordCount.get() >= limit) {
127-
return null;
128-
}
129-
InternalRow next = iterator.next();
130-
if (next != null) {
131-
recordCount.incrementAndGet();
132-
}
133-
return next;
134-
}
135-
136-
@Override
137-
public void releaseBatch() {
138-
iterator.releaseBatch();
139-
}
140-
};
141-
}
142-
143-
@Override
144-
public void close() throws IOException {
145-
reader.close();
146-
}
147-
};
148-
}
149105
}

‎paimon-core/src/main/java/org/apache/paimon/table/source/KeyValueTableRead.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,15 @@
2121
import org.apache.paimon.CoreOptions;
2222
import org.apache.paimon.KeyValue;
2323
import org.apache.paimon.annotation.VisibleForTesting;
24+
import org.apache.paimon.catalog.TableQueryAuthResult;
2425
import org.apache.paimon.data.InternalRow;
2526
import org.apache.paimon.disk.IOManager;
2627
import org.apache.paimon.operation.MergeFileSplitRead;
2728
import org.apache.paimon.operation.RawFileSplitRead;
2829
import org.apache.paimon.operation.SplitRead;
2930
import org.apache.paimon.predicate.Predicate;
3031
import org.apache.paimon.predicate.TopN;
32+
import org.apache.paimon.reader.LimitRecordReader;
3133
import org.apache.paimon.reader.RecordReader;
3234
import org.apache.paimon.schema.TableSchema;
3335
import org.apache.paimon.table.source.splitread.IncrementalChangelogReadProvider;
@@ -134,6 +136,30 @@ public InnerTableRead withLimit(int limit) {
134136
return this;
135137
}
136138

139+
@Override
140+
public RecordReader<InternalRow> createReader(List<Split> splits) throws IOException {
141+
return LimitRecordReader.limit(super.createReader(splits), limit);
142+
}
143+
144+
@Override
145+
public RecordReader<InternalRow> createReader(Split split) throws IOException {
146+
if (limit == null || limit <= 0 || !hasLateAppliedRowFilter(split)) {
147+
return LimitRecordReader.limit(super.createReader(split), limit);
148+
}
149+
150+
Integer savedLimit = this.limit;
151+
this.limit = null;
152+
initialized().forEach(r -> r.withLimit(null));
153+
try {
154+
return LimitRecordReader.limit(super.createReader(split), savedLimit);
155+
} finally {
156+
this.limit = savedLimit;
157+
if (savedLimit != null) {
158+
initialized().forEach(r -> r.withLimit(savedLimit));
159+
}
160+
}
161+
}
162+
137163
@Override
138164
public TableRead withIOManager(IOManager ioManager) {
139165
initialized().forEach(r -> r.withIOManager(ioManager));
@@ -152,6 +178,14 @@ public RecordReader<InternalRow> reader(Split split) throws IOException {
152178
throw new RuntimeException("Should not happen.");
153179
}
154180

181+
private static boolean hasLateAppliedRowFilter(Split split) {
182+
if (split instanceof QueryAuthSplit) {
183+
TableQueryAuthResult authResult = ((QueryAuthSplit) split).authResult();
184+
return authResult != null && authResult.extractPredicate() != null;
185+
}
186+
return false;
187+
}
188+
155189
public static RecordReader<InternalRow> unwrap(
156190
RecordReader<KeyValue> reader, Map<String, String> schemaOptions) {
157191
return new RecordReader<InternalRow>() {

0 commit comments

Comments
 (0)