alibaba--zvec/python/tests/test_jieba_default_dict.py

# Copyright 2025-present the zvec project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""End-to-end: jieba FTS works without any user configuration.

`import zvec` is supposed to register the wheel-bundled jieba dict
directory via `set_default_jieba_dict_dir`. With that in place a user can
declare an FTS field with `tokenizer_name="jieba"`, leave `extra_params`
empty, and Chinese full-text search just works.

Falls back to GTEST_SKIP-equivalent when running against a build that did
not bundle the dict (e.g., source-tree dev install without the install
step). In that case CI will rely on the C++ unit tests instead.
"""

from __future__ import annotations

import os
import sys

import pytest
import zvec
from zvec import (
    Collection,
    CollectionOption,
    DataType,
    Doc,
    FieldSchema,
    FtsIndexParam,
)
from zvec.model.param.query import Fts, Query


def _bundled_dict_dir() -> str:
    """Path zvec.__init__ would have registered; empty when not bundled."""
    return zvec.get_default_jieba_dict_dir()


def _bundled_dict_files_exist() -> bool:
    """Whether the registered default actually contains the dict files.

    `importlib.resources` happily returns a path even when the data dir was
    not installed (e.g. source-tree dev runs); only an installed wheel has
    the files on disk.
    """
    import os

    base = _bundled_dict_dir()
    if not base:
        return False
    return os.path.isfile(os.path.join(base, "jieba.dict.utf8")) and os.path.isfile(
        os.path.join(base, "hmm_model.utf8")
    )


@pytest.fixture(scope="module", autouse=True)
def _require_bundled_dict():
    if not _bundled_dict_files_exist():
        pytest.skip(
            "Bundled jieba dict not found at zvec/data/jieba_dict/ — "
            "this test requires an installed wheel (not a source-tree dev "
            "build without the install step).",
        )


@pytest.fixture(scope="function")
def jieba_collection(tmp_path_factory) -> Collection:
    """FTS-only collection using jieba tokenizer and no explicit dict path."""
    # env-var shadows GlobalConfig in the priority chain.
    if os.environ.get("ZVEC_JIEBA_DICT_DIR"):
        pytest.skip("ZVEC_JIEBA_DICT_DIR shadows the bundled default")
    temp_dir = tmp_path_factory.mktemp("zvec_jieba_default")
    collection_path = temp_dir / "fts_jieba"

    schema = zvec.CollectionSchema(
        name="fts_jieba_default",
        fields=[
            FieldSchema("title", DataType.STRING, nullable=False),
            FieldSchema(
                "content",
                DataType.STRING,
                nullable=False,
                # Deliberately omit extra_params — the bundled default must
                # be picked up via GlobalConfig.jieba_dict_dir.
                index_param=FtsIndexParam(
                    tokenizer_name="jieba",
                    filters=["lowercase"],
                ),
            ),
        ],
    )

    coll = zvec.create_and_open(
        path=str(collection_path),
        schema=schema,
        option=CollectionOption(read_only=False, enable_mmap=True),
    )
    assert coll is not None
    try:
        yield coll
    finally:
        try:
            coll.destroy()
        except Exception as e:
            print(f"Warning: failed to destroy collection: {e}")


def test_jieba_works_without_explicit_dict_path(jieba_collection: Collection):
    """User opens collection, inserts CJK doc, searches — no init() / no
    extra_params / no env var / no manual setter call. Just `import zvec`."""
    docs = [
        Doc(id="pk_1", fields={"title": "t1", "content": "中华人民共和国成立"}),
        Doc(id="pk_2", fields={"title": "t2", "content": "无关文档"}),
    ]
    insert_results = jieba_collection.insert(docs)
    assert all(r.ok() for r in insert_results)

    hits = jieba_collection.query(
        queries=Query(field_name="content", fts=Fts(match_string="中华")),
        topk=10,
    )
    ids = {doc.id for doc in hits}
    assert "pk_1" in ids
    assert "pk_2" not in ids


def test_default_dict_dir_is_registered_on_import():
    """Sanity check: zvec.__init__ registered a non-empty default."""
    assert _bundled_dict_dir() != ""


def test_user_can_override_default_at_runtime():
    """zvec.set_default_jieba_dict_dir can be called any time to override."""
    saved = zvec.get_default_jieba_dict_dir()
    try:
        zvec.set_default_jieba_dict_dir("/tmp/zvec/jieba-override")
        assert zvec.get_default_jieba_dict_dir() == "/tmp/zvec/jieba-override"
    finally:
        zvec.set_default_jieba_dict_dir(saved)


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="os.environ writes may not propagate across CRT to zvec.pyd",
)
def test_env_var_overrides_global_config(monkeypatch, tmp_path_factory):
    """ZVEC_JIEBA_DICT_DIR beats GlobalConfig in jieba's resolution chain."""
    bundled = _bundled_dict_dir()
    monkeypatch.setenv("ZVEC_JIEBA_DICT_DIR", bundled)
    saved_global = zvec.get_default_jieba_dict_dir()
    try:
        zvec.set_default_jieba_dict_dir("/zvec/intentionally/missing/global")

        temp_dir = tmp_path_factory.mktemp("zvec_jieba_env")
        schema = zvec.CollectionSchema(
            name="fts_jieba_env",
            fields=[
                FieldSchema("title", DataType.STRING, nullable=False),
                FieldSchema(
                    "content",
                    DataType.STRING,
                    nullable=False,
                    index_param=FtsIndexParam(
                        tokenizer_name="jieba",
                        filters=["lowercase"],
                    ),
                ),
            ],
        )
        coll = zvec.create_and_open(
            path=str(temp_dir / "fts_jieba_env"),
            schema=schema,
            option=CollectionOption(read_only=False, enable_mmap=True),
        )
        assert coll is not None
        try:
            results = coll.insert(
                [
                    Doc(id="pk_1", fields={"title": "t", "content": "搜索引擎技术"}),
                ]
            )
            assert all(r.ok() for r in results)
            hits = coll.query(
                queries=Query(field_name="content", fts=Fts(match_string="搜索")),
                topk=10,
            )
            assert {d.id for d in hits} == {"pk_1"}
        finally:
            coll.destroy()
    finally:
        zvec.set_default_jieba_dict_dir(saved_global)