This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit b9e8d12f6d60526a13e7d1e803324ef8f13d25b9 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])
