[python] Infer multimodal vector columns by query dimension

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"):

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",