This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 1032567c4a [python] Vectorize batch raw vector scoring (#9796)
1032567c4a is described below
commit 1032567c4a9c356dc3045c31ebfbdb0397037586
Author: chaoyang <[email protected]>
AuthorDate: Tue Sep 15 11:15:01 2026 +0800
[python] Vectorize batch raw vector scoring (#9796)
---
.../pypaimon/table/source/vector_search_read.py | 113 +++++++++++++++------
.../pypaimon/tests/batch_vector_raw_scan_test.py | 12 +--
.../pypaimon/tests/vector_scoring_test.py | 26 ++++-
3 files changed, 113 insertions(+), 38 deletions(-)
diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py
b/paimon-python/pypaimon/table/source/vector_search_read.py
index b05f6bb8b4..0148e147bc 100644
--- a/paimon-python/pypaimon/table/source/vector_search_read.py
+++ b/paimon-python/pypaimon/table/source/vector_search_read.py
@@ -682,16 +682,15 @@ class
BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl,
with _ClosableArrowBatchReader(reader, batches) as batch_reader:
for batch in batch_reader:
row_ids = batch.column(SpecialFields.ROW_ID.name).to_pylist()
- vectors = batch.column(self._vector_column.name).to_pylist()
- for row_id, stored in zip(row_ids, vectors):
- if stored is None:
- continue
- stored_vector = _to_vector_list(stored)
- for query, heap in zip(self._query_vectors, heaps):
- _check_vector_dimension(query, stored_vector)
- _offer_score(
- heap, self._limit, row_id,
- _compute_score(query, stored_vector, metric))
+ vectors = batch.column(self._vector_column.name)
+ for start, query_index, scores in _iter_arrow_batch_scores(
+ vectors, self._query_vectors, metric
+ ):
+ heap = heaps[query_index]
+ block_row_ids = row_ids[start:start + len(scores)]
+ for row_id, score in zip(block_row_ids, scores):
+ if score is not None:
+ _offer_score(heap, self._limit, row_id, score)
del batch, row_ids, vectors
return [_scored_result(heap) for heap in heaps]
@@ -898,29 +897,11 @@ def _score_block_size(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)
+ matrix = _arrow_score_matrix(block, len(query))
+ scores = None if matrix is None else _compute_scores(query, matrix,
metric)
if scores is None:
for vector in block.to_pylist():
if vector is None:
@@ -933,6 +914,78 @@ def _iter_arrow_scores(vectors, query, metric):
yield from scores
+def _iter_arrow_batch_scores(vectors, queries, metric):
+ """Yield vectorized scores for one Arrow batch without expanding its rows.
+
+ The float64 row matrix is materialized once per bounded block and reused
for
+ every query. Only the scratch matrix for the active query is retained, so
+ memory does not grow with the number of queries.
+ """
+ import numpy as np
+
+ if not queries:
+ return
+ block_size = min(_score_block_size(query) for query in queries)
+ dimension = len(queries[0])
+ fast_queries = metric in ("l2", "cosine", "inner_product")
+ if fast_queries:
+ for query in queries:
+ try:
+ query_array = np.asarray(query, dtype=np.float64)
+ except (ValueError, TypeError, OverflowError):
+ fast_queries = False
+ break
+ if (query_array.ndim != 1 or len(query_array) != dimension or
+ not dimension or not np.isfinite(query_array).all()):
+ fast_queries = False
+ break
+
+ for start in range(0, len(vectors), block_size):
+ block = vectors.slice(start, block_size)
+ matrix = _arrow_score_matrix(block, dimension) if fast_queries else
None
+ if matrix is not None and not np.isfinite(matrix).all():
+ matrix = None
+ if matrix is None:
+ rows = block.to_pylist()
+ for query_index, query in enumerate(queries):
+ yield start, query_index, _score_rows(rows, query, metric)
+ continue
+
+ last_query = len(queries) - 1
+ for query_index, query in enumerate(queries):
+ # Scoring mutates its owned matrix. Reuse the original for the last
+ # query and keep only one scratch copy for every preceding query.
+ working = matrix if query_index == last_query else matrix.copy()
+ scores = _compute_scores(query, working, metric)
+ assert scores is not None
+ yield start, query_index, scores
+
+
+def _arrow_score_matrix(vectors, dimension):
+ """Return one regular float32 Arrow vector block as owned float64 rows."""
+ import numpy as np
+ import pyarrow as pa
+
+ if isinstance(vectors, pa.ChunkedArray):
+ vectors = vectors.combine_chunks()
+ dtype = vectors.type
+ if not ((pa.types.is_list(dtype) or pa.types.is_large_list(dtype) or
+ pa.types.is_fixed_size_list(dtype)) and not vectors.null_count):
+ return None
+ flat = vectors.flatten()
+ if not pa.types.is_float32(flat.type) or flat.null_count:
+ return None
+ if pa.types.is_fixed_size_list(dtype):
+ regular = dtype.list_size == dimension
+ else:
+ offsets = vectors.offsets.to_numpy(zero_copy_only=False)
+ regular = bool(np.all(np.diff(offsets) == dimension))
+ if not regular or not dimension:
+ return None
+ return flat.to_numpy(zero_copy_only=False).astype(np.float64).reshape(
+ len(vectors), dimension)
+
+
def _score_rows(vectors, query, metric):
import numpy as np
diff --git a/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py
b/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py
index ad84a85316..6bc5b4f818 100644
--- a/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py
+++ b/paimon-python/pypaimon/tests/batch_vector_raw_scan_test.py
@@ -126,7 +126,7 @@ class BatchVectorRawScanTest(BatchModeMixin,
DataEvolutionTestBase, unittest.Tes
generators = []
batch_sizes = []
module = 'pypaimon.table.source.vector_search_read'
- from pypaimon.table.source.vector_search_read import _compute_score
+ from pypaimon.table.source.vector_search_read import _compute_scores
def batches(table_read, splits):
arrow, generator = original(table_read, splits)
@@ -135,21 +135,21 @@ class BatchVectorRawScanTest(BatchModeMixin,
DataEvolutionTestBase, unittest.Tes
def tracked_generator(table_read, *args):
source = original_generator(table_read, *args)
- expected_scores = 0
+ expected_score_blocks = 0
try:
for batch in source:
batch_sizes.append(batch.num_rows)
- expected_scores += batch.num_rows * 2
+ expected_score_blocks += 2
yield batch
- self.assertEqual(expected_scores, score.call_count)
+ self.assertEqual(expected_score_blocks, score.call_count)
finally:
source.close()
with patch.object(TableRead, '_new_arrow_batch_reader', batches), \
patch.object(TableRead, '_arrow_batch_generator',
tracked_generator), \
- patch(module + '._compute_score', wraps=_compute_score) as
score:
+ patch(module + '._compute_scores', wraps=_compute_scores) as
score:
reader._read_raw_batch_search([Range(0, 3)], None, 'ivf-flat')
- self.assertEqual(8, score.call_count)
+ self.assertEqual(4, score.call_count)
self.assertEqual(1, len(generators))
self.assertGreater(len(batch_sizes), 1)
self.assertIsNone(generators[0].gi_frame)
diff --git a/paimon-python/pypaimon/tests/vector_scoring_test.py
b/paimon-python/pypaimon/tests/vector_scoring_test.py
index b900efd2a5..b9ef51e0a3 100644
--- a/paimon-python/pypaimon/tests/vector_scoring_test.py
+++ b/paimon-python/pypaimon/tests/vector_scoring_test.py
@@ -22,8 +22,8 @@ import numpy as np
import pyarrow as pa
from pypaimon.table.source.vector_search_read import (
- DataEvolutionVectorRead, _compute_score, _compute_scores,
_iter_arrow_scores,
- _score_block_size, _score_rows,
+ DataEvolutionVectorRead, _compute_score, _compute_scores,
+ _iter_arrow_batch_scores, _iter_arrow_scores, _score_block_size,
_score_rows,
)
from pypaimon.table.special_fields import SpecialFields
from pypaimon.tests.vector_search_filter_test import _StubTable, _field
@@ -70,6 +70,28 @@ class VectorScoringTest(unittest.TestCase):
expected = [_compute_score(query, row, metric) for row in
values[2:1029]]
self.assertEqual(expected, list(_iter_arrow_scores(data,
query, metric)))
+ def test_batch_arrow_scores_exactly_match_scalar_accumulation(self):
+ rng = np.random.default_rng(17)
+ values = rng.standard_normal((2051, 128)).astype(np.float32)
+ vectors = pa.array(values.tolist(), type=pa.list_(pa.float32(), 128))
+ queries = rng.standard_normal((5, 128)).astype(np.float32).tolist()
+ for metric in ("l2", "cosine", "inner_product"):
+ with self.subTest(metric=metric):
+ actual = [[] for unused in queries]
+ for start, query_index, scores in _iter_arrow_batch_scores(
+ vectors, queries, metric
+ ):
+ self.assertEqual(start, len(actual[query_index]))
+ actual[query_index].extend(scores)
+ for query, scores in zip(queries, actual):
+ expected = [
+ _compute_score(query, row, metric)
+ for row in values.tolist()
+ ]
+ self.assertEqual(expected, scores)
+ self.assertEqual(
+ np.array(expected).tobytes(),
np.array(scores).tobytes())
+
def test_null_and_unsupported_data_preserve_scalar_behavior(self):
for dtype in (pa.list_(pa.float32()), pa.list_(pa.float64()),
pa.list_(pa.int64())):
array = pa.array([[1, 2], None, [0, 0]], type=dtype)
