[core] Enforce LIMIT at the table-read level for primary-key tables

[wwj6591812]

Background

ReadBuilder.withLimit propagates LIMIT to TableRead. For primary-key tables,
KeyValueTableRead previously forwarded the limit to lower-level SplitRead
implementations.

This makes the behavior depend on individual split readers. A per-split limit is
also insufficient for multi-split reads, and becomes more complicated when row
filters, such as query authorization filters, are applied outside the split reader.

Since RecordReader and TableRead are lazy, the limit can be enforced directly
on the reader returned by KeyValueTableRead, after table-level processing.

Why this PR

This PR makes the TableRead.withLimit contract explicit for primary-key tables
by applying the limit at the KeyValueTableRead boundary.

The goal is to provide correct and consistent limit behavior with a simpler
execution path. This PR does not claim an additional general SQL runtime
optimization beyond the existing lazy reader behavior.

What changes

• Add a reusable LimitRecordReader in paimon-common.
• Apply the limit to the final reader returned by
  KeyValueTableRead#createReader(Split) and
  KeyValueTableRead#createReader(List<Split>).
• Stop forwarding the limit from KeyValueTableRead to lower-level SplitRead
  implementations. MergeFileSplitRead no longer contains limit-specific logic.
• Apply the limit after table-level filters and query authorization filters.
• Enforce one limit across the concatenated reader for multi-split reads.
• Reuse LimitRecordReader in FormatTableRead as a small refactor with no
  behavior change.

Stage applied: the table-read output boundary, outside individual split readers
and after table-level filtering.

Tests

• PrimaryKeySimpleTableTest — covers empty plans, single-split and multi-split
  reads, non-primary-key filters, and query authorization filters.
• KeyValueTableReadLimitITCase — covers Flink batch SQL LIMIT on a
  primary-key table with multiple level-0 files, including WHERE correctness.

Limitations

• In distributed engines, the global SQL LIMIT is still coordinated by the
  engine-level limit operator.
• This change does not reduce merge work required before the first output row is
  produced.

[JingsongLi]

This override is not reached by the normal KeyValueTableRead path. MergeFileSplitReadProvider wraps this reader with SplitRead.convert(...), and that adapter does not override/forward withLimit, so KeyValueTableRead.config() ends up calling the default no-op method on the wrapper. As a result, SQL/table reads still do not apply the new merge-read limit optimization even though the lower-level store.newRead().withLimit(...) test passes. Please forward withLimit in SplitRead.convert (and add a table-read-path test) so this actually takes effect for production reads.

[wwj6591812]

Thanks for the review, @JingsongLi . You are absolutely right — the withLimit override in MergeFileSplitRead was not reachable through the normal KeyValueTableRead path because MergeFileSplitReadProvider wraps the reader via SplitRead.convert(...), and the adapter did not forward withLimit (and withTopN) to the underlying reader.
I have fixed this by adding withTopN and withLimit forwarding in SplitRead.convert(...). I also added a new test testReadWithLimitThroughTableReadPath in PrimaryKeySimpleTableTest to verify that the limit is actually applied through the full table read path (KeyValueTableRead -> MergeFileSplitReadProvider -> SplitRead.convert). Please take another look.

[JingsongLi]

The limit is scoped to this split reader, but TableRead.createReader(List<Split>) just concatenates one reader per split. If a plan still contains multiple DataSplits (for example multiple buckets/partitions, or any case where scan-side limit pushdown cannot shrink the plan to one split), each split can now emit up to limit rows, so table.newRead().withLimit(10).createReader(plan.splits()) may return 20/30/... rows instead of 10. The new table-read test only covers the default single-bucket plan, so it does not catch this. Please either keep the read-side limit global across the concatenated plan, or only enable this optimization when the plan is guaranteed to contain a single split, and add a multi-split regression test.

[wwj6591812]

Thanks for catching this, @JingsongLi .

You're right — applying the limit inside each split reader is incorrect when TableRead.createReader(List) concatenates multiple splits, because each split could emit up to limit rows.

I've fixed this with a hybrid approach:

Single-split plans: keep the merge-read limit optimization (early truncation during merge).
Multi-split plans: disable the per-split merge-read limit and apply a global LimitRecordReader at the KeyValueTableRead.createReader(List) level, so the total row count is capped correctly.
Added testReadWithLimitThroughTableReadPathMultiSplit in PrimaryKeySimpleTableTest (BUCKET=4, overlapping L0 files across buckets) to verify that withLimit(10).createReader(plan.splits()) returns exactly 10 rows when the plan contains multiple splits.

Please take another look. Thanks!

[JingsongLi]

This only wraps the concatenated reader when there are multiple splits, but MergeFileSplitRead.effectiveReadLimit() can also return null for a single split (for example with a non-primary-key/value filter, forceKeepDelete, partial-update, aggregate, or deletion vectors). In those cases table.newRead().withLimit(10).createReader(plan.splits()) no longer enforces the read limit at all when the plan has one split; the new MergeFileSplitReadTest#testWithLimitDisabledByNonPrimaryKeyFilter even shows the lower-level reader returning all rows once the merge-read optimization is disabled. Please keep a global LimitRecordReader fallback at the table-read level whenever the per-split merge-read limit is not applied, and add a single-split + non-PK filter regression test through KeyValueTableRead.

[wwj6591812]

Thanks for the review, @JingsongLi .

You're right — the global LimitRecordReader wrapper was only applied for multi-split plans, so when a single-split plan hit a case where MergeFileSplitRead.effectiveReadLimit() returns null (e.g. non-PK filter), withLimit(n) no longer enforced the limit at the table-read level.

I've added SplitRead#isMergeReadLimitActive() and updated KeyValueTableRead#createReader(List) to apply a table-level LimitRecordReader fallback whenever merge-read limit is not active (including single-split + non-PK filter). I also added testReadWithLimitThroughTableReadPathWithNonPrimaryKeyFilter in PrimaryKeySimpleTableTest to cover this through the full KeyValueTableRead path. Please take another look when you have a chance.

[JingsongLi]

createConcatenatedReader(splits) is lazy when there is more than one split: it only stores suppliers that call createReader(split) later. This method restores applyMergeReadLimit in the finally block before any of those suppliers are consumed, so the later per-split readers are created with withLimit(limit) enabled again, despite this multi-split path trying to disable it. The outer LimitRecordReader still caps the final count, but rows can be dropped before late filters that are applied outside MergeFileSplitRead (for example row-level auth filters from QueryAuthSplit in AbstractDataTableRead.authedReader): each split may be truncated before the auth filter runs, and the global wrapper cannot recover the filtered-out rows. Please keep merge-read limit disabled until the lazy readers are actually consumed (or materialize the split readers while disabled), and add a multi-split regression test with a late/applied-after-read filter.

[wwj6591812]

Thanks for the careful review, @JingsongLi .

You're right — createReaderWithGlobalLimit restored applyMergeReadLimit in the finally block before the lazy ConcatRecordReader suppliers were consumed, so per-split readers could still be created with merge-read limit enabled again. That could truncate rows before late filters such as QueryAuthSplit auth filtering, and the outer LimitRecordReader cannot recover those rows.

I've fixed this by materializing each split reader while merge-read limit is disabled, then wrapping the concatenated reader with the global LimitRecordReader. I also added testReadWithLimitThroughTableReadPathMultiSplitWithQueryAuthFilter to cover the multi-split + late row-filter case. Please take another look when you have a chance.

[JingsongLi]

This single-split fast path only checks whether the underlying MergeFileSplitRead considers its own pushed filters safe, but it does not know about row filters carried by QueryAuthSplit. Those filters are applied later in AbstractDataTableRead.authedReader() after reader(split) has already been created, so a one-split authorized read can truncate to limit rows during merge read and then drop some of them in the auth filter, returning fewer than the requested limit even when more authorized rows exist later in the split. Please route row-filtering QueryAuthSplits through the table-level limit fallback with merge-read limit disabled, and cover the single-split auth-filter case as well.

[wwj6591812]

Thanks for pointing this out, @JingsongLi .

You're right — the single-split fast path only checked MergeFileSplitRead.isMergeReadLimitActive() and did not account for row filters applied later in AbstractDataTableRead.authedReader() via QueryAuthSplit, which could return fewer than the requested limit.

I've extended the routing logic with needsTableLevelLimit() so splits carrying a row-level QueryAuthSplit filter also go through the table-level limit fallback with merge-read limit disabled. I added testReadWithLimitThroughTableReadPathWithQueryAuthFilter for the single-split auth-filter case. Please take another look when you have a chance.

[JingsongLi]

This check runs before the split reader is normally initialized. In the common table-read path (table.newRead().withLimit(n).createReader(plan.splits())), none of the SplitReadProvider lazy fields has been materialized yet, so initialized() is empty and isMergeReadLimitActive() returns false. That sends even a safe single-split merge read through createReaderWithGlobalLimit, which disables applyMergeReadLimit and means the new merge-read limit optimization does not actually run for the normal production path unless the same TableRead happened to create a reader earlier. Could we initialize/check the provider that matches the single split before deciding this fallback, and add a regression that verifies the safe single-split table-read path really applies the merge-read limit instead of only the outer LimitRecordReader?

[wwj6591812]

Thanks for catching this. You are right — isMergeReadLimitActive() was checked against only already-initialized providers, so in the normal table.newRead().withLimit(n).createReader(plan.splits()) path no provider was materialized yet and safe single-split reads incorrectly fell back to createReaderWithGlobalLimit().

I updated needsTableLevelLimit() to resolve and initialize the provider that matches the single split before deciding the fallback. I also added a regression test (testSingleSplitTableReadUsesMergeReadLimitNotTableLevelLimit) to verify that the safe single-split table-read path keeps merge-read limit active instead of relying only on the outer LimitRecordReader.

[JingsongLi]

This path now assumes the limited read has at least one split. An empty table, a partition/bucket filter that prunes everything, or a scan filter that produces an empty plan can still call table.newRead().withLimit(n).createReader(plan.splits()) with splits empty; then needsTableLevelLimit reaches splits.get(0) and throws IndexOutOfBoundsException instead of returning an empty reader. Please handle the empty-list case before checking the single-split merge-limit status, and add a regression for withLimit over an empty plan.

[wwj6591812]

Thanks for the follow-up. @JingsongLi

You're right — when filters or an empty table produce no splits, needsTableLevelLimit() could call splits.get(0) and throw IndexOutOfBoundsException.

I fixed this by short-circuiting in createReader() when splits is empty (returning an empty reader directly) and adding the same guard in needsTableLevelLimit(). I also added testReadWithLimitOnEmptyPlan as a regression test for withLimit over an empty plan.

[JingsongLi]

createReaderWithGlobalLimit keeps applyMergeReadLimit disabled until the returned reader is closed, but this flag is shared by the reusable KeyValueTableRead instance. A caller can create another reader from the same table read before this global-limit reader is closed, and that second reader will see applyMergeReadLimit=false, route through the fallback, and initialize unrelated split readers with the merge-read limit disabled. If reader.close() throws, the flag is also never restored. This makes the new optimization depend on reader close ordering. Please scope the disablement to the readers created for this fallback (for example, materialize/capture the per-split readers while disabled and restore in a finally, or pass a per-call flag to the suppliers) instead of leaving shared table-read state changed until close.

[wwj6591812]

Thanks for the review, @JingsongLi .

You're right — keeping applyMergeReadLimit disabled until the returned reader is closed was unsafe because this flag is shared by the reusable KeyValueTableRead instance. A second reader created before the first one is closed, or a failed close(), could leave unrelated reads with merge-read limit disabled.

I've fixed this by materializing each split reader while merge-read limit is disabled, wrapping them with the table-level LimitRecordReader, and restoring applyMergeReadLimit in a finally block before returning. The lazy ConcatRecordReader suppliers now return the pre-created split readers, so the disablement is scoped to this fallback path without depending on reader close ordering.

I also added testGlobalLimitReaderRestoresMergeReadLimitImmediately to verify that applyMergeReadLimit is restored immediately after createReader() returns on the multi-split global-limit path. Please take another look when you have a chance.

[JingsongLi]

I don't quite understand why it's so complicated, just judge for low level Read whether it's possible to perform limit pruning inside, is that enough?

[wwj6591812]

I don't quite understand why it's so complicated, just judge for low level Read whether it's possible to perform limit pruning inside, is that enough?

Thanks for the review. @JingsongLi

After re-reading the code, I agree — merge-read limit safety can be decided entirely in MergeFileSplitRead, without orchestration in the upper layer.

Changes:
1、MergeFileSplitRead: uses effectiveReadLimit() / mergeReadLimitDisabledReason() to decide whether limit pruning is safe during merge read.
2、KeyValueTableRead: only passes limit to split reads and wraps multi-split reads with LimitRecordReader for correct SQL semantics.
3、Query-auth splits are the only exception: filters run after split read, so merge-read limit is skipped there.

Removed withApplyMergeReadLimit / isMergeReadLimitActive and related tests; functional UT coverage is unchanged.

[JingsongLi]

For the single-split API this drops the limit completely when query auth has a row filter. The method clears this.limit to avoid applying it before authResult.doAuth(...), but then returns super.createReader(split) directly, so table.newRead().withLimit(10).createReader(queryAuthSplit) can read all authorized rows. The new tests cover createReader(List<Split>), whose outer LimitRecordReader restores the total limit, but callers and source readers also use the public createReader(Split) path. Could we wrap this returned reader with LimitRecordReader.limit(..., savedLimit) after auth filtering, so the limit is still enforced after the late row filter?

[wwj6591812]

Thanks for the review, @JingsongLi.
You're right — in the single-split createReader(Split) path we cleared the merge-read limit for query-auth splits but forgot to re-apply the table-level limit after auth filtering. I've wrapped the returned reader with LimitRecordReader.limit(..., savedLimit) after super.createReader(split), and added testReadWithLimitThroughSingleSplitCreateReaderWithQueryAuthFilter to cover the public single-split API (the existing tests only exercised createReader(List)). Please take another look when you have a chance.

[JingsongLi]

This still drops the limit on the public single-split API whenever the merge-read limit is disabled for reasons other than query auth. For example, table.newRead().withFilter(nonPrimaryKeyFilter).withLimit(10).createReader(split) reaches this branch, MergeFileSplitRead.effectiveReadLimit() returns null because of the non-PK filter, and there is no outer LimitRecordReader, so the reader can return all matching rows. The same applies to other disabled cases such as forceKeepDelete, partial-update/aggregation, or deletion vectors. Please apply the table-level limit fallback for any split whose lower-level read cannot apply the merge-read limit, not only for QueryAuthSplit, and add a createReader(Split) regression for a non-PK filter/disabled-limit case.

[wwj6591812]

Thanks for the review, @JingsongLi.
You're right — the single-split createReader(Split) path only applied the table-level LimitRecordReader fallback for QueryAuthSplit, so when merge-read limit was disabled for other reasons (e.g. non-PK filter, forceKeepDelete, partial-update/aggregation, or deletion vectors), the public single-split API could return all matching rows with no limit enforced. I've generalized the routing via MergeFileSplitRead#isMergeReadLimitActive() so any split whose lower-level read cannot apply merge-read limit uses the table-level fallback, and added testReadWithLimitThroughSingleSplitCreateReaderWithNonPrimaryKeyFilter to cover the public single-split API. Please take another look when you have a chance.

[JingsongLi]

I still don't understand why you need this method. If it can be done on any layer, then limit it. What's wrong?

[wwj6591812]

Thanks for the suggestion.

I simplified this part and removed the isMergeReadLimitActive plumbing from SplitRead / KeyValueTableRead.

Now KeyValueTableRead always keeps a table-level LimitRecordReader as the semantic limit guard, while MergeFileSplitRead still applies the read-side limit internally only when it is safe. For row-filtering QueryAuthSplit, I still disable the lower-level limit during reader creation and apply the table-level limit after auth filtering, because that filter is applied after the split reader. The related limit regression tests pass.

[JingsongLi]

Remove mergeReadLimitDisabledReason and logMergeReadLimitOnce and mergeReadLimitLogEmitted.

[wwj6591812]

Thanks for the suggestion. I removed mergeReadLimitDisabledReason, logMergeReadLimitOnce, and mergeReadLimitLogEmitted, and kept the safety checks directly in effectiveReadLimit() without the one-time logging. I also ran mvn spotless:apply and force-pushed the updated branch.

[JingsongLi]

Thanks for the PR. I have one concern about the actual optimization benefit here.

RecordReader / TableRead are lazy. If the upper Limit operator already pulls only N records and then stops/closes the source, the underlying split reader should not be fully drained anyway. In that case, wrapping MergeFileSplitRead with LimitRecordReader mostly duplicates the upper-layer limit behavior and may not bring a real SQL runtime optimization.

The current tests prove that TableRead.withLimit returns only N rows when the reader is fully consumed, but they do not demonstrate reduced IO / fewer opened data files / less merge work in an actual engine LIMIT query.

Could you clarify the concrete execution path where the Paimon reader would still be drained despite an outer LIMIT? If the main purpose is to fix the TableRead.withLimit contract for PK merge-read, I think the PR description should state that more explicitly. If it is intended as a performance optimization, it would be better to add a benchmark or a test showing fewer opened file readers / merge sections compared with the previous behavior.

[wwj6591812]

Thanks for the PR. I have one concern about the actual optimization benefit here.

RecordReader / TableRead are lazy. If the upper Limit operator already pulls only N records and then stops/closes the source, the underlying split reader should not be fully drained anyway. In that case, wrapping MergeFileSplitRead with LimitRecordReader mostly duplicates the upper-layer limit behavior and may not bring a real SQL runtime optimization.

The current tests prove that TableRead.withLimit returns only N rows when the reader is fully consumed, but they do not demonstrate reduced IO / fewer opened data files / less merge work in an actual engine LIMIT query.

Could you clarify the concrete execution path where the Paimon reader would still be drained despite an outer LIMIT? If the main purpose is to fix the TableRead.withLimit contract for PK merge-read, I think the PR description should state that more explicitly. If it is intended as a performance optimization, it would be better to add a benchmark or a test showing fewer opened file readers / merge sections compared with the previous behavior.

Thanks for pointing this out. I agree that an outer SQL LIMIT can already stop a lazy TableRead / RecordReader, so this PR should not be presented as a general SQL runtime optimization. I updated the focus to the TableRead.withLimit contract for PK merge-read when the returned reader is fully consumed. I also added a focused test with a tracking FileReaderFactory: with the split-level limit it only opens the first needed file reader, while the unbounded merge read opens all file readers when drained. Filters still disable this split-level limit for correctness.

[wwj6591812]

@JingsongLi

Thanks for the suggestion. I revisited the execution path and changed the implementation to handle LIMIT directly at the outer KeyValueTableRead boundary.

The limit is no longer forwarded to SplitRead or handled in MergeFileSplitRead. KeyValueTableRead now wraps the final reader returned by createReader(Split) and createReader(List) with LimitRecordReader, after table-level filters and query authorization. The list path also enforces one limit across all concatenated splits.

This relies on the existing lazy reader behavior and keeps the PR focused on the TableRead.withLimit contract rather than claiming an additional general SQL runtime optimization. I have also updated the PR title and description
accordingly.

Please take another look. Thanks.

[JingsongLi]

+1