Skip to content

feat(table): plan a PK-vector search from a decoded bucket split - #757

Merged
JingsongLi merged 8 commits into
apache:mainfrom
JunRuiLee:feat/pk-vector-plan-from-bucket-split
Sep 1, 2026
Merged

feat(table): plan a PK-vector search from a decoded bucket split#757
JingsongLi merged 8 commits into
apache:mainfrom
JunRuiLee:feat/pk-vector-plan-from-bucket-split

Conversation

@JunRuiLee

@JunRuiLee JunRuiLee commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

Third step of #755: run a primary-key vector search over bucket splits an engine planned elsewhere, not only over a plan read from this table's index manifest.

Independent of #752 — the diff is against main alone.

PkVectorScan::plan_for_bucket_vector_splits takes the splits as authoritative: their payloads, row ranges and pinned snapshot are used as given, no manifest is read, and only the partition conjuncts are re-applied, since a caller may narrow further than the planner that produced them. Bucket grouping and segment selection reuse the manifest route.

Four decisions worth checking:

  • Path resolution. _EXTERNAL_PATH wins; otherwise the bucket directory the split serialized when index-file-in-data-file-dir is set, else <table>/index. Java records an external path only for an index stored outside the table — PkVectorAnnSegmentFile writes null otherwise — so an ordinary bucket-local payload carries none. This duplicates only the small path decision fix(table): resolve index files by external path and bucket layout #752 centralizes; whichever lands second should consolidate this call site, and the option accessor added here is byte-identical to fix(table): resolve index files by external path and bucket layout #752's.
  • The read is limited to the rows the split allows. The plan carries the selection as normalized Vec<RowRange>, and both paths that touch a data file read through those ranges: the residual predicate, and the exact fallback — which needs no predicate to be handed a narrow split. Positions are recovered by walking the selection in step with the emitted rows, and the two are checked against each other, so an over-long read fails loudly instead of returning a filtered answer.
  • Absence means opposite things in the two formats. Java omits files its pre-filter did not narrow, while the search kernel reads a missing entry as "no rows". An omitted file therefore becomes an explicit full-file range; an explicitly empty list still excludes its file, and such a file is registered empty without being opened.
  • Rejected rather than worked around: no splits, splits pinning different snapshots (checked before pruning, so an inconsistent input cannot hide behind an empty plan), two splits for one bucket, a nested data split carrying its own row ranges, and an unlisted file whose row count is unknown — reading ROW_COUNT_UNKNOWN as "no rows" would drop the file from the search silently.

A payload's deletion_vectors_ranges is ignored: Java reserves that field for deletion-vector index files and takes a read's deletion vectors from the data split.

What the range read saves is format-dependent, and the reader's contract says so: mosaic skips a row group before touching its column data, parquet skips pages through the offset index, .row prunes blocks. Avro is the exception — it loads the whole file and deserializes every record before applying the selection, so there a narrow selection saves only what comes after decoding. Making Avro prune physically needs a block-aware reader and is not part of this PR.

Deliberately not here:

  • plan_for_bucket_vector_splits has no in-tree caller yet — the C entry point is the next step, and it carries #[allow(dead_code)] meanwhile. Happy to fold the caller in if you would rather not merge an uncalled entry point.
  • An unrestricted split still materializes a whole-file allow-list for the search kernel, where Java returns no mask at all, which costs Lumina its batch search path. Saying "no mask" instead changes the plan's own type, so it belongs with the C entry point that first makes this route reachable.

@JunRuiLee
JunRuiLee force-pushed the feat/pk-vector-plan-from-bucket-split branch from 2ef30b3 to aaac4b9 Compare August 28, 2026 07:47
`plan_and_search_pk_candidates_batch` resolved the query parameters, read the
index manifest into a plan, and searched that plan in one body, so a caller that
already holds a plan could not reuse the search path.

Split it into three pieces with no behavior change:

- `resolve_pk_vector_search_params` — the query-level parameters and the
  pre-filter guard: everything resolvable from the schema, the options and the
  queries alone, before planning.
- `search_pk_raw_candidates_batch_with_plan` — search a supplied plan and return
  each query's raw indexed and exact candidate lists. Plan-dependent concurrency
  (segment count, batch-index parallelism, range-read bound) is derived from the
  plan actually being searched, so a narrowed plan can never be searched under
  limits computed for a wider one.
- `search_pk_candidates_batch_with_plan` — the raw layer plus the optional exact
  rerank of the approximate candidates and the merge into one best-first list per
  query.

`plan_and_search_pk_candidates_batch` keeps its signature and becomes a wrapper
over the three. The empty-plan short circuit moves into the raw layer, still
ahead of backend resolution, so a table with no searchable data does not error on
an unrecognized index type.
A `BucketVectorSearchSplit` already carries everything a search needs for one
bucket: the payload files, the rows each data file allows, and the snapshot the
whole plan is pinned to. Planning could only be driven the other way round, by
reading this table's index manifest, so a search could not be run over splits an
engine planned elsewhere.

`PkVectorScan::plan_for_bucket_vector_splits` builds a plan from such splits
instead. The splits are authoritative -- no manifest is read -- and only the
partition conjuncts of the scan's filter are re-applied, since a caller may
narrow the query further than the planner that produced the splits. Bucket
grouping, current-segment selection and exact-fallback eligibility reuse the
manifest route's `plan_from_inputs`, so both routes pick segments the same way.

A payload's index file is resolved where Java put it: its `_EXTERNAL_PATH` when it
records one, and otherwise the bucket directory the split serialized when the
table sets `index-file-in-data-file-dir`, or the table `index/` directory. Java
records an external path only for an index stored outside the table --
`PkVectorAnnSegmentFile` writes `null` unless its path factory is external -- so an
ordinary bucket-local payload carries none, and looking for it under `index/`
would not find it. The manifest route keeps its existing `index/` assumption; the
option is read only for splits an engine supplied.

Four inputs are rejected rather than planned around:

- No splits at all, which pins no snapshot to report, and the plan's snapshot id
  has to stay authoritative even when nothing is searchable.
- Splits pinning different snapshots, checked before partition pruning so an
  inconsistent input cannot hide behind an empty plan.
- Two splits for one bucket, which would search its rows twice. Java emits one
  split per bucket, but independently decoded buffers cannot enforce that.
- A nested data split carrying its own row ranges, which would be a second
  authority over which physical rows are readable, free to disagree with the
  per-file ranges the bucket form carries.

Row ranges become a per-split allow-list of physical positions on the plan, and
the search intersects it with the residual predicate's allow-list: both sides
list what is permitted, so a position needs to survive both. The normalization is
where the two formats disagree -- Java records ranges only for the files its own
pre-filter narrowed and omits the rest, while the search kernel reads a missing
entry as "no rows allowed" -- so an omitted file is turned into an explicit
full-file range. An empty list stays empty and excludes its file.

A payload's `deletion_vectors_ranges` is ignored on purpose. Java reserves that
field for deletion-vector index files, builds vector payloads through the
overload that leaves it null, and takes a read's deletion vectors from the
bucket's data split, so a value there describes something the payload is not.

Planning from the Java golden fixture is covered end to end: the external payload
path wins over both directory layouts, the five-billion-byte size survives, and a
six-row file listed as rows 0-1 and 4-5 plans to exactly those positions.
A producer that restricts the readable rows of only some data files leaves the
rest unrestricted, and an adapter has to say so explicitly, as an allow-list
covering the whole file. Building live row ids then walked that list one position
at a time, costing an insert per row of the file, where the same statement made by
omitting the residual entirely takes a single range insert.

Recognize the whole-file shape and insert one range instead. A list whose length
equals the file's row count and whose maximum is the last position can only be the
full set, so the check also subsumes the per-position bound check it replaces.
Deletion vectors still apply: the shortcut only replaces how positions enter the
live set, not what happens to them afterwards.
- `plan_for_bucket_vector_splits` carried its whole doc block twice.
- `into_parts` claimed the index metadata moves rather than being cloned, which
  `plan_from_inputs` does not honor: it clones `_INDEX_META` out of the
  `GlobalIndexMeta` it is handed. Narrowed the claim to what the method itself does.
- `manifest_planning_leaves_positions_unrestricted` could not observe what its name
  claims: it drives `plan_from_inputs`, which returns splits and not a
  `PkVectorScanPlan`, so `physical_row_ranges_by_split` is out of its reach and the
  only assertion left was a split count that `builds_one_split_per_bucket_with_data`
  already makes.
@JunRuiLee
JunRuiLee force-pushed the feat/pk-vector-plan-from-bucket-split branch from aaac4b9 to 90433b9 Compare August 31, 2026 09:57
@JunRuiLee
JunRuiLee marked this pull request as ready for review August 31, 2026 09:59
// built from engine-supplied bucket splits carries the physical positions each
// file is limited to; a plan read from the index manifest carries none. Both
// sides list what is permitted, so combining them is an intersection.
let residual_by_split = intersect_row_allow_lists(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Apply the supplied physical ranges before scanning residuals

For an engine-supplied bucket split, plan.physical_row_ranges_by_split can restrict a large data file to a very small range. When a data predicate is also present, the code calls residual_positions_by_file first; that function invokes read_single_file_stream(..., None, None) and scans every physical row in each active file. Only here, afterward, are positions outside the supplied range discarded. Thus a 10-row split over a billion-row file turns into a full-file read for every residual-filtered vector query, defeating the distributed split and creating an extreme I/O/latency regression. Java evaluates the residual through an IndexedSplit built from the candidate ranges. Please pass the physical allow-list into residual evaluation, skip empty files, and read/evaluate only those ranges; a counting-reader test should assert that rows outside the engine range are never read.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 3b9ba6f. The plan now carries the selection as normalized Vec<RowRange> rather than materialized positions, and both paths that touch a data file read through those ranges — the residual, and the exact fallback, which you did not mention but has the same defect and does not even need a predicate to hit it. Positions are recovered by walking the selection in step with the emitted rows and the two are checked against each other, so a read that overshoots fails loudly instead of quietly returning a filtered answer. A file the plan lists no rows for is registered empty without being opened.

DataFileReader already had a file-local ranges path — it coalesced positions into ranges internally — so it grew a ranges entry point rather than a new one. I did not use the generic row_ranges argument: that one translates through first_row_id, which is not the file-local space these ranges are in.

On the counting test: it is at the format layer, where the reading happens. A selection inside one row group requests strictly fewer byte ranges and fewer bytes than a full read. One caveat I should state rather than let you find — what the selection saves depends on the format. Mosaic skips a row group before touching its column data, parquet skips pages through the offset index, .row prunes blocks, but Avro loads the whole file and deserializes every record before applying the selection, so for an Avro data file the full read remains. Java allows file.format=avro for a vector column, so that is reachable. Making Avro prune physically needs a block-aware reader; I left it stated rather than half-done.

The intersection you pointed at stays. It cannot remove anything once the residual was evaluated over the same ranges, which is the invariant it now records.

An engine-supplied bucket split can restrict a large data file to a handful of
rows, but the residual predicate was evaluated by reading every physical row of
every active file and discarding what fell outside the split's ranges afterwards.
A ten-row split over a billion-row file therefore cost a full scan on every
residual-filtered query, which is the opposite of what the split is for. Java
evaluates the residual through an `IndexedSplit` built from the same candidate
ranges.

The plan now carries the selection as normalized `Vec<RowRange>` rather than
materialized positions, since that is what a read is limited by; expanding a
whole-file range of a large file into positions costs memory no reader needs.
`DataFileReader` grew a ranges entry point beside the positions one — it coalesced
positions into ranges internally anyway — and the residual read goes through it.
Positions are recovered by walking the selection in step with the emitted rows
instead of counting from zero, and the two are checked against each other: with no
pushdown predicate and no deletion vector the read emits exactly what was selected,
so a mismatch means the assumption broke.

A file the plan lists no rows for is registered empty without being read at all.
The intersection against the residual stays: it cannot remove anything once the
residual was evaluated over the same ranges, which is the invariant it now states.
The exact fallback had the same defect as the residual scan, and did not even need
a data predicate to hit it: it opened every active file in full and let
`is_excluded` reject the rows outside the split's ranges afterwards. Scoring is not
the expensive part of that — reading the file is.

`search_file` now takes the plan's selection for the file and reads through the
ranges entry point, recovering each row's physical position from the selection the
same way the residual scan does. `is_excluded` still applies on top, since it also
folds in a residual predicate and the deletion vector, but it can no longer be the
only thing standing between a ten-row split and a full file read. An empty
selection returns without opening the file at all.

Three tests, one per layer, because a result-only assertion cannot tell
range-limited reading from full-read-then-filter:

- the mosaic reader requests strictly fewer byte ranges, and fewer bytes, when the
  selection sits inside one row group. That is the granularity of the win: a row
  group is skipped before its column data is touched, so a narrow range does not
  seek within a row group.
- the residual scan over a selection returns only positions inside it, and the rows
  it excludes would otherwise have matched the predicate.
- the exact fallback returns the allowed positions rather than the nearest rows,
  which are outside the selection.

The first two would fail loudly rather than quietly under a full read: positions come
from walking the selection, so an over-long read exhausts it and reports that.
… saves

Two corrections to the range pushdown.

An unlisted file was read as "the whole file" through a helper that mapped any
non-positive row count to an empty selection, so `DataFileMeta::ROW_COUNT_UNKNOWN`
(-1) silently dropped the file from the search. Before the pushdown that conversion
failed loudly, and it has to keep failing: the protocol decoder only validates the
row count of a file it carries ranges for, which leaves the unlisted ones to this
helper.

What a range read saves is also format-dependent, and the reader's contract now says
so rather than implying it is universal. Mosaic skips a row group before touching its
column data, parquet skips pages through the offset index, `.row` prunes blocks. Avro
is the exception: it loads the whole file and deserializes every record before
applying the selection, so there a narrow selection saves only what comes after
decoding. Java allows `file.format=avro` for a vector column, so that case is
reachable rather than hypothetical; making Avro prune physically needs a block-aware
reader and is not part of this change.

Also narrows the exact fallback's comment about what its row-count check proves: a
selected read vouches for the ranges it asked for and cannot notice a file truncated
elsewhere, which a full read could.

@JingsongLi JingsongLi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1

@JingsongLi
JingsongLi merged commit 120ff9e into apache:main Sep 1, 2026
13 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants