[python] Stream batch vector refinement candidates

[TheR1sing3un]

Purpose

Batch ANN refinement reads the union of candidate vectors into a complete Arrow table, expands every vector into Python lists, and retains a row-ID-to-vector dictionary while scoring each query. Large candidate unions cause substantial memory amplification.

Stream candidate Arrow batches and score only the rows belonging to each query, updating bounded Top-K heaps as each batch arrives. Candidate membership is shared across split workers, existing table read parallelism is retained, and each worker owns one streaming reader with bounded scoring blocks. Snapshot selection, score calculations, null handling, and row-ID tie breaking remain unchanged.

Tests

• python -m pytest -q pypaimon/tests/batch_vector_streaming_refine_test.py pypaimon/tests/vector_search_filter_test.py pypaimon/tests/vector_scoring_test.py pypaimon/tests/batch_vector_raw_scan_test.py pypaimon/tests/vector_metric_consistency_test.py: 119 passed.
• New coverage checks all three metrics, sliced List/LargeList/FixedSizeList arrays, overlapping and empty candidate sets, missing/null rows, invalid non-candidates, bounded scoring blocks, batch release before the next read, snapshot forwarding, parallel split merging, and reader cleanup on scoring/read failures.
• Flake8, license-header checks, and git diff --check passed.

Benchmark

macOS 26.4.1 arm64, Python 3.9.6, NumPy 2.0.2, PyArrow 19.0.1. Each query has 512 candidates and returns Top-K=64 with refine factor 8. The temporary harness reads real, uncompressed local Parquet data in 1,024-row batches and measures candidate reading, conversion, scoring, and Top-K selection. ANN candidate generation and Paimon scan planning are excluded. Each variant runs in three fresh processes with warmed filesystem pages and one read worker; values are medians and RSS is peak process resident memory.

The materialized-Arrow ablation keeps all candidate Arrow data resident but uses the new Arrow scoring path. The streaming variant additionally reads and releases bounded batches.

Workload	Variant	Time (s)	Peak RSS (MiB)
16 queries x 128 dimensions; 8,192 unique candidates	Existing Python lists	0.164936	204.9
Same	Materialized Arrow	0.136279	161.6
Same	Streaming Arrow	0.145794	153.7
64 queries x 384 dimensions; 32,768 unique candidates	Existing Python lists	1.845079	913.4
Same	Materialized Arrow	0.226318	297.9
Same	Streaming Arrow	0.226809	172.4
64 queries x 384 dimensions; 512 shared candidates	Existing Python lists	0.255893	159.1
Same	Materialized Arrow	0.181830	157.9
Same	Streaming Arrow	0.179778	157.9

For the large disjoint candidate workload, the final implementation is 8.14x faster and reduces peak RSS by 81.1%. The ablation shows that avoiding Python vector expansion supplies most of the speedup, while streaming further reduces retained Arrow memory. Small or heavily overlapping candidate sets have smaller memory benefits. All variants returned identical row IDs and score bytes.

End-to-end gains depend on candidate overlap, vector dimensions, storage latency, batch size, and read parallelism. Multiple read workers retain one batch and per-query heaps per worker.