Join a description table into image similarity results

README.md

@@ -315,6 +315,28 @@ python -m multimodal_toolkit.image.workflow.query \
   --lance-uri s3://contacts/image_poc/assets.lance \
   --image-from face_001.jpg
 
+# Join a description table (doc_id → description) into similarity results.
+# The table can be a plain parquet/jsonl/csv file — no ingestion needed:
+#
+#   import pyarrow as pa, pyarrow.parquet as pq
+#   pq.write_table(pa.table({
+#       "doc_id": ["face_001.jpg", "group_photo.jpg"],
+#       "description": ["清晰正面人像", "两人合影"],
+#   }), "descriptions.parquet")
+#
+# Results gain a `description` column (left join: images without a
+# description stay in the results with description = null).
+python -m multimodal_toolkit.image.workflow.query \
+  --lance-uri s3://contacts/image_poc/assets.lance \
+  --text "合影" \
+  --desc-table descriptions.parquet
+
+# With --sql the description table is registered as `descriptions`:
+python -m multimodal_toolkit.image.workflow.query \
+  --lance-uri s3://contacts/image_poc/assets.lance \
+  --sql "SELECT i.doc_id, d.description FROM images i LEFT JOIN descriptions d ON i.doc_id = d.doc_id WHERE i.has_face = true" \
+  --desc-table descriptions.parquet
+
 # Stage 5 — manage (shared entry point)
 python -m multimodal_toolkit.workflow.manage \
   --lance-uri s3://contacts/image_poc/assets.lance --before 2025-01-01

multimodal_toolkit/image/workflow/query.py

@@ -5,6 +5,9 @@
   --text       文本搜图（ChineseCLIP 文本向量 → image_embedding ANN）
   --image-path 本地图片搜相似图
   --image-from 表内 doc_id 搜相似图
+  --desc-table 描述表（doc_id → description，parquet/jsonl/csv/lance 均可）：
+               查询结果按 doc_id left join 出 description 列；--sql 模式下
+               该表以 ``descriptions`` 为名注册进 SQL 作用域，可手写 JOIN
 """
 from __future__ import annotations
 
@@ -41,7 +44,44 @@ def _doc_id_filter(doc_id: str) -> str:
     return f"doc_id = '{escaped}'"
 
 
-def scalar_query(lance_uri: str, where: str | None = None, top_k: int = 100) -> list[dict]:
+def _read_desc_table(uri: str):
+    """读取描述表（doc_id → description），按后缀选 reader。
+
+    描述表可以是任意 Daft 能读的表：parquet/jsonl/csv 文件即可当表用，
+    不必先落成 Lance；无后缀匹配时按 Lance 表读。
+    """
+    import daft
+
+    io_config = daft_io_config()
+    low = uri.rstrip("/").lower()
+    if low.endswith(".parquet"):
+        df = daft.read_parquet(uri, io_config=io_config)
+    elif low.endswith(".jsonl") or low.endswith(".ndjson") or low.endswith(".json"):
+        df = daft.read_json(uri, io_config=io_config)
+    elif low.endswith(".csv"):
+        df = daft.read_csv(uri, io_config=io_config)
+    else:
+        df = daft.read_lance(uri, io_config=io_config)
+    return df.select("doc_id", "description")
+
+
+def _maybe_join_description(df, cols: list[str], desc_table: str | None):
+    """有描述表时按 doc_id left join，返回 (df, cols)。
+
+    left join 保证没有描述的图片不会从结果里消失（description 为 null）。
+    """
+    if not desc_table:
+        return df, cols
+    desc = _read_desc_table(desc_table)
+    return df.join(desc, on="doc_id", how="left"), [*cols, "description"]
+
+
+def scalar_query(
+    lance_uri: str,
+    where: str | None = None,
+    top_k: int = 100,
+    desc_table: str | None = None,
+) -> list[dict]:
     """标量过滤查询（过滤条件经 Daft 下推到 Lance scanner，不全表扫描）。"""
     import daft
 
@@ -51,28 +91,39 @@ def scalar_query(lance_uri: str, where: str | None = None, top_k: int = 100) ->
     df = daft.read_lance(lance_uri, io_config=daft_io_config(), **kwargs)
     names = set(df.schema().column_names())
     cols = [c for c in DEFAULT_COLUMNS if c in names]
+    df, cols = _maybe_join_description(df.select(*cols), cols, desc_table)
     rows = df.select(*cols).limit(top_k).collect().to_pydict()
     return _rows_from_pydict(rows)
 
 
-def sql_query(lance_uri: str, sql: str, top_k: int = 100) -> list[dict]:
+def sql_query(
+    lance_uri: str,
+    sql: str,
+    top_k: int = 100,
+    desc_table: str | None = None,
+) -> list[dict]:
     """对图片表执行任意 Daft SQL SELECT（表在 SQL 里叫 ``images``）。
 
+    传入 desc_table 时，描述表以 ``descriptions`` 为名一并注册进 SQL 作用域。
+
     示例::
 
         SELECT doc_id, blur_score, face_count
         FROM images
         WHERE has_face = true AND is_blurry = false
         ORDER BY blur_score ASC
 
-        SELECT has_face, COUNT(*) AS cnt, AVG(blur_score) AS avg_blur
-        FROM images
-        GROUP BY has_face
+        SELECT i.doc_id, d.description
+        FROM images i LEFT JOIN descriptions d ON i.doc_id = d.doc_id
+        WHERE i.has_face = true
     """
     import daft
 
     images = daft.read_lance(lance_uri, io_config=daft_io_config())
-    rows = daft.sql(sql, images=images).limit(top_k).collect().to_pydict()
+    tables: dict = {"images": images}
+    if desc_table:
+        tables["descriptions"] = _read_desc_table(desc_table)
+    rows = daft.sql(sql, **tables).limit(top_k).collect().to_pydict()
     return _rows_from_pydict(rows)
 
 
@@ -82,6 +133,7 @@ def _vector_query(
     top_k: int = 10,
     where: str | None = None,
     distance_range: tuple[float, float] | None = None,
+    desc_table: str | None = None,
 ) -> list[dict]:
     import daft
     import pyarrow as pa
@@ -102,6 +154,9 @@ def _vector_query(
     df = daft.read_lance(lance_uri, io_config=daft_io_config(), default_scan_options=scan_options)
     names = set(df.schema().column_names())
     cols = [c for c in DEFAULT_COLUMNS if c in names]
+    # 先 select 收窄列再 join：ANN 结果只有 top_k 行，join 成本可忽略；
+    # 最后再 select 一次固定列序（join 可能改变列顺序）。
+    df, cols = _maybe_join_description(df.select(*cols), cols, desc_table)
     rows = df.select(*cols).limit(top_k).collect().to_pydict()
     return _rows_from_pydict(rows)
 
@@ -112,13 +167,14 @@ def text_query(
     top_k: int = 10,
     where: str | None = None,
     distance_range: tuple[float, float] | None = None,
+    desc_table: str | None = None,
 ) -> list[dict]:
     from multimodal_toolkit.image.embedding import get_embedder
 
     q_vec = get_embedder().embed_text(text)
     if q_vec is None:
         raise ValueError("text query is empty")
-    return _vector_query(lance_uri, q_vec, top_k, where, distance_range)
+    return _vector_query(lance_uri, q_vec, top_k, where, distance_range, desc_table)
 
 
 def image_path_query(
@@ -127,14 +183,15 @@ def image_path_query(
     top_k: int = 10,
     where: str | None = None,
     distance_range: tuple[float, float] | None = None,
+    desc_table: str | None = None,
 ) -> list[dict]:
     from multimodal_toolkit.image.embedding import get_embedder
 
     with open(image_path, "rb") as fp:
         q_vec = get_embedder().embed_image_bytes(fp.read())
     if q_vec is None:
         raise ValueError(f"image query cannot be embedded: {image_path}")
-    return _vector_query(lance_uri, q_vec, top_k, where, distance_range)
+    return _vector_query(lance_uri, q_vec, top_k, where, distance_range, desc_table)
 
 
 def image_doc_query(
@@ -143,6 +200,7 @@ def image_doc_query(
     top_k: int = 10,
     where: str | None = None,
     distance_range: tuple[float, float] | None = None,
+    desc_table: str | None = None,
 ) -> list[dict]:
     import daft
 
@@ -159,7 +217,7 @@ def image_doc_query(
     )
     if not query_rows.get("image_embedding") or query_rows["image_embedding"][0] is None:
         raise ValueError(f"query_doc_id not found or has null image_embedding: {query_doc_id}")
-    return _vector_query(lance_uri, query_rows["image_embedding"][0], top_k, where, distance_range)
+    return _vector_query(lance_uri, query_rows["image_embedding"][0], top_k, where, distance_range, desc_table)
 
 
 def run(
@@ -171,17 +229,18 @@ def run(
     image_path: str | None = None,
     image_from: str | None = None,
     distance_range: tuple[float, float] | None = None,
+    desc_table: str | None = None,
 ) -> None:
     if sql:
-        results = sql_query(lance_uri, sql, top_k)
+        results = sql_query(lance_uri, sql, top_k, desc_table)
     elif text:
-        results = text_query(lance_uri, text, top_k, where, distance_range)
+        results = text_query(lance_uri, text, top_k, where, distance_range, desc_table)
     elif image_path:
-        results = image_path_query(lance_uri, image_path, top_k, where, distance_range)
+        results = image_path_query(lance_uri, image_path, top_k, where, distance_range, desc_table)
     elif image_from:
-        results = image_doc_query(lance_uri, image_from, top_k, where, distance_range)
+        results = image_doc_query(lance_uri, image_from, top_k, where, distance_range, desc_table)
     else:
-        results = scalar_query(lance_uri, where, top_k)
+        results = scalar_query(lance_uri, where, top_k, desc_table)
     for row in results:
         print(row)
 
@@ -197,6 +256,11 @@ def main() -> None:
     parser.add_argument("--image-from", help="doc_id in the image table to use as a similarity query")
     parser.add_argument("--distance-min", type=float, help="minimum vector distance for ANN results")
     parser.add_argument("--distance-max", type=float, help="maximum vector distance for ANN results")
+    parser.add_argument(
+        "--desc-table",
+        help="description table (doc_id, description) as parquet/jsonl/csv/lance; "
+        "left-joined into results, and registered as `descriptions` for --sql",
+    )
     args = parser.parse_args()
     vector_modes = [bool(args.text), bool(args.image_path), bool(args.image_from)]
     if sum(vector_modes) > 1:
@@ -219,6 +283,7 @@ def main() -> None:
         args.image_path,
         args.image_from,
         distance_range,
+        args.desc_table,
     )
 
 

tests/image/test_desc_join.py

@@ -0,0 +1,115 @@
+"""场景测试：以文/以图搜图时，按 doc_id left join 描述表输出 description。
+
+描述表在代码中生成为 parquet 文件（doc_id → description），覆盖三种情况：
+正常映射、图片无描述（left join 后为 null）、描述表中多余的 doc_id（被忽略）。
+"""
+from __future__ import annotations
+
+import pathlib
+import tempfile
+
+import lance
+import pyarrow as pa
+import pyarrow.parquet as pq
+import pytest
+
+from multimodal_toolkit.image.workflow.query import (
+    image_doc_query,
+    scalar_query,
+    sql_query,
+    text_query,
+)
+
+DIM = 4
+
+ROWS = {
+    "doc_id": ["img_a", "img_b", "img_c"],
+    "status": ["ok", "ok", "ok"],
+    "has_face": [True, False, True],
+    "blur_score": [500.0, 50.0, 300.0],
+    "is_blurry": [False, True, False],
+}
+
+EMBEDDINGS = [
+    [1.0, 0.0, 0.0, 0.0],
+    [0.0, 1.0, 0.0, 0.0],
+    [0.0, 0.9, 0.1, 0.0],
+]
+
+# img_b 故意没有描述；img_zzz 是表里不存在的图片。
+DESCRIPTIONS = {
+    "doc_id": ["img_a", "img_c", "img_zzz"],
+    "description": ["清晰正面人像", "两人合影", "不存在的图"],
+}
+
+
+@pytest.fixture(scope="module")
+def scenario() -> dict:
+    tmp = pathlib.Path(tempfile.mkdtemp())
+    lance_uri = str(tmp / "images.lance")
+    tbl = pa.table(
+        {
+            **ROWS,
+            "image_embedding": pa.array(EMBEDDINGS, type=pa.list_(pa.float32(), DIM)),
+        }
+    )
+    lance.write_dataset(tbl, lance_uri)
+
+    # 描述表：代码中生成 parquet，文件本身即可当表参与 join
+    desc_path = str(tmp / "descriptions.parquet")
+    pq.write_table(pa.table(DESCRIPTIONS), desc_path)
+
+    return {"lance_uri": lance_uri, "desc_table": desc_path}
+
+
+class _FakeImageEmbedder:
+    def embed_text(self, text):
+        return [0.0, 1.0, 0.0, 0.0] if text else None
+
+    def embed_image_bytes(self, image_bytes):
+        return [1.0, 0.0, 0.0, 0.0] if image_bytes else None
+
+
+def test_text_query_joins_description(monkeypatch, scenario):
+    import multimodal_toolkit.image.embedding as embedding
+
+    monkeypatch.setattr(embedding, "get_embedder", lambda: _FakeImageEmbedder())
+    rows = text_query(scenario["lance_uri"], "合影", top_k=3, desc_table=scenario["desc_table"])
+    by_id = {r["doc_id"]: r for r in rows}
+    assert by_id["img_c"]["description"] == "两人合影"
+    assert by_id["img_a"]["description"] == "清晰正面人像"
+    # left join：没有描述的图片不丢，description 为 null
+    assert by_id["img_b"]["description"] is None
+
+
+def test_image_doc_query_joins_description(scenario):
+    rows = image_doc_query(scenario["lance_uri"], "img_a", top_k=1, desc_table=scenario["desc_table"])
+    assert rows[0]["doc_id"] == "img_a"
+    assert rows[0]["description"] == "清晰正面人像"
+
+
+def test_extra_doc_id_in_desc_table_is_ignored(scenario):
+    rows = scalar_query(scenario["lance_uri"], top_k=10, desc_table=scenario["desc_table"])
+    assert sorted(r["doc_id"] for r in rows) == ["img_a", "img_b", "img_c"]  # img_zzz 不出现
+
+
+def test_sql_query_manual_join(scenario):
+    rows = sql_query(
+        scenario["lance_uri"],
+        "SELECT i.doc_id, d.description FROM images i "
+        "LEFT JOIN descriptions d ON i.doc_id = d.doc_id "
+        "WHERE i.has_face = true ORDER BY i.doc_id",
+        desc_table=scenario["desc_table"],
+    )
+    assert [(r["doc_id"], r["description"]) for r in rows] == [
+        ("img_a", "清晰正面人像"),
+        ("img_c", "两人合影"),
+    ]
+
+
+def test_no_desc_table_keeps_existing_behavior(monkeypatch, scenario):
+    import multimodal_toolkit.image.embedding as embedding
+
+    monkeypatch.setattr(embedding, "get_embedder", lambda: _FakeImageEmbedder())
+    rows = text_query(scenario["lance_uri"], "合影", top_k=3)
+    assert all("description" not in r for r in rows)