Skip to content

Commit 50f7d50

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

8 files changed

Lines changed: 648 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: 102 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;
@@ -54,6 +56,9 @@
5456
import org.apache.paimon.utils.ProjectedRow;
5557
import org.apache.paimon.utils.UserDefinedSeqComparator;
5658

59+
import org.slf4j.Logger;
60+
import org.slf4j.LoggerFactory;
61+
5762
import javax.annotation.Nullable;
5863

5964
import java.io.IOException;
@@ -74,6 +79,8 @@
7479
*/
7580
public class MergeFileSplitRead implements SplitRead<KeyValue> {
7681

82+
private static final Logger LOG = LoggerFactory.getLogger(MergeFileSplitRead.class);
83+
7784
private final TableSchema tableSchema;
7885
private final FileIO fileIO;
7986
private final KeyValueFileReaderFactory.Builder readerFactoryBuilder;
@@ -82,14 +89,18 @@ public class MergeFileSplitRead implements SplitRead<KeyValue> {
8289
private final MergeSorter mergeSorter;
8390
private final List<String> sequenceFields;
8491
private final boolean sequenceOrder;
92+
private final CoreOptions coreOptions;
8593

8694
@Nullable private RowType readKeyType;
8795
@Nullable private RowType outerReadType;
8896

8997
@Nullable private List<Predicate> filtersForKeys;
9098
@Nullable private List<Predicate> filtersForAll;
9199

100+
@Nullable private Integer limit;
101+
92102
private boolean forceKeepDelete = false;
103+
private boolean mergeReadLimitLogEmitted = false;
93104

94105
public MergeFileSplitRead(
95106
CoreOptions options,
@@ -100,6 +111,7 @@ public MergeFileSplitRead(
100111
MergeFunctionFactory<KeyValue> mfFactory,
101112
KeyValueFileReaderFactory.Builder readerFactoryBuilder) {
102113
this.tableSchema = schema;
114+
this.coreOptions = options;
103115
this.readerFactoryBuilder = readerFactoryBuilder;
104116
this.fileIO = readerFactoryBuilder.fileIO();
105117
this.keyComparator = keyComparator;
@@ -177,6 +189,12 @@ public MergeFileSplitRead forceKeepDelete() {
177189
return this;
178190
}
179191

192+
@Override
193+
public MergeFileSplitRead withLimit(@Nullable Integer limit) {
194+
this.limit = limit;
195+
return this;
196+
}
197+
180198
@Override
181199
public MergeFileSplitRead withFilter(Predicate predicate) {
182200
if (predicate == null) {
@@ -312,6 +330,7 @@ public RecordReader<KeyValue> createMergeReader(
312330
reader = new DropDeleteReader(reader);
313331
}
314332

333+
reader = LimitRecordReader.limit(reader, effectiveReadLimit());
315334
return projectOuter(projectKey(reader));
316335
}
317336

@@ -334,7 +353,89 @@ public RecordReader<KeyValue> createNoMergeReader(
334353
suppliers.add(() -> readerFactory.createRecordReader(file));
335354
}
336355

337-
return projectOuter(ConcatRecordReader.create(suppliers));
356+
return LimitRecordReader.limit(
357+
projectOuter(ConcatRecordReader.create(suppliers)), effectiveReadLimit());
358+
}
359+
360+
/**
361+
* Limit on merge read is only safe when filters are fully applied before truncation. This
362+
* aligns with {@link KeyValueFileStoreScan#limitPushdownEnabled()} and additionally rejects
363+
* non-primary-key filters, which are not pushed down to overlapping L0 sections (see {@link
364+
* #withFilter(Predicate)}).
365+
*/
366+
@Nullable
367+
private Integer effectiveReadLimit() {
368+
if (limit == null || limit <= 0) {
369+
return null;
370+
}
371+
372+
String disabledReason = mergeReadLimitDisabledReason();
373+
if (disabledReason != null) {
374+
logMergeReadLimitOnce(disabledReason);
375+
return null;
376+
}
377+
378+
logMergeReadLimitOnce(null);
379+
return limit;
380+
}
381+
382+
/** Returns why merge read limit is unsafe, or {@code null} if limit can be applied. */
383+
@Nullable
384+
private String mergeReadLimitDisabledReason() {
385+
if (forceKeepDelete) {
386+
return "forceKeepDelete is enabled";
387+
}
388+
389+
MergeEngine mergeEngine = coreOptions.mergeEngine();
390+
if (mergeEngine == MergeEngine.PARTIAL_UPDATE) {
391+
return "merge-engine is partial-update";
392+
}
393+
if (mergeEngine == MergeEngine.AGGREGATE) {
394+
return "merge-engine is aggregation";
395+
}
396+
397+
if (coreOptions.deletionVectorsEnabled()) {
398+
return "deletion-vectors is enabled";
399+
}
400+
401+
if (hasNonPrimaryKeyFilter()) {
402+
return "non-primary-key filter is present";
403+
}
404+
405+
return null;
406+
}
407+
408+
private void logMergeReadLimitOnce(@Nullable String disabledReason) {
409+
if (mergeReadLimitLogEmitted) {
410+
return;
411+
}
412+
mergeReadLimitLogEmitted = true;
413+
if (disabledReason != null) {
414+
LOG.info(
415+
"Merge read limit {} is disabled: {}. Limit will not be applied during merge read.",
416+
limit,
417+
disabledReason);
418+
} else {
419+
LOG.info("Applying merge read limit {} during merge read.", limit);
420+
}
421+
}
422+
423+
private boolean hasNonPrimaryKeyFilter() {
424+
if (filtersForAll == null || filtersForAll.isEmpty()) {
425+
return false;
426+
}
427+
428+
List<String> primaryKeys = tableSchema.trimmedPrimaryKeys();
429+
Set<String> nonPrimaryKeys =
430+
tableSchema.fieldNames().stream()
431+
.filter(name -> !primaryKeys.contains(name))
432+
.collect(Collectors.toSet());
433+
for (Predicate filter : filtersForAll) {
434+
if (containsFields(filter, nonPrimaryKeys)) {
435+
return true;
436+
}
437+
}
438+
return false;
338439
}
339440

340441
/**

‎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
}

0 commit comments

Comments
 (0)