diff --git a/paimon-python/pypaimon/multimodal/table.py b/paimon-python/pypaimon/multimodal/table.py index 40674aeb07ef..d7b8706b5139 100644 --- a/paimon-python/pypaimon/multimodal/table.py +++ b/paimon-python/pypaimon/multimodal/table.py @@ -416,7 +416,8 @@ def search( return VectorQuery( read_table, vector=vector, - vector_column=column or _infer_vector_column(schema, "column"), + vector_column=_resolve_vector_column( + schema, column, len(vector)), vector_options=options, pre_filter=pre_filter, ) @@ -434,7 +435,13 @@ def search_vectors( self.raw_table, snapshot_id=snapshot_id, tag_name=tag_name) schema = _target_schema(read_table) vectors = _coerce_vectors(vectors) - vector_column = column or _infer_vector_column(schema, "column") + dimension = len(vectors[0]) + if any(len(vector) != dimension for vector in vectors): + raise ValueError( + "search_vectors requires all query vectors to have the same " + "dimension.") + vector_column = _resolve_vector_column( + schema, column, dimension) return BatchVectorQuery( read_table, vectors=vectors, @@ -1135,13 +1142,47 @@ def _coerce_full_text_query(query, method, schema, column=None): raise ValueError("%s requires a text string or query mapping." % method) -def _infer_vector_column(schema: pa.Schema, parameter: str = "vector_column"): - columns = [ - field.name +def _resolve_vector_column( + schema: pa.Schema, + column: Optional[str], + dimension: int) -> str: + if column is not None: + try: + field = schema.field(column) + except KeyError as e: + raise ValueError( + "Vector column '%s' not found in table schema." % column) from e + if not pa.types.is_fixed_size_list(field.type): + raise ValueError( + "Column '%s' is not a fixed-size vector column." % column) + if field.type.list_size != dimension: + raise ValueError( + "Vector dimension %d does not match column '%s' dimension %d." + % (dimension, column, field.type.list_size)) + return column + + candidates = [ + (field.name, field.type.list_size) for field in schema if pa.types.is_fixed_size_list(field.type) ] - return _infer_single_column(columns, "vector", parameter) + matches = [ + name + for name, column_dimension in candidates + if column_dimension == dimension + ] + if len(matches) == 1: + return matches[0] + if not matches: + available = ", ".join( + "%s(%d)" % candidate for candidate in candidates) or "none" + raise ValueError( + "No vector column found with dimension %d; available vector " + "columns: %s." + % (dimension, available)) + raise ValueError( + "Multiple vector columns found with dimension %d: %s; pass column." + % (dimension, ", ".join(matches))) def _infer_text_column(schema: pa.Schema, parameter: str = "text_column"): diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py b/paimon-python/pypaimon/tests/multimodal_table_test.py index ddf84f5696e1..e99fe194371b 100644 --- a/paimon-python/pypaimon/tests/multimodal_table_test.py +++ b/paimon-python/pypaimon/tests/multimodal_table_test.py @@ -2137,6 +2137,98 @@ def execute_local(self): self.assertEqual("category", calls["pre_filter"].field) self.assertEqual(["lake"], calls["pre_filter"].literals) + def test_searches_infer_vector_column_by_query_dimension(self): + docs = self.conn.create_table( + "docs", + schema=_schema({ + "id": pa.int32(), + "image_embedding": _vector(2), + "text_embedding": _vector(3), + }), + options=_PARQUET_OPTIONS, + ) + docs.add([ + { + "id": 1, + "image_embedding": [1.0, 0.0], + "text_embedding": [1.0, 0.0, 0.0], + }, + { + "id": 2, + "image_embedding": [0.0, 1.0], + "text_embedding": [0.0, 1.0, 0.0], + }, + ]) + + result = ( + docs.search([0.0, 1.0, 0.0]) + .select(["id"]) + .limit(1) + .to_list() + ) + batch_result = ( + docs.search_vectors([[0.0, 1.0, 0.0]]) + .select(["id"]) + .limit(1) + .to_list() + ) + + self.assertEqual([{"id": 2}], result) + self.assertEqual([[{"id": 2}]], batch_result) + + def test_search_reports_vector_column_dimension_errors(self): + docs = self.conn.create_table( + "docs", + schema=_schema({ + "id": pa.int32(), + "content": pa.string(), + "image_embedding": _vector(2), + "text_embedding": _vector(3), + }), + options=_PARQUET_OPTIONS, + ) + + with self.assertRaisesRegex( + ValueError, + r"No vector column found with dimension 4;.*" + r"image_embedding\(2\).*text_embedding\(3\)"): + docs.search([1.0, 0.0, 0.0, 0.0]) + + with self.assertRaisesRegex(ValueError, "not found in table schema"): + docs.search([1.0, 0.0], column="missing") + + with self.assertRaisesRegex(ValueError, "not a fixed-size vector"): + docs.search([1.0, 0.0], column="content") + + with self.assertRaisesRegex( + ValueError, + "Vector dimension 3 does not match column " + "'image_embedding' dimension 2"): + docs.search([1.0, 0.0, 0.0], column="image_embedding") + + with self.assertRaisesRegex(ValueError, "same dimension"): + docs.search_vectors([ + [1.0, 0.0, 0.0], + [1.0, 0.0], + ]) + + def test_search_requires_column_for_same_dimension_vectors(self): + docs = self.conn.create_table( + "docs", + schema=_schema({ + "id": pa.int32(), + "title_embedding": _vector(2), + "body_embedding": _vector(2), + }), + options=_PARQUET_OPTIONS, + ) + + with self.assertRaisesRegex( + ValueError, + r"Multiple vector columns found with dimension 2: " + r"title_embedding, body_embedding; pass column\."): + docs.search([1.0, 0.0]) + def test_search_pre_filter_rejects_predicate_object(self): docs = self.conn.create_table( "docs",