diff --git a/README.md b/README.md index fb6acd6..9ae820e 100644 --- a/README.md +++ b/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 diff --git a/multimodal_toolkit/image/workflow/query.py b/multimodal_toolkit/image/workflow/query.py index c72bac0..e60e6fa 100644 --- a/multimodal_toolkit/image/workflow/query.py +++ b/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,13 +91,21 @@ 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 @@ -65,14 +113,17 @@ def sql_query(lance_uri: str, sql: str, top_k: int = 100) -> list[dict]: 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,6 +183,7 @@ 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 @@ -134,7 +191,7 @@ def image_path_query( 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, ) diff --git a/tests/image/test_desc_join.py b/tests/image/test_desc_join.py new file mode 100644 index 0000000..09ee9ba --- /dev/null +++ b/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)