[python] Vectorize raw vector scoring and refinement in bounded blocks

paimon-python/README.md

@@ -302,3 +302,16 @@ only support `seek` and `read` remain serialized. Workers are created lazily
 and released when the index reader closes; separate readers have separate
 budgets. This option controls index I/O, not shard search or native compute
 threads.
+
+
+# Vector fallback scoring and refinement
+
+Raw vector fallback and refinement score regular FLOAT vectors in bounded
+blocks using NumPy. List, large-list and fixed-size-list Arrow arrays are
+supported, including slices and multiple chunks. Null or unsupported blocks
+use the scalar path. Candidate filters are applied before scoring.
+
+L2 and cosine retain scalar accumulation order. Inner product retains Python
+`sum` semantics, including its behavior on newer Python versions. Existing
+Top-K tie-breaking rules are preserved. The same scoring path is used for raw and
+refined primary-key vector results.

paimon-python/pypaimon/table/source/primary_key_vector_read.py

@@ -25,7 +25,7 @@
 from pypaimon.table.source.primary_key_vector_scan import PrimaryKeyVectorScanPlan
 from pypaimon.table.source.vector_search_read import DataEvolutionVectorRead
 from pypaimon.table.source.vector_search_read import (
-    _check_vector_dimension, _compute_score, _to_vector_list)
+    _iter_arrow_scores)
 from pypaimon.read.split import DataSplit
 from pypaimon.globalindex.indexed_split import IndexedSplit
 from pypaimon.deletionvectors.deletion_vector import DeletionVector
@@ -131,7 +131,7 @@ def reranked_iter():
                     if key[:3] == (partition, data_split.bucket, data_file_name))
                 position_iter = iter(positions)
                 for batch in reader.to_arrow([split]).to_batches():
-                    for stored in batch.column(0).to_pylist():
+                    for score in _iter_arrow_scores(batch.column(0), self._query_vector, metric):
                         try:
                             row_position = next(position_iter)
                         except StopIteration:
@@ -144,14 +144,11 @@ def reranked_iter():
                             raise ValueError(
                                 "Primary-key vector rerank read unexpected position %s."
                                 % (key,))
-                        if stored is None:
+                        if score is None:
                             raise ValueError(
                                 "Primary-key vector candidate %s contains a null vector."
                                 % (key,))
-                        stored = _to_vector_list(stored)
-                        _check_vector_dimension(self._query_vector, stored)
-                        yield candidate.with_score(_compute_score(
-                            self._query_vector, stored, metric))
+                        yield candidate.with_score(score)
                 try:
                     next(position_iter)
                     raise ValueError(
@@ -196,21 +193,18 @@ def _raw_candidates(self, plan):
                     if _allowed(split, data_file.file_name, position))
                 position_iter = iter(positions)
                 for batch in reader.to_arrow([read_split]).to_batches():
-                    for stored in batch.column(0).to_pylist():
+                    for score in _iter_arrow_scores(batch.column(0), self._query_vector, metric):
                         try:
                             row_position = next(position_iter)
                         except StopIteration:
                             raise ValueError(
                                 "Raw vector read returned an unexpected row.")
-                        if stored is None:
+                        if score is None:
                             continue
-                        stored = _to_vector_list(stored)
-                        _check_vector_dimension(self._query_vector, stored)
                         yield PrimaryKeySearchPosition(
                             _partition_bytes(split.data_split.partition),
                             split.data_split.bucket, data_file.file_name,
-                            row_position, _compute_score(
-                                self._query_vector, stored, metric))
+                            row_position, score)
                 try:
                     next(position_iter)
                     raise ValueError(

paimon-python/pypaimon/table/source/vector_search_read.py

@@ -306,21 +306,23 @@ def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector,
 
         top_k_heap = []
         metric = self._search_metric(index_type)
-        row_ids = table.column(SpecialFields.ROW_ID.name).to_pylist()
-        vectors = table.column(self._vector_column.name).to_pylist()
-        for row_id, stored in zip(row_ids, vectors):
-            if score_candidates is not None and row_id not in score_candidates:
-                continue
-            if stored is None:
-                continue
-            stored_vector = _to_vector_list(stored)
-            _check_vector_dimension(query_vector, stored_vector)
-            _offer_score(
-                top_k_heap,
-                self._limit,
-                row_id,
-                _compute_score(query_vector, stored_vector, metric),
-            )
+        block_size = _score_block_size(query_vector)
+        for start in range(0, table.num_rows, block_size):
+            block = table.slice(start, block_size)
+            row_ids = block.column(SpecialFields.ROW_ID.name).to_pylist()
+            vectors = block.column(self._vector_column.name)
+            if score_candidates is not None:
+                positions = [i for i, row_id in enumerate(row_ids)
+                             if row_id in score_candidates]
+                if not positions:
+                    continue
+                vectors = vectors.take(positions)
+                row_ids = [row_ids[i] for i in positions]
+            for row_id, score in zip(
+                row_ids, _iter_arrow_scores(vectors, query_vector, metric)
+            ):
+                if score is not None:
+                    _offer_score(top_k_heap, self._limit, row_id, score)
         return _scored_result(top_k_heap)
 
     def _read_raw_vectors(self, candidates, include_filter=True, snapshot=None):
@@ -368,17 +370,24 @@ def _plan_raw_read(self, raw_row_ranges, include_filter, snapshot=None):
 
     def _score_raw_vectors(self, candidates, raw_vectors, query_vector, metric, top_k):
         top_k_heap = []
+        row_ids, vectors = [], []
+
+        def offer_block():
+            for row_id, score in zip(row_ids, _score_rows(vectors, query_vector, metric)):
+                _offer_score(top_k_heap, top_k, row_id, score)
+
+        block_size = _score_block_size(query_vector)
         for row_id in candidates:
             stored_vector = raw_vectors.get(row_id)
             if stored_vector is None:
                 continue
-            _check_vector_dimension(query_vector, stored_vector)
-            _offer_score(
-                top_k_heap,
-                top_k,
-                row_id,
-                _compute_score(query_vector, stored_vector, metric),
-            )
+            row_ids.append(row_id)
+            vectors.append(stored_vector)
+            if len(vectors) == block_size:
+                offer_block()
+                row_ids, vectors = [], []
+        if vectors:
+            offer_block()
         return _scored_result(top_k_heap)
 
     def _read_raw_refine_search(self, candidates, query_vector, index_type=None,
@@ -646,8 +655,7 @@ def _read_raw_batch_search(self, raw_row_ranges, pre_filter,
             return [_scored_result(heap) for heap in heaps]
 
         table_read, splits = self._plan_raw_read(raw_row_ranges, True, snapshot)
-        metric = _raw_search_metric(
-            self._table, self._vector_column, self._options, index_type)
+        metric = self._search_metric(index_type)
         workers = min(len(splits), table_read._resolve_parallelism(None, len(splits)))
         if workers <= 1:
             return self._score_raw_splits(table_read, splits, metric)
@@ -884,6 +892,116 @@ def _normalize_metric(metric):
     return str(metric).lower().replace("-", "_")
 
 
+def _score_block_size(query):
+    # Target 8 MiB per float64 matrix, allowing at least one vector.
+    return max(1, min(1024, (1 << 20) // max(1, len(query))))
+
+
+def _iter_arrow_scores(vectors, query, metric):
+    import numpy as np
+    import pyarrow as pa
+
+    block_size = _score_block_size(query)
+    for start in range(0, len(vectors), block_size):
+        block = vectors.slice(start, block_size)
+        if isinstance(block, pa.ChunkedArray):
+            block = block.combine_chunks()
+        dtype = block.type
+        scores = None
+        if (pa.types.is_list(dtype) or pa.types.is_large_list(dtype) or
+                pa.types.is_fixed_size_list(dtype)) and not block.null_count:
+            flat = block.flatten()
+            if pa.types.is_float32(flat.type) and not flat.null_count:
+                if pa.types.is_fixed_size_list(dtype):
+                    regular = dtype.list_size == len(query)
+                else:
+                    offsets = block.offsets.to_numpy(zero_copy_only=False)
+                    regular = bool(np.all(np.diff(offsets) == len(query)))
+                if regular and len(query):
+                    matrix = flat.to_numpy(zero_copy_only=False).astype(np.float64).reshape(
+                        len(block), len(query))
+                    scores = _compute_scores(query, matrix, metric)
+        if scores is None:
+            for vector in block.to_pylist():
+                if vector is None:
+                    yield None
+                else:
+                    vector = _to_vector_list(vector)
+                    _check_vector_dimension(query, vector)
+                    yield _compute_score(query, vector, metric)
+        else:
+            yield from scores
+
+
+def _score_rows(vectors, query, metric):
+    import numpy as np
+
+    scores = None
+    if vectors and all(vector is not None for vector in vectors):
+        try:
+            matrix = np.array(vectors, dtype=np.float64)
+        except (ValueError, TypeError, OverflowError):
+            matrix = None
+        if matrix is not None:
+            scores = _compute_scores(query, matrix, metric)
+    if scores is not None:
+        return scores
+    result = []
+    for vector in vectors:
+        if vector is None:
+            result.append(None)
+        else:
+            vector = _to_vector_list(vector)
+            _check_vector_dimension(query, vector)
+            result.append(_compute_score(query, vector, metric))
+    return result
+
+
+def _compute_scores(query, matrix, metric):
+    """Score an owned float64 matrix while preserving scalar reduction semantics.
+
+    L2/cosine use left-to-right accumulation, as in the scalar loops. Inner
+    product uses Python's sum on precomputed products, retaining its behavior
+    across Python versions (including compensated summation on Python 3.12+).
+    Returning None selects the scalar fallback.
+    """
+    import numpy as np
+
+    try:
+        query = np.asarray(query, dtype=np.float64)
+    except (ValueError, TypeError, OverflowError):
+        return None
+    if (matrix.ndim != 2 or query.ndim != 1 or not len(query) or
+            matrix.shape[1] != len(query) or not np.isfinite(matrix).all() or
+            not np.isfinite(query).all()):
+        return None
+    if metric not in ("l2", "cosine", "inner_product"):
+        return None
+    with np.errstate(over="ignore", invalid="ignore", divide="ignore"):
+        if metric == "l2":
+            np.subtract(query, matrix, out=matrix)
+            np.square(matrix, out=matrix)
+            np.add.accumulate(matrix, axis=1, out=matrix)
+            return (1.0 / (1.0 + matrix[:, -1])).tolist()
+        if metric == "cosine":
+            norms = np.square(matrix)
+            np.add.accumulate(norms, axis=1, out=norms)
+            stored_norms = norms[:, -1].copy()
+            del norms
+            query_norm = 0.0
+            for value in query:
+                value = float(value)
+                query_norm += value * value
+        np.multiply(matrix, query, out=matrix)
+        if metric == "inner_product":
+            return [sum(row) for row in matrix.tolist()]
+        matrix[:, 0] += 0.0  # Match a scalar accumulator initialized to +0.0.
+        np.add.accumulate(matrix, axis=1, out=matrix)
+        denominators = [query_norm ** 0.5 * float(norm) ** 0.5 for norm in stored_norms]
+        return [0.0 if denominator == 0 else float(dot) / denominator
+                for dot, denominator in zip(matrix[:, -1], denominators)]
+
+
 def _compute_score(query, stored, metric):
     if metric == "l2":
         sum_sq = 0.0

paimon-python/pypaimon/tests/primary_key_global_index_golden_test.py

@@ -148,3 +148,50 @@ def test_java_primary_key_raw_only_uses_column_metric(catalog, metric, expected_
     assert rows.column("id").to_pylist() == [expected_id]
     assert result.positions[0].score == _compute_score(
         query, rows.column("embedding")[0].as_py(), metric)
+
+
+def test_java_primary_key_vector_refinement_matches_scalar(catalog):
+    from unittest import mock
+    from pypaimon.table.source import vector_search_read as scoring
+
+    _require_native("paimon_vindex")
+    table = catalog.get_table("default.test_pk_vector_golden")
+
+    def search():
+        return (table.new_vector_search_builder()
+                .with_vector_column("embedding")
+                .with_query_vector([1.0, 0.0, 0.0, 0.0])
+                .with_option("ivf.refine_factor", "2")
+                .with_limit(2).execute_local())
+
+    with mock.patch.object(scoring, "_compute_scores", lambda *args: None):
+        expected = search()
+    with mock.patch.object(scoring, "_compute_scores", wraps=scoring._compute_scores) as fast:
+        actual = search()
+    assert actual.positions == expected.positions
+    assert actual.positions
+    assert fast.call_count > 0
+
+
+def test_java_primary_key_raw_vectors_match_scalar(catalog):
+    from dataclasses import replace
+    from unittest import mock
+    from pypaimon.table.source import vector_search_read as scoring
+    from pypaimon.table.source.primary_key_vector_scan import PrimaryKeyVectorScanPlan
+
+    table = catalog.get_table("default.test_pk_vector_golden")
+    builder = (table.new_vector_search_builder()
+               .with_vector_column("embedding")
+               .with_query_vector([1.0, 0.0, 0.0, 0.0]).with_limit(2))
+    plan = builder.new_vector_search_scan().scan()
+    raw_plan = PrimaryKeyVectorScanPlan(plan.snapshot_id, [
+        replace(split, payloads=(), uncovered_data_files=tuple(
+            file.file_name for file in split.data_split.files)) for split in plan.splits()])
+    reader = builder.new_vector_search_read()
+    with mock.patch.object(scoring, "_compute_scores", lambda *args: None):
+        expected = list(reader._raw_candidates(raw_plan))
+    with mock.patch.object(scoring, "_compute_scores", wraps=scoring._compute_scores) as fast:
+        actual = list(reader._raw_candidates(raw_plan))
+    assert actual == expected
+    assert actual
+    assert fast.call_count > 0