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 dca0ab89db [python] Keep batch vector scoring vectorized for null rows
(#9934)
dca0ab89db is described below
commit dca0ab89db4bb04c8934cc793b911de26602c72d
Author: chaoyang <[email protected]>
AuthorDate: Fri Sep 18 12:00:56 2026 +0800
[python] Keep batch vector scoring vectorized for null rows (#9934)
---
.../pypaimon/table/source/vector_search_read.py | 14 +++++-
.../pypaimon/tests/vector_scoring_test.py | 50 ++++++++++++++++++++++
2 files changed, 63 insertions(+), 1 deletion(-)
diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py
b/paimon-python/pypaimon/table/source/vector_search_read.py
index c473d36308..178a7ea907 100644
--- a/paimon-python/pypaimon/table/source/vector_search_read.py
+++ b/paimon-python/pypaimon/table/source/vector_search_read.py
@@ -1014,7 +1014,14 @@ def _iter_arrow_batch_scores(vectors, queries, metric):
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
+ positions = None
+ score_vectors = block
+ if fast_queries and block.null_count:
+ # Compact only this block; null rows must not disable vectorized
+ # scoring for the remaining rows or change their result positions.
+ positions = [i for i, valid in
enumerate(block.is_valid().to_pylist()) if valid]
+ score_vectors = block.take(np.asarray(positions, dtype=np.int64))
+ matrix = _arrow_score_matrix(score_vectors, dimension) if fast_queries
else None
if matrix is not None and not np.isfinite(matrix).all():
matrix = None
if matrix is None:
@@ -1030,6 +1037,11 @@ def _iter_arrow_batch_scores(vectors, queries, metric):
working = matrix if query_index == last_query else matrix.copy()
scores = _compute_scores(query, working, metric)
assert scores is not None
+ if positions is not None:
+ restored = [None] * len(block)
+ for position, score in zip(positions, scores):
+ restored[position] = score
+ scores = restored
yield start, query_index, scores
diff --git a/paimon-python/pypaimon/tests/vector_scoring_test.py
b/paimon-python/pypaimon/tests/vector_scoring_test.py
index b9ef51e0a3..5ad0c1fd7e 100644
--- a/paimon-python/pypaimon/tests/vector_scoring_test.py
+++ b/paimon-python/pypaimon/tests/vector_scoring_test.py
@@ -110,6 +110,56 @@ class VectorScoringTest(unittest.TestCase):
self.assertEqual(1.0, scores[1])
self.assertEqual([1.0], list(_iter_arrow_scores(pa.array([[]],
type=pa.list_(pa.float32())), [], "l2")))
+ def test_batch_null_vectors_preserve_scores_without_scalar_fallback(self):
+ rng = np.random.default_rng(37)
+ values = rng.standard_normal((2053, 7)).astype(np.float32).tolist()
+ for i in (0, 2, 511, 1025, 2051, 2052):
+ values[i] = None
+ queries = [[1.0] * 7, [0.0] * 7, [-0.5] * 7]
+ for dtype in (pa.list_(pa.float32()), pa.large_list(pa.float32()),
pa.list_(pa.float32(), 7)):
+ sliced = pa.array(values, type=dtype).slice(2, 2050)
+ chunked = pa.chunked_array([sliced.slice(0, 513),
sliced.slice(513)])
+ for vectors in (sliced, chunked):
+ for metric in ("l2", "cosine", "inner_product"):
+ with self.subTest(dtype=dtype, chunked=isinstance(vectors,
pa.ChunkedArray), metric=metric):
+ expected = [[None if row is None else
_compute_score(query, row, metric)
+ for row in values[2:2052]] for query in
queries]
+ actual = [[] for _ in queries]
+ with
mock.patch("pypaimon.table.source.vector_search_read._compute_score",
+ side_effect=AssertionError("valid rows
must use vectorized scoring")):
+ 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)
+ self.assertEqual(expected, actual)
+
+ def
test_batch_all_null_vectors_and_invalid_children_keep_existing_behavior(self):
+ queries = [[1.0, 0.0], [0.0, 0.0]]
+ for dtype in (pa.list_(pa.float32()), pa.large_list(pa.float32()),
pa.list_(pa.float32(), 2)):
+ for metric in ("l2", "cosine", "inner_product"):
+ result = list(_iter_arrow_batch_scores(pa.array([None, None],
type=dtype), queries, metric))
+ self.assertEqual([(0, 0, [None, None]), (0, 1, [None, None])],
result)
+ for rows, exception, message in (
+ ([None, [1.0]], ValueError, "dimension mismatch"),
+ ([None, [1.0, None]], TypeError, "NoneType")):
+ with self.subTest(rows=rows):
+ vectors = pa.array(rows, type=pa.list_(pa.float32()))
+ with self.assertRaisesRegex(exception, message):
+ list(_iter_arrow_batch_scores(vectors, queries, "l2"))
+ vectors = pa.array([None, [float("nan"), 0.0], [1.0, 0.0]],
type=pa.list_(pa.float32()))
+ scores = list(_iter_arrow_batch_scores(vectors, queries, "l2"))[0][2]
+ self.assertIsNone(scores[0])
+ self.assertTrue(np.isnan(scores[1]))
+ self.assertEqual(1.0, scores[2])
+
+ def test_batch_null_parent_ignores_invalid_child_values(self):
+ vectors = pa.Array.from_buffers(
+ pa.list_(pa.float32(), 2), 2, [pa.py_buffer(b"\x02")],
+ children=[pa.array([float("nan"), None, 1.0, 2.0],
type=pa.float32())])
+ with
mock.patch("pypaimon.table.source.vector_search_read._compute_score",
+ side_effect=AssertionError("null children must not
disable the fast path")):
+ result = list(_iter_arrow_batch_scores(vectors, [[1.0, 2.0]],
"l2"))
+ self.assertEqual([(0, 0, [None, 1.0])], result)
+
def test_raw_search_filters_before_scoring_and_preserves_ties(self):
column = _field(1, "embedding", "FLOAT")
reader = DataEvolutionVectorRead(_StubTable([column], []), 2, column,
[1, 0])
