mirror of
https://github.com/milvus-io/milvus.git
synced 2026-10-29 01:16:06 +08:00
## What Consume the producer-owned error classification at the segcore boundary and make the whole C++→Go classification drift-proof, so a segcore error is classified as **input** (caller's fault, non-retriable), **transient** (retriable) or **permanent** (non-retriable) instead of flattening to `UnexpectedError(2001)` or carrying the wrong retry default. Design + tracking: #50903. ## Changes - **T1** — register the storage fallback pair in `pkg/util/merr/segcore.go`: `StorageError(2044)` non-retriable, `StorageTransientError(2045)` retriable. - **T2** — `KnowhereStatusToErrorCode` → a switch with **no `default` + `-Werror=switch`** over the full `knowhere::Status`; add build-path variant `KnowhereBuildStatusToErrorCode` so a build-time OOM / disk read stays **retriable** instead of collapsing into a permanent `IndexBuildError`. - **T3/T4** — `ArrowStatusToErrorCode` delegates to the producer's `milvus_storage::ToSegcoreError` (retires milvus's duplicate mapper); audited and routed **25 storage arrow-status sites** that were collapsing to `2001` through the single mapper (extracted to `storage/StatusToErrorCode.h`), always preserving the arrow sub-code in the message. - **T5** — unmapped-code observability: `UnmappedSegcoreCodeTotal{code}` counter + rate-limited WARN via an observer hook (merr is a leaf package); registered on QueryNode and DataNode. Unknown code degrades to non-retriable, never panics. - **T6** — codegen + compile-time enforcement: a generated `SegcoreCode` type (from milvus-common's `EasyAssert.h`) + an exhaustive `classForCode` switch marked `//exhaustive:enforce`, with the `exhaustive` golangci-lint enabled opt-in — a new C++ code that is not classified fails lint (the C++→Go analog of `-Werror=switch`). - **§3 B-tier** — classify `marisa` and `simdjson` errors (build/load/parse) instead of collapsing to `2001`, sub-code in the message; simdjson optional-access (`NO_SUCH_FIELD`/`INCORRECT_TYPE`) stays a benign skip; the `loon_ffi` FFI boundary is untouched. - **Boundary hardening (adversarial self-review of this PR's own diff)** — closed the escapes that would defeat the mapping above: a `throw e;` slicing rethrow in `LoadWithStrategy` that destroyed the very codes the columnar-read mapping attaches (bare `throw;` now), the same slice in `MinioChunkManager::PreCheck`; `GetCoreMetrics` / `EstimateLoadIndexResource` / init-and-config entry points that could let an exception cross the C ABI and terminate the process; and every remaining extern-C entry that caught only `std::exception` now ends in `catch(...)` via the shared `CGoCatch.h` macros. - **Pin + semantics** — bump `milvus-storage_VERSION` to `11f8a36` (the milvus-io/milvus-storage#574 merge, which also contains #575) and align the no-detail `IOError` expectation with the settled semantics: the producer tags every known-transient failure with a retryable `ExtendStatusDetail`, so a bare `IOError` with no detail is unclassified and deliberately falls back to permanent `StorageError(2044)` — a stripped-detail NotFound now degrades to non-retriable (safe) instead of retriable (retry storm on a permanent 404). - **Wire pass-through (client-visible)** — a segcore error now reaches the client with its ORIGINAL code (2009 stays 2009, 2024 stays 2024) instead of collapsing to the `ErrSegcore(2000)` umbrella with the real code buried in the message. Family identity for `errors.Is` is preserved via inner/Unwrap; input/system/retriable classification unchanged. Guardrails: only in-band (2000-2099) codes pass through (garbage still collapses to 2000); cross-family mappings (2046 → wire 110) keep their sentinel's code. `ErrSegcoreUnsupported`/`ErrSegcorePretendFinished` move to the C++ values they represent (2001→2003, 2002→2033) — their old numbers squatted on C++ UnexpectedError/NotImplemented and would false-match under code-based `errors.Is`. Verified end-to-end on a live standalone (ef<k reaches the client as 2042, unsupported tokenizer as 2001); the three e2e assertions pinning the old 2000 updated. - **Remaining code-destroying sites** — the three classes that still swallowed a producer's classification before the cgo boundary are now gone from `internal/core/src` and `internal/core/thirdparty`: status-consuming `AssertInfo` (104 → 0, incl. ~47 arrow builder paths whose commonest failure is OOM, now retriable `MemAllocateFailed` instead of a permanent 2001), bare `throw std::runtime_error/logic_error/bad_alloc` (68 → 0 — these were not `SegcoreError`, so they collapsed to 2001 *and* falsely fired the untyped-exception observer), and `throw fmt::format(...)` (12 → 0 — it throws a `std::string`, which `catch (std::exception&)` cannot see at all). tantivy's 73 `AssertInfo(res.result_->success, ...)` (plus 10 raw-`RustResult` stragglers found later) now classify the rust error — originally by its Display prefix, since replaced by a proper `#[repr(i32)]` discriminant carried in `RustResult.error_code` (see the Aug-10 update below). Typed `ThrowInfo` sites: 894 → 1081. The ~1500 genuine invariant asserts are untouched — 2001 is correct for them. The long-standing FIXME about `err_code` not surviving the nested LOON FFI boundary is also resolved, delegating to `milvus_storage::ToSegcoreErrorCode` rather than duplicating its table. ## Verification **Verified in this PR:** - **Mapping correctness (unit-tested, in-process):** `test_knowhere_status_mapping.cpp` / `test_storage_error_code.cpp` / `test_exec.cpp` cover every mapper branch (knowhere Status incl. the build variant, arrow/extend status incl. `AwsErrorNotFound→ObjectNotExist(2017)`, permanent-S3 vs transient), plus `FailureCStatus` code preservation and both observer hooks firing. - **Code projection to Go (one hop, unit-tested):** `segcore_test.go` pins `classForCode` for every generated code and asserts `merr.Status(err).GetRetriable()` for transient codes; the T6 generator is idempotent and the `exhaustive` lint fails on an unclassified code. - **Full C++ suite:** 8213/8223 unit tests pass locally (10 skipped; Azure connectivity tests excluded), 8648 in CI, rebased on current master (one pre-existing, unrelated concurrency test excluded: `GrowingConcurrentReopenTest` deadlocks deterministically on current master with or without this PR — rwlock writer starvation in growing-segment reopen code this PR does not touch; reported separately). - **Static audit (grep-verifiable):** every storage arrow-status consumption site on the read path routes through `ArrowStatusToErrorCode`, and every extern-C boundary ends in a `catch(...)` tail. **Explicitly NOT verified here (follow-up):** - **Runtime fault injection.** No S3 throttle / 404 / OOM / corrupt-file failure has been triggered end-to-end in a running cluster. Transient codes reach Go with `retriable=true` (unit-tested projection), but the downstream consumption — `lb_policy` replica reroute on `merr.IsRetryableErr`, index/analyze scheduler retry — is pre-existing logic from #50221 and has **not** been driven by a real segcore transient error in this PR. This PR preserves classification for observability and correct retry defaults; the retry behavior itself is exercised only by its own pre-existing tests. ## Dependencies - ~~milvus-common `StorageTransientError(2045)` — zilliztech/milvus-common#102~~ **merged**. - ~~milvus-storage `ToSegcoreError` / packed `ExtendStatusCode` — milvus-io/milvus-storage#575 + #574~~ **merged; pin bumped in-tree to `11f8a36`**. - ~~knowhere three-way classification — zilliztech/knowhere#1704~~ **merged** (the milvus-side `KnowhereStatusToErrorCode` → thin delegate to knowhere's own `ToSegcoreErrorCode` is a follow-up, gated on a knowhere version bump). - ~~milvus-common untyped-cgo-exception observer — zilliztech/milvus-common#112~~ **merged and released as `1.0.0-1fd1160`; the pin now points at the published package.** All dependencies are in. ## Update (Aug 10) — full-population audit, LOON path, runtime observability The originally deferred FFI/LOON path is now **done on the milvus side**, and the audit was extended from the three grep-able classes to the *entire* 2001-producing population: - **Every remaining 2001 site read.** All 1,517 `AssertInfo` (four sweeps: errno fingerprint, failure-keyword messages, condition morphology, and finally **data provenance** — does the guarded value come from disk/network?) and all 198 explicit `ThrowInfo(UnexpectedError)` sites. ~290 were externally-triggerable and now carry typed codes: file/remote IO -> `FileOpen/Create/Read/WriteFailed` (retriable), mmap/allocation -> `MmapError`/`MemAllocateFailed` (retriable), persisted-format damage (CRC/magic/parquet meta/index-meta keys) -> `DataFormatBroken`, deployment config -> `ConfigInvalid`, request content -> `InvalidParameter`, a cancel-race -> `FollyCancel`. The ~1,400 kept sites are genuine invariants or cgo contracts where 2001 is the correct report. - **Two infinite-retry bugs.** Statically-impossible conditions (index_type x metric blacklist, per-type metric allowlists, json/geometry index gates) threw 2001 -> generic retry -> the build task spun forever; they now throw `Unsupported`, which `getStateFromError` maps to a terminal `JobStateFailed`. Missing `index_type`/`metric_type`/`min_gram`/`max_gram` keys in persisted index meta had the same loop on the load path; they are `DataFormatBroken` now. - **knowhere `expected<>` bypasses closed** (8 sites in `QueryResult.h`/`CachedSearchIterator`): iterator failures went through `AssertInfo` and discarded the Status knowhere had already classified; they now route through `KnowhereStatusToErrorCode`, so an OOM/disk failure during search iteration stays retriable. Preflight rewraps in `segment_c`/`boost_score` similarly preserved the original `SegcoreError` code instead of flattening to 2001+string. - **tantivy discriminant over the FFI.** `RustResult` now carries `error_code` (`#[repr(i32)] TantivyBindingErrorCode`, cbindgen-exported); the C++ mapper switches on the enum instead of parsing the Display text, and the inner `tantivy::TantivyError` is discriminated too (`IoError/Open*Error` -> Io/retriable, `DataCorruption/IncompatibleIndex` -> DataCorruption). Wording changes on the rust side can no longer silently degrade classification. - **LOON / FFI path (the deferred item), milvus side complete.** The Go funnel `HandleLoonFFIResult` dropped `err_code` entirely and wrapped every failure as `ErrLoonTransient` — a 404/access-denied/corrupt-data retried as transient. It now classifies by the producer's own `loon_ffi_is_retryable_errcode`; permanent failures carry the new `ErrLoonPermanent` and terminate retry loops (`pack_writer_v3` via `retry.Unrecoverable`; the external-refresh manager guard extended so behavior does not invert). On the C++ side `LoonErrCodeToErrorCode` is the single classification entry (low band -> hand table, extend band -> producer's `ToSegcoreErrorCode`, unknown -> producer's retryable probe), unifying the two previously-divergent `ThrowIfFFIError` helpers — `LOON_FILE_NOT_FOUND(12)` now converges to `ObjectNotExist(2017)` on both integration paths. Remaining LOON items (e.g. promoting FileNotFound into `ExtendStatusCode`) live in the milvus-storage repo. - **Regression guards.** `scripts/check_segcore_error_boundaries.sh` wired into `make static-check`: every `throw` in `internal/core/src` must carry a milvus ErrorCode (zero-tolerance; currently 0 violations); vendored `fmindex::` is confined to its boundary files; knowhere/arrow/milvus_storage/tantivy are ratcheted by a checked-in file-set baseline (new consumer files fail the check; shrinking is free). - **Runtime observability for what is left.** `milvus_cgo_unexpected_segcore_origin_total{origin="<file>:<line>"}` counts every 2001 crossing the cgo boundary by its C++ source location (parsed from the ` at file:line` suffix `AssertInfo` already emits, build paths collapsed to repo-relative). A site that fires in production names itself — reclassification becomes evidence-driven instead of re-reading ~1,400 asserts. Site count for the 2001 family: 1,955 on master -> 1,525 on this branch; the delta is reclassification into actionable codes, not deletion of checks. ## Deferred - milvus-storage-side LOON improvements: promote `LOON_FILE_NOT_FOUND` into `ExtendStatusCode`, category byte (design §4.7) — tracked in the storage repo. - knowhere-side: thin-delegate `KnowhereStatusToErrorCode` to knowhere's own `ToSegcoreErrorCode`, gated on a knowhere version bump. issue: #50903 --------- Signed-off-by: Zack <noreply@zilliz.com> Co-authored-by: Zack <noreply@zilliz.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: xiaofanluan <xf@hjjaq.com>
3504 lines
143 KiB
Python
3504 lines
143 KiB
Python
import random
|
|
import threading
|
|
import time
|
|
from collections import Counter
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
import utils.util_pymilvus as ut
|
|
from base.client_base import TestcaseBase
|
|
from common import common_func as cf
|
|
from common import common_type as ct
|
|
from common.common_type import CaseLabel, CheckTasks
|
|
from common.text_generator import ICUTextGenerator, KoreanTextGenerator
|
|
from faker import Faker
|
|
from pymilvus import (
|
|
CollectionSchema,
|
|
DataType,
|
|
FieldSchema,
|
|
)
|
|
from utils.util_log import test_log as log
|
|
|
|
Faker.seed(19530)
|
|
|
|
fake_en = Faker("en_US")
|
|
fake_zh = Faker("zh_CN")
|
|
fake_de = Faker("de_DE")
|
|
fake_jp = Faker("ja_JP")
|
|
fake_ko = Faker("ko_KR")
|
|
|
|
|
|
# patch faker to generate text with specific distribution
|
|
cf.patch_faker_text(fake_en, cf.en_vocabularies_distribution)
|
|
cf.patch_faker_text(fake_zh, cf.zh_vocabularies_distribution)
|
|
|
|
pd.set_option("expand_frame_repr", False)
|
|
|
|
prefix = "query"
|
|
exp_res = "exp_res"
|
|
count = "count(*)"
|
|
default_term_expr = f"{ct.default_int64_field_name} in [0, 1]"
|
|
default_mix_expr = 'int64 >= 0 && varchar >= "0"'
|
|
default_expr = f"{ct.default_int64_field_name} >= 0"
|
|
default_invalid_expr = "varchar >= 0"
|
|
default_string_term_expr = f'{ct.default_string_field_name} in ["0", "1"]'
|
|
default_index_params = ct.default_index
|
|
binary_index_params = ct.default_binary_index
|
|
|
|
default_entities = ut.gen_entities(ut.default_nb, is_normal=True)
|
|
default_pos = 5
|
|
json_field = ct.default_json_field_name
|
|
default_int_field_name = ct.default_int64_field_name
|
|
default_float_field_name = "float"
|
|
default_string_field_name = "varchar"
|
|
|
|
|
|
class TestQueryParams(TestcaseBase):
|
|
"""
|
|
test Query interface
|
|
query(collection_name, expr, output_fields=None, partition_names=None, timeout=None)
|
|
"""
|
|
|
|
@pytest.fixture(scope="function", params=[True, False])
|
|
def enable_dynamic_field(self, request):
|
|
yield request.param
|
|
|
|
@pytest.fixture(scope="function", params=[True, False])
|
|
def random_primary_key(self, request):
|
|
yield request.param
|
|
|
|
@pytest.mark.tags(CaseLabel.L2)
|
|
def test_query_invalid(self):
|
|
"""
|
|
target: test query with invalid term expression
|
|
method: query with invalid term expr
|
|
expected: raise exception
|
|
"""
|
|
collection_w, entities = self.init_collection_general(prefix, insert_data=True, nb=10)[0:2]
|
|
term_expr = f"{default_int_field_name} in {entities[:default_pos]}"
|
|
error = {ct.err_code: 999, ct.err_msg: "cannot parse expression: int64 in"}
|
|
collection_w.query(term_expr, check_task=CheckTasks.err_res, check_items=error)
|
|
|
|
# check missing the template variable
|
|
expr = "int64 in {value_0}"
|
|
expr_params = {"value_1": [0, 1]}
|
|
error = {ct.err_code: 999, ct.err_msg: "the value of expression template variable name {value_0} is not found"}
|
|
collection_w.query(expr=expr, expr_params=expr_params, check_task=CheckTasks.err_res, check_items=error)
|
|
|
|
# check the template variable type dismatch
|
|
expr = "int64 in {value_0}"
|
|
expr_params = {"value_0": 1}
|
|
error = {ct.err_code: 999, ct.err_msg: "the value of term expression template variable {value_0} is not array"}
|
|
collection_w.query(expr=expr, expr_params=expr_params, check_task=CheckTasks.err_res, check_items=error)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query(self, enable_dynamic_field):
|
|
"""
|
|
target: test query
|
|
method: query with term expr
|
|
expected: verify query result
|
|
"""
|
|
# create collection, insert default_nb, load collection
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix, insert_data=True, enable_dynamic_field=enable_dynamic_field
|
|
)[0:2]
|
|
pos = 5
|
|
if enable_dynamic_field:
|
|
int_values = []
|
|
for vector in vectors[0]:
|
|
vector = vector[ct.default_int64_field_name]
|
|
int_values.append(vector)
|
|
res = [{ct.default_int64_field_name: int_values[i]} for i in range(pos)]
|
|
else:
|
|
int_values = vectors[0][ct.default_int64_field_name].values.tolist()
|
|
res = vectors[0].iloc[0:pos, :1].to_dict("records")
|
|
|
|
term_expr = f"{ct.default_int64_field_name} in {int_values[:pos]}"
|
|
collection_w.query(
|
|
term_expr,
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_auto_id_collection(self):
|
|
"""
|
|
target: test query with auto_id=True collection
|
|
method: test query with auto id
|
|
expected: query result is correct
|
|
"""
|
|
self._connect()
|
|
df = cf.gen_default_dataframe_data()
|
|
df[ct.default_int64_field_name] = None
|
|
(
|
|
insert_res,
|
|
_,
|
|
) = self.collection_wrap.construct_from_dataframe(
|
|
cf.gen_unique_str(prefix), df, primary_field=ct.default_int64_field_name, auto_id=True
|
|
)
|
|
assert self.collection_wrap.num_entities == ct.default_nb
|
|
ids = insert_res[1].primary_keys
|
|
pos = 5
|
|
res = df.iloc[:pos, :1].to_dict("records")
|
|
self.collection_wrap.create_index(ct.default_float_vec_field_name, index_params=ct.default_flat_index)
|
|
self.collection_wrap.load()
|
|
|
|
# query with all primary keys
|
|
term_expr_1 = f"{ct.default_int64_field_name} in {ids[:pos]}"
|
|
for i in range(5):
|
|
res[i][ct.default_int64_field_name] = ids[i]
|
|
self.collection_wrap.query(
|
|
term_expr_1,
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": self.collection_wrap.primary_field.name},
|
|
)
|
|
|
|
# query with part primary keys
|
|
term_expr_2 = f"{ct.default_int64_field_name} in {[ids[0], 0]}"
|
|
self.collection_wrap.query(
|
|
term_expr_2,
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res[:1], "pk_name": self.collection_wrap.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L2)
|
|
def test_query_non_string_expr(self):
|
|
"""
|
|
target: test query with non-string expr
|
|
method: query with non-string expr, eg 1, [] ..
|
|
expected: raise exception
|
|
"""
|
|
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
|
|
exprs = [1, 2.0, [], {}, ()]
|
|
error = {ct.err_code: 0, ct.err_msg: "The type of expr must be string"}
|
|
for expr in exprs:
|
|
collection_w.query(expr, check_task=CheckTasks.err_res, check_items=error)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
@pytest.mark.skip(reason="repeat with test_query, waiting for other expr")
|
|
def test_query_expr_term(self):
|
|
"""
|
|
target: test query with TermExpr
|
|
method: query with TermExpr
|
|
expected: query result is correct
|
|
"""
|
|
collection_w, vectors = self.init_collection_general(prefix, insert_data=True)[0:2]
|
|
res = vectors[0].iloc[:2, :1].to_dict("records")
|
|
collection_w.query(
|
|
default_term_expr,
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.fixture(scope="function", params=[0, 10, 100])
|
|
def offset(self, request):
|
|
yield request.param
|
|
|
|
@pytest.mark.tags(CaseLabel.L2)
|
|
@pytest.mark.skip("not stable")
|
|
def test_query_during_upsert(self):
|
|
"""
|
|
target: test query during upsert
|
|
method: 1. create a collection and query
|
|
2. query during upsert
|
|
3. compare two query results
|
|
expected: the two query results is the same
|
|
"""
|
|
upsert_nb = 1000
|
|
expr = f"int64 >= 0 && int64 <= {upsert_nb}"
|
|
collection_w = self.init_collection_general(prefix, True)[0]
|
|
res1 = collection_w.query(expr, output_fields=[default_float_field_name])[0]
|
|
|
|
def do_upsert():
|
|
data = cf.gen_default_data_for_upsert(upsert_nb)[0]
|
|
collection_w.upsert(data=data)
|
|
|
|
t = threading.Thread(target=do_upsert, args=())
|
|
t.start()
|
|
res2 = collection_w.query(expr, output_fields=[default_float_field_name])[0]
|
|
t.join()
|
|
assert [res1[i][default_float_field_name] for i in range(upsert_nb)] == [
|
|
res2[i][default_float_field_name] for i in range(upsert_nb)
|
|
]
|
|
|
|
|
|
class TestQueryOperation(TestcaseBase):
|
|
"""
|
|
******************************************************************
|
|
The following cases are used to test query interface operations
|
|
******************************************************************
|
|
"""
|
|
|
|
@pytest.mark.tags(CaseLabel.L2)
|
|
def test_query_expr_all_term_array(self):
|
|
"""
|
|
target: test query with all array term expr
|
|
method: query with all array value
|
|
expected: verify query result
|
|
"""
|
|
|
|
# init a collection and insert data
|
|
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, insert_data=True)[0:3]
|
|
|
|
# data preparation
|
|
int_values = vectors[0][ct.default_int64_field_name].values.tolist()
|
|
term_expr = f"{ct.default_int64_field_name} in {int_values}"
|
|
check_vec = vectors[0].iloc[:, [0]][0 : len(int_values)].to_dict("records")
|
|
|
|
# query all array value
|
|
collection_w.query(
|
|
term_expr,
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: check_vec, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_expr_half_term_array(self):
|
|
"""
|
|
target: test query with half array term expr
|
|
method: query with half array value
|
|
expected: verify query result
|
|
"""
|
|
|
|
half = ct.default_nb // 2
|
|
collection_w, partition_w, df_partition, df_default = self.insert_entities_into_two_partitions_in_half(half)
|
|
|
|
int_values = df_default[ct.default_int64_field_name].values.tolist()
|
|
term_expr = f"{ct.default_int64_field_name} in {int_values}"
|
|
res, _ = collection_w.query(term_expr)
|
|
assert len(res) == len(int_values)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_expr_repeated_term_array(self):
|
|
"""
|
|
target: test query with repeated term array on primary field with unique value
|
|
method: query with repeated array value
|
|
expected: return hit entities, no repeated
|
|
"""
|
|
collection_w, vectors, binary_raw_vectors = self.init_collection_general(prefix, insert_data=True)[0:3]
|
|
int_values = [0, 0, 0, 0]
|
|
term_expr = f"{ct.default_int64_field_name} in {int_values}"
|
|
res, _ = collection_w.query(term_expr)
|
|
assert len(res) == 1
|
|
assert res[0][ct.default_int64_field_name] == int_values[0]
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_dup_ids_dup_term_array(self):
|
|
"""
|
|
target: test query on duplicate primary keys with dup term array
|
|
method: 1.create collection and insert dup primary keys
|
|
2.query with dup term array
|
|
expected: todo
|
|
"""
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix))
|
|
df = cf.gen_default_dataframe_data(nb=100)
|
|
df[ct.default_int64_field_name] = 0
|
|
mutation_res, _ = collection_w.insert(df)
|
|
assert mutation_res.primary_keys == df[ct.default_int64_field_name].tolist()
|
|
collection_w.create_index(ct.default_float_vec_field_name, index_params=ct.default_flat_index)
|
|
collection_w.load()
|
|
term_expr = f"{ct.default_int64_field_name} in {[0, 0, 0]}"
|
|
res = df.iloc[:, :2].to_dict("records")
|
|
collection_w.query(
|
|
term_expr,
|
|
output_fields=["*"],
|
|
check_items=CheckTasks.check_query_results,
|
|
check_task={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_search_multi_logical_exprs(self):
|
|
"""
|
|
target: test the scenario which search with many logical expressions
|
|
method: 1. create collection
|
|
3. search with the expr that like: int64 == 0 || int64 == 1 ........
|
|
expected: run successfully
|
|
"""
|
|
c_name = cf.gen_unique_str(prefix)
|
|
collection_w = self.init_collection_wrap(name=c_name)
|
|
df = cf.gen_default_dataframe_data()
|
|
collection_w.insert(df)
|
|
collection_w.create_index(ct.default_float_vec_field_name, index_params=ct.default_flat_index)
|
|
collection_w.load()
|
|
|
|
multi_exprs = " || ".join(f"{default_int_field_name} == {i}" for i in range(60))
|
|
|
|
collection_w.load()
|
|
vectors_s = [[random.random() for _ in range(ct.default_dim)] for _ in range(ct.default_nq)]
|
|
limit = 1000
|
|
_, check_res = collection_w.search(
|
|
vectors_s[: ct.default_nq], ct.default_float_vec_field_name, ct.default_search_params, limit, multi_exprs
|
|
)
|
|
assert check_res is True
|
|
|
|
|
|
class TestQueryString(TestcaseBase):
|
|
"""
|
|
******************************************************************
|
|
The following cases are used to test query with string
|
|
******************************************************************
|
|
"""
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_string_expr_with_binary(self):
|
|
"""
|
|
target: test query string expr with binary
|
|
method: query string expr with binary
|
|
expected: verify query successfully
|
|
"""
|
|
collection_w, vectors = self.init_collection_general(prefix, insert_data=True, is_binary=True, is_index=False)[
|
|
0:2
|
|
]
|
|
collection_w.create_index(ct.default_binary_vec_field_name, binary_index_params)
|
|
collection_w.load()
|
|
assert collection_w.has_index()[0]
|
|
res, _ = collection_w.query(default_string_term_expr, output_fields=[ct.default_binary_vec_field_name])
|
|
assert len(res) == 2
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
@pytest.mark.skip(reason="issue 24637")
|
|
def test_query_after_insert_multi_threading(self):
|
|
"""
|
|
target: test data consistency after multi threading insert
|
|
method: multi threads insert, and query, compare queried data with original
|
|
expected: verify data consistency
|
|
"""
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix))
|
|
thread_num = 4
|
|
threads = []
|
|
primary_keys = []
|
|
df_list = []
|
|
|
|
# prepare original data for parallel insert
|
|
for i in range(thread_num):
|
|
df = cf.gen_default_dataframe_data(ct.default_nb, start=i * ct.default_nb)
|
|
df_list.append(df)
|
|
primary_key = df[ct.default_int64_field_name].values.tolist()
|
|
primary_keys.append(primary_key)
|
|
|
|
def insert(thread_i):
|
|
log.debug(f"In thread-{thread_i}")
|
|
mutation_res, _ = collection_w.insert(df_list[thread_i])
|
|
assert mutation_res.insert_count == ct.default_nb
|
|
assert mutation_res.primary_keys == primary_keys[thread_i]
|
|
|
|
for i in range(thread_num):
|
|
x = threading.Thread(target=insert, args=(i,))
|
|
threads.append(x)
|
|
x.start()
|
|
for t in threads:
|
|
t.join()
|
|
assert collection_w.num_entities == ct.default_nb * thread_num
|
|
|
|
# Check data consistency after parallel insert
|
|
collection_w.create_index(ct.default_float_vec_field_name, index_params=ct.default_flat_index)
|
|
collection_w.load()
|
|
df_dict_list = []
|
|
for df in df_list:
|
|
df_dict_list += df.to_dict("records")
|
|
output_fields = ["*"]
|
|
expression = "int64 >= 0"
|
|
collection_w.query(
|
|
expression,
|
|
output_fields=output_fields,
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: df_dict_list, "pk_name": collection_w.primary_field.name, "with_vec": True},
|
|
)
|
|
|
|
|
|
class TestQueryCount(TestcaseBase):
|
|
"""
|
|
test query count(*)
|
|
"""
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_count_compact_merge(self):
|
|
"""
|
|
target: test count after compact merge segments
|
|
method: 1. init 2 segments with same channel
|
|
2. compact
|
|
3. count
|
|
expected: verify count
|
|
"""
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), shards_num=1)
|
|
|
|
# init two segments
|
|
tmp_nb = 100
|
|
segment_num = 2
|
|
for i in range(segment_num):
|
|
df = cf.gen_default_dataframe_data(nb=tmp_nb, start=i * tmp_nb)
|
|
collection_w.insert(df)
|
|
collection_w.flush()
|
|
|
|
collection_w.create_index(ct.default_float_vec_field_name, ct.default_index)
|
|
collection_w.compact()
|
|
collection_w.wait_for_compaction_completed()
|
|
# recreate index wait for compactTo indexed
|
|
collection_w.create_index(ct.default_float_vec_field_name, ct.default_index)
|
|
|
|
collection_w.load()
|
|
segment_info, _ = self.utility_wrap.get_query_segment_info(collection_w.name)
|
|
assert len(segment_info) == 1
|
|
|
|
# count after compact
|
|
collection_w.query(
|
|
expr=default_expr,
|
|
output_fields=[ct.default_count_output],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: [{count: tmp_nb * segment_num}], "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L2)
|
|
def test_count_compact_delete(self):
|
|
"""
|
|
target: test count after delete-compact
|
|
method: 1. init segments
|
|
2. delete half ids and compact
|
|
3. count
|
|
expected: verify count
|
|
"""
|
|
# create -> index -> insert
|
|
collection_w = self.init_collection_wrap(cf.gen_unique_str(prefix), shards_num=1)
|
|
collection_w.create_index(ct.default_float_vec_field_name, index_params=ct.default_flat_index)
|
|
df = cf.gen_default_dataframe_data()
|
|
insert_res, _ = collection_w.insert(df)
|
|
|
|
# delete half entities, flush
|
|
half_expr = f"{ct.default_int64_field_name} in {[i for i in range(ct.default_nb // 2)]}"
|
|
collection_w.delete(half_expr)
|
|
assert collection_w.num_entities == ct.default_nb
|
|
|
|
# compact
|
|
collection_w.compact()
|
|
collection_w.wait_for_compaction_completed()
|
|
|
|
# load and count
|
|
collection_w.load()
|
|
collection_w.query(
|
|
expr=default_expr,
|
|
output_fields=[ct.default_count_output],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: [{count: ct.default_nb // 2}], "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
@pytest.mark.parametrize("index", ct.all_index_types[10:12])
|
|
def test_counts_expression_sparse_vectors(self, index):
|
|
"""
|
|
target: test count with expr
|
|
method: count with expr
|
|
expected: verify count
|
|
"""
|
|
self._connect()
|
|
c_name = cf.gen_unique_str(prefix)
|
|
schema = cf.gen_default_sparse_schema()
|
|
collection_w = self.init_collection_wrap(c_name, schema=schema)
|
|
data = cf.gen_default_list_sparse_data()
|
|
collection_w.insert(data)
|
|
params = cf.get_index_params_params(index)
|
|
index_params = {"index_type": index, "metric_type": "IP", "params": params}
|
|
collection_w.create_index(ct.default_sparse_vec_field_name, index_params, index_name=index)
|
|
collection_w.load()
|
|
collection_w.query(
|
|
expr=default_expr,
|
|
output_fields=[count],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: [{count: ct.default_nb}], "pk_name": collection_w.primary_field.name},
|
|
)
|
|
expr = "int64 > 50 && int64 < 100 && float < 75"
|
|
collection_w.query(
|
|
expr=expr,
|
|
output_fields=[count],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: [{count: 24}], "pk_name": collection_w.primary_field.name},
|
|
)
|
|
batch_size = 100
|
|
collection_w.query_iterator(
|
|
batch_size=batch_size,
|
|
expr=default_expr,
|
|
check_task=CheckTasks.check_query_iterator,
|
|
check_items={"count": ct.default_nb, "batch_size": batch_size, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
@pytest.mark.repeat(3)
|
|
@pytest.mark.skip(reason="issue #36538")
|
|
def test_count_query_search_after_release_partition_load(self):
|
|
"""
|
|
target: test query count(*) after release collection and load partition
|
|
method: 1. create a collection and 2 partitions with nullable and default value fields
|
|
2. insert data
|
|
3. load one partition
|
|
4. delete half data in each partition
|
|
5. release the collection and load one partition
|
|
6. search
|
|
expected: No exception
|
|
"""
|
|
# insert data
|
|
collection_w = self.init_collection_general(prefix, True, 200, partition_num=1, is_index=True)[0]
|
|
collection_w.query(
|
|
expr="",
|
|
output_fields=[ct.default_count_output],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={"exp_res": [{ct.default_count_output: 200}], "pk_name": collection_w.primary_field.name},
|
|
)
|
|
collection_w.release()
|
|
partition_w1, partition_w2 = collection_w.partitions
|
|
# load
|
|
partition_w1.load()
|
|
# delete data
|
|
delete_ids = [i for i in range(50, 150)]
|
|
collection_w.delete(f"int64 in {delete_ids}")
|
|
# release
|
|
collection_w.release()
|
|
# partition_w1.load()
|
|
collection_w.load(partition_names=[partition_w1.name])
|
|
# search on collection, partition1, partition2
|
|
collection_w.query(
|
|
expr="",
|
|
output_fields=[ct.default_count_output],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={"exp_res": [{ct.default_count_output: 50}], "pk_name": collection_w.primary_field.name},
|
|
)
|
|
partition_w1.query(
|
|
expr="",
|
|
output_fields=[ct.default_count_output],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={"exp_res": [{ct.default_count_output: 50}], "pk_name": collection_w.primary_field.name},
|
|
)
|
|
vectors = [[random.random() for _ in range(ct.default_dim)] for _ in range(ct.default_nq)]
|
|
collection_w.search(
|
|
vectors[:1],
|
|
ct.default_float_vec_field_name,
|
|
ct.default_search_params,
|
|
200,
|
|
partition_names=[partition_w2.name],
|
|
check_task=CheckTasks.err_res,
|
|
check_items={ct.err_code: 1, ct.err_msg: "not loaded"},
|
|
)
|
|
|
|
|
|
class TestQueryNoneAndDefaultData(TestcaseBase):
|
|
"""
|
|
test Query interface with none and default data
|
|
query(collection_name, expr, output_fields=None, partition_names=None, timeout=None)
|
|
"""
|
|
|
|
@pytest.fixture(scope="function", params=[True, False])
|
|
def enable_dynamic_field(self, request):
|
|
yield request.param
|
|
|
|
@pytest.fixture(scope="function", params=["STL_SORT", "INVERTED"])
|
|
def numeric_scalar_index(self, request):
|
|
yield request.param
|
|
|
|
@pytest.fixture(scope="function", params=["TRIE", "INVERTED", "BITMAP"])
|
|
def varchar_scalar_index(self, request):
|
|
yield request.param
|
|
|
|
@pytest.fixture(scope="function", params=[0, 0.5, 1])
|
|
def null_data_percent(self, request):
|
|
yield request.param
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_by_normal_with_none_data(self, enable_dynamic_field, null_data_percent):
|
|
"""
|
|
target: test query with none data
|
|
method: query with term expr with nullable fields, insert data including none
|
|
expected: verify query result
|
|
"""
|
|
# create collection, insert default_nb, load collection
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix,
|
|
insert_data=True,
|
|
enable_dynamic_field=enable_dynamic_field,
|
|
nullable_fields={default_float_field_name: null_data_percent},
|
|
)[0:2]
|
|
pos = 5
|
|
if enable_dynamic_field:
|
|
int_values, float_values = [], []
|
|
for vector in vectors[0]:
|
|
int_values.append(vector[ct.default_int64_field_name])
|
|
float_values.append(vector[default_float_field_name])
|
|
res = [
|
|
{ct.default_int64_field_name: int_values[i], default_float_field_name: float_values[i]}
|
|
for i in range(pos)
|
|
]
|
|
else:
|
|
int_values = vectors[0][ct.default_int64_field_name].values.tolist()
|
|
res = vectors[0].iloc[0:pos, :2].to_dict("records")
|
|
|
|
term_expr = f"{ct.default_int64_field_name} in {int_values[:pos]}"
|
|
collection_w.query(
|
|
term_expr,
|
|
output_fields=[ct.default_int64_field_name, default_float_field_name],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_by_expr_none_with_none_data(self, enable_dynamic_field, null_data_percent):
|
|
"""
|
|
target: test query by none expr with nullable fields, insert data including none
|
|
method: query by expr None after inserting data including none
|
|
expected: verify query result
|
|
"""
|
|
# create collection, insert default_nb, load collection
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix,
|
|
insert_data=True,
|
|
enable_dynamic_field=enable_dynamic_field,
|
|
nullable_fields={default_float_field_name: null_data_percent},
|
|
)[0:2]
|
|
pos = 5
|
|
if enable_dynamic_field:
|
|
int_values, float_values = [], []
|
|
for vector in vectors[0]:
|
|
int_values.append(vector[ct.default_int64_field_name])
|
|
float_values.append(vector[default_float_field_name])
|
|
res = [
|
|
{ct.default_int64_field_name: int_values[i], default_float_field_name: float_values[i]}
|
|
for i in range(pos)
|
|
]
|
|
else:
|
|
res = vectors[0].iloc[0:pos, :2].to_dict("records")
|
|
|
|
term_expr = ""
|
|
collection_w.query(
|
|
term_expr,
|
|
output_fields=[ct.default_int64_field_name, default_float_field_name],
|
|
limit=pos,
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_by_nullable_field_with_none_data(self):
|
|
"""
|
|
target: test query with nullable fields expr, insert data including none into nullable Fields
|
|
method: query by nullable field expr after inserting data including none
|
|
expected: verify query result
|
|
"""
|
|
# create collection, insert default_nb, load collection
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix, insert_data=True, enable_dynamic_field=True, nullable_fields={default_float_field_name: 0.5}
|
|
)[0:2]
|
|
pos = 5
|
|
int_values, float_values = [], []
|
|
for vector in vectors[0]:
|
|
int_values.append(vector[ct.default_int64_field_name])
|
|
float_values.append(vector[default_float_field_name])
|
|
res = [
|
|
{ct.default_int64_field_name: int_values[i], default_float_field_name: float_values[i]} for i in range(pos)
|
|
]
|
|
|
|
term_expr = f"{default_float_field_name} < {pos}"
|
|
collection_w.query(
|
|
term_expr,
|
|
output_fields=[ct.default_int64_field_name, default_float_field_name],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_after_none_data_all_field_datatype(
|
|
self, varchar_scalar_index, numeric_scalar_index, null_data_percent
|
|
):
|
|
"""
|
|
target: test query after different index on scalar fields
|
|
method: query after different index on nullable fields
|
|
expected: verify query result
|
|
"""
|
|
# 1. initialize with data
|
|
nullable_fields = {
|
|
ct.default_int32_field_name: null_data_percent,
|
|
ct.default_int16_field_name: null_data_percent,
|
|
ct.default_int8_field_name: null_data_percent,
|
|
ct.default_bool_field_name: null_data_percent,
|
|
ct.default_float_field_name: null_data_percent,
|
|
ct.default_double_field_name: null_data_percent,
|
|
ct.default_string_field_name: null_data_percent,
|
|
}
|
|
# 2. create collection, insert default_nb
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix, True, 1000, is_all_data_type=True, is_index=False, nullable_fields=nullable_fields
|
|
)[0:2]
|
|
# 3. create index on vector field and load
|
|
index = "HNSW"
|
|
params = cf.get_index_params_params(index)
|
|
default_index = {"index_type": index, "params": params, "metric_type": "COSINE"}
|
|
vector_name_list = cf.extract_vector_field_name_list(collection_w)
|
|
vector_name_list.append(ct.default_float_vec_field_name)
|
|
for vector_name in vector_name_list:
|
|
collection_w.create_index(vector_name, default_index)
|
|
# 4. create index on scalar field with None data
|
|
scalar_index_params = {"index_type": varchar_scalar_index, "params": {}}
|
|
collection_w.create_index(ct.default_string_field_name, scalar_index_params)
|
|
# 5. create index on scalar field with default data
|
|
scalar_index_params = {"index_type": numeric_scalar_index, "params": {}}
|
|
collection_w.create_index(ct.default_int64_field_name, scalar_index_params)
|
|
collection_w.create_index(ct.default_int32_field_name, scalar_index_params)
|
|
collection_w.create_index(ct.default_int16_field_name, scalar_index_params)
|
|
collection_w.create_index(ct.default_int8_field_name, scalar_index_params)
|
|
if numeric_scalar_index != "STL_SORT":
|
|
collection_w.create_index(ct.default_bool_field_name, scalar_index_params)
|
|
collection_w.create_index(ct.default_float_field_name, scalar_index_params)
|
|
collection_w.load()
|
|
pos = 5
|
|
int64_values, float_values = [], []
|
|
scalar_fields = vectors[0]
|
|
for i in range(pos):
|
|
int64_values.append(scalar_fields[0][i])
|
|
float_values.append(scalar_fields[5][i])
|
|
res = [
|
|
{ct.default_int64_field_name: int64_values[i], default_float_field_name: float_values[i]}
|
|
for i in range(pos)
|
|
]
|
|
|
|
term_expr = f"0 <= {ct.default_int64_field_name} < {pos}"
|
|
collection_w.query(
|
|
term_expr,
|
|
output_fields=[ct.default_int64_field_name, ct.default_float_field_name],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_default_value_with_insert(self, enable_dynamic_field):
|
|
"""
|
|
target: test query normal case with default value set
|
|
method: create connection, collection with default value set, insert and query
|
|
expected: query successfully and verify query result
|
|
"""
|
|
# 1. initialize with data
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix,
|
|
True,
|
|
enable_dynamic_field=enable_dynamic_field,
|
|
default_value_fields={ct.default_float_field_name: np.float32(10.0)},
|
|
)[0:2]
|
|
pos = 5
|
|
if enable_dynamic_field:
|
|
int_values, float_values = [], []
|
|
for vector in vectors[0]:
|
|
int_values.append(vector[ct.default_int64_field_name])
|
|
float_values.append(vector[default_float_field_name])
|
|
res = [
|
|
{ct.default_int64_field_name: int_values[i], default_float_field_name: float_values[i]}
|
|
for i in range(pos)
|
|
]
|
|
else:
|
|
int_values = vectors[0][ct.default_int64_field_name].values.tolist()
|
|
res = vectors[0].iloc[0:pos, :2].to_dict("records")
|
|
|
|
term_expr = f"{ct.default_int64_field_name} in {int_values[:pos]}"
|
|
# 2. query
|
|
collection_w.query(
|
|
term_expr,
|
|
output_fields=[ct.default_int64_field_name, default_float_field_name],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_default_value_without_insert(self, enable_dynamic_field):
|
|
"""
|
|
target: test query normal case with default value set
|
|
method: create connection, collection with default value set, no insert and query
|
|
expected: query successfully and verify query result
|
|
"""
|
|
# 1. initialize with data
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix,
|
|
False,
|
|
enable_dynamic_field=enable_dynamic_field,
|
|
default_value_fields={ct.default_float_field_name: np.float32(10.0)},
|
|
)[0:2]
|
|
|
|
term_expr = f"{ct.default_int64_field_name} > 0"
|
|
# 2. query
|
|
collection_w.query(
|
|
term_expr,
|
|
output_fields=[ct.default_int64_field_name, default_float_field_name],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: [], "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_after_default_data_all_field_datatype(self, varchar_scalar_index, numeric_scalar_index):
|
|
"""
|
|
target: test query after different index on default value data
|
|
method: test query after different index on default value and corresponding search params
|
|
expected: query successfully and verify query result
|
|
"""
|
|
# 1. initialize with data
|
|
default_value_fields = {
|
|
ct.default_int32_field_name: np.int32(1),
|
|
ct.default_int16_field_name: np.int32(2),
|
|
ct.default_int8_field_name: np.int32(3),
|
|
ct.default_bool_field_name: True,
|
|
ct.default_float_field_name: np.float32(10.0),
|
|
ct.default_double_field_name: 10.0,
|
|
ct.default_string_field_name: "1",
|
|
}
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix,
|
|
True,
|
|
1000,
|
|
partition_num=1,
|
|
is_all_data_type=True,
|
|
is_index=False,
|
|
default_value_fields=default_value_fields,
|
|
)[0:2]
|
|
# 2. create index on vector field and load
|
|
index = "HNSW"
|
|
params = cf.get_index_params_params(index)
|
|
default_index = {"index_type": index, "params": params, "metric_type": "L2"}
|
|
vector_name_list = cf.extract_vector_field_name_list(collection_w)
|
|
vector_name_list.append(ct.default_float_vec_field_name)
|
|
for vector_name in vector_name_list:
|
|
collection_w.create_index(vector_name, default_index)
|
|
# 3. create index on scalar field with None data
|
|
scalar_index_params = {"index_type": varchar_scalar_index, "params": {}}
|
|
collection_w.create_index(ct.default_string_field_name, scalar_index_params)
|
|
# 4. create index on scalar field with default data
|
|
scalar_index_params = {"index_type": numeric_scalar_index, "params": {}}
|
|
collection_w.create_index(ct.default_int64_field_name, scalar_index_params)
|
|
collection_w.create_index(ct.default_int32_field_name, scalar_index_params)
|
|
collection_w.create_index(ct.default_int16_field_name, scalar_index_params)
|
|
collection_w.create_index(ct.default_int8_field_name, scalar_index_params)
|
|
if numeric_scalar_index != "STL_SORT":
|
|
collection_w.create_index(ct.default_bool_field_name, scalar_index_params)
|
|
collection_w.create_index(ct.default_float_field_name, scalar_index_params)
|
|
collection_w.load()
|
|
pos = 5
|
|
int64_values, float_values = [], []
|
|
scalar_fields = vectors[0]
|
|
for i in range(pos):
|
|
int64_values.append(scalar_fields[0][i])
|
|
float_values.append(scalar_fields[5][i])
|
|
res = [
|
|
{ct.default_int64_field_name: int64_values[i], default_float_field_name: float_values[i]}
|
|
for i in range(pos)
|
|
]
|
|
|
|
term_expr = f"0 <= {ct.default_int64_field_name} < {pos}"
|
|
# 5. query
|
|
collection_w.query(
|
|
term_expr,
|
|
output_fields=[ct.default_int64_field_name, ct.default_float_field_name],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
@pytest.mark.skip(reason="issue #36003")
|
|
def test_query_both_default_value_non_data(self, enable_dynamic_field):
|
|
"""
|
|
target: test query normal case with default value set
|
|
method: create connection, collection with default value set, insert and query
|
|
expected: query successfully and verify query result
|
|
"""
|
|
# 1. initialize with data
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix,
|
|
True,
|
|
enable_dynamic_field=enable_dynamic_field,
|
|
nullable_fields={ct.default_float_field_name: 1},
|
|
default_value_fields={ct.default_float_field_name: np.float32(10.0)},
|
|
)[0:2]
|
|
pos = 5
|
|
if enable_dynamic_field:
|
|
int_values, float_values = [], []
|
|
for vector in vectors[0]:
|
|
int_values.append(vector[ct.default_int64_field_name])
|
|
float_values.append(vector[default_float_field_name])
|
|
res = [
|
|
{ct.default_int64_field_name: int_values[i], default_float_field_name: float_values[i]}
|
|
for i in range(pos)
|
|
]
|
|
else:
|
|
res = vectors[0].iloc[0:pos, :2].to_dict("records")
|
|
|
|
term_expr = f"{ct.default_float_field_name} in [10.0]"
|
|
collection_w.query(
|
|
term_expr,
|
|
output_fields=[ct.default_int64_field_name, default_float_field_name],
|
|
limit=pos,
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
@pytest.mark.tags(CaseLabel.GPU)
|
|
def test_query_after_different_index_with_params_none_default_data(
|
|
self, varchar_scalar_index, numeric_scalar_index, null_data_percent
|
|
):
|
|
"""
|
|
target: test query after different index
|
|
method: test query after different index on none default data
|
|
expected: query successfully and verify query result
|
|
"""
|
|
# 1. initialize with data
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix,
|
|
True,
|
|
1000,
|
|
partition_num=1,
|
|
is_all_data_type=True,
|
|
is_index=False,
|
|
nullable_fields={ct.default_string_field_name: null_data_percent},
|
|
default_value_fields={ct.default_float_field_name: np.float32(10.0)},
|
|
)[0:2]
|
|
# 2. create index on vector field and load
|
|
index = "HNSW"
|
|
params = cf.get_index_params_params(index)
|
|
default_index = {"index_type": index, "params": params, "metric_type": "COSINE"}
|
|
vector_name_list = cf.extract_vector_field_name_list(collection_w)
|
|
vector_name_list.append(ct.default_float_vec_field_name)
|
|
for vector_name in vector_name_list:
|
|
collection_w.create_index(vector_name, default_index)
|
|
# 3. create index on scalar field with None data
|
|
scalar_index_params = {"index_type": varchar_scalar_index, "params": {}}
|
|
collection_w.create_index(ct.default_string_field_name, scalar_index_params)
|
|
# 4. create index on scalar field with default data
|
|
scalar_index_params = {"index_type": numeric_scalar_index, "params": {}}
|
|
collection_w.create_index(ct.default_float_field_name, scalar_index_params)
|
|
collection_w.load()
|
|
pos = 5
|
|
int64_values, float_values = [], []
|
|
scalar_fields = vectors[0]
|
|
for i in range(pos):
|
|
int64_values.append(scalar_fields[0][i])
|
|
float_values.append(scalar_fields[5][i])
|
|
res = [
|
|
{ct.default_int64_field_name: int64_values[i], default_float_field_name: float_values[i]}
|
|
for i in range(pos)
|
|
]
|
|
|
|
term_expr = f"{ct.default_int64_field_name} in {int64_values[:pos]}"
|
|
# 5. query
|
|
collection_w.query(
|
|
term_expr,
|
|
output_fields=[ct.default_int64_field_name, ct.default_float_field_name],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_iterator_with_none_data(self, null_data_percent):
|
|
"""
|
|
target: test query iterator normal with none data
|
|
method: 1. query iterator
|
|
2. check the result, expect pk
|
|
expected: query successfully
|
|
"""
|
|
# 1. initialize with data
|
|
batch_size = 100
|
|
collection_w = self.init_collection_general(
|
|
prefix, True, is_index=False, nullable_fields={ct.default_string_field_name: null_data_percent}
|
|
)[0]
|
|
collection_w.create_index(ct.default_float_vec_field_name, {"metric_type": "L2"})
|
|
collection_w.load()
|
|
# 2. search iterator
|
|
expr = "int64 >= 0"
|
|
collection_w.query_iterator(
|
|
batch_size,
|
|
expr=expr,
|
|
check_task=CheckTasks.check_query_iterator,
|
|
check_items={"count": ct.default_nb, "pk_name": collection_w.primary_field.name, "batch_size": batch_size},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
@pytest.mark.skip(reason="issue #36123")
|
|
def test_query_normal_none_data_partition_key(self, enable_dynamic_field, null_data_percent):
|
|
"""
|
|
target: test query normal case with none data inserted
|
|
method: create connection, collection with nullable fields, insert data including none, and query
|
|
expected: query successfully and verify query result
|
|
"""
|
|
# 1. initialize with data
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix,
|
|
True,
|
|
enable_dynamic_field=enable_dynamic_field,
|
|
nullable_fields={ct.default_float_field_name: null_data_percent},
|
|
is_partition_key=ct.default_float_field_name,
|
|
)[0:2]
|
|
pos = 5
|
|
if enable_dynamic_field:
|
|
int_values, float_values = [], []
|
|
for vector in vectors[0]:
|
|
int_values.append(vector[ct.default_int64_field_name])
|
|
float_values.append(vector[default_float_field_name])
|
|
res = [
|
|
{ct.default_int64_field_name: int_values[i], default_float_field_name: float_values[i]}
|
|
for i in range(pos)
|
|
]
|
|
else:
|
|
int_values = vectors[0][ct.default_int64_field_name].values.tolist()
|
|
res = vectors[0].iloc[0:pos, :2].to_dict("records")
|
|
|
|
term_expr = f"{ct.default_int64_field_name} in {int_values[:pos]}"
|
|
collection_w.query(
|
|
term_expr,
|
|
output_fields=[ct.default_int64_field_name, default_float_field_name],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
@pytest.mark.skip(reason="issue #36538")
|
|
def test_query_none_count(self, null_data_percent):
|
|
"""
|
|
target: test query count(*) with None and default data
|
|
method: 1. create a collection and 2 partitions with nullable and default value fields
|
|
2. insert data
|
|
3. load one partition
|
|
4. delete half data in each partition
|
|
5. release the collection and load one partition
|
|
6. search
|
|
expected: No exception
|
|
"""
|
|
# insert data
|
|
collection_w = self.init_collection_general(
|
|
prefix,
|
|
True,
|
|
200,
|
|
partition_num=1,
|
|
is_index=True,
|
|
nullable_fields={ct.default_float_field_name: null_data_percent},
|
|
default_value_fields={ct.default_string_field_name: "data"},
|
|
)[0]
|
|
collection_w.query(
|
|
expr="",
|
|
output_fields=[ct.default_count_output],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={"exp_res": [{ct.default_count_output: 200}], "pk_name": collection_w.primary_field.name},
|
|
)
|
|
collection_w.release()
|
|
partition_w1, partition_w2 = collection_w.partitions
|
|
# load
|
|
partition_w1.load()
|
|
# delete data
|
|
delete_ids = [i for i in range(50, 150)]
|
|
collection_w.delete(f"int64 in {delete_ids}")
|
|
# release
|
|
collection_w.release()
|
|
# partition_w1.load()
|
|
collection_w.load(partition_names=[partition_w1.name])
|
|
# search on collection, partition1, partition2
|
|
collection_w.query(
|
|
expr="",
|
|
output_fields=[ct.default_count_output],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={"exp_res": [{ct.default_count_output: 50}], "pk_name": collection_w.primary_field.name},
|
|
)
|
|
partition_w1.query(
|
|
expr="",
|
|
output_fields=[ct.default_count_output],
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={"exp_res": [{ct.default_count_output: 50}], "pk_name": collection_w.primary_field.name},
|
|
)
|
|
vectors = [[random.random() for _ in range(ct.default_dim)] for _ in range(ct.default_nq)]
|
|
collection_w.search(
|
|
vectors[:1],
|
|
ct.default_float_vec_field_name,
|
|
ct.default_search_params,
|
|
200,
|
|
partition_names=[partition_w2.name],
|
|
check_task=CheckTasks.err_res,
|
|
check_items={ct.err_code: 1, ct.err_msg: "not loaded"},
|
|
)
|
|
|
|
|
|
class TestQueryTextMatch(TestcaseBase):
|
|
"""
|
|
******************************************************************
|
|
The following cases are used to test query text match
|
|
******************************************************************
|
|
"""
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
@pytest.mark.parametrize("enable_partition_key", [True, False])
|
|
@pytest.mark.parametrize("enable_inverted_index", [True, False])
|
|
@pytest.mark.parametrize("tokenizer", ["standard"])
|
|
def test_query_text_match_en_normal(self, tokenizer, enable_inverted_index, enable_partition_key):
|
|
"""
|
|
target: test text match normal
|
|
method: 1. enable text match and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: text match successfully and result is correct
|
|
"""
|
|
analyzer_params = {
|
|
"tokenizer": tokenizer,
|
|
}
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
is_partition_key=enable_partition_key,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 3000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
if tokenizer == "jieba":
|
|
language = "zh"
|
|
fake = fake_zh
|
|
else:
|
|
language = "en"
|
|
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"word": fake.word().lower(),
|
|
"sentence": fake.sentence().lower(),
|
|
"paragraph": fake.paragraph().lower(),
|
|
"text": fake.text().lower(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
# only if the collection is flushed, the inverted index ca be applied.
|
|
# growing segment may be not applied, although in strong consistency.
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
if enable_inverted_index:
|
|
collection_w.create_index("word", {"index_type": "INVERTED"})
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
# query single field for one token
|
|
for field in text_fields:
|
|
most_common_tokens = wf_map[field].most_common(10)
|
|
mid = len(most_common_tokens) // 2
|
|
idx = random.randint(0, max(0, mid - 1))
|
|
token = most_common_tokens[idx][0]
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
assert len(res) > 0
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
assert token in r[field]
|
|
|
|
# verify inverted index
|
|
if enable_inverted_index:
|
|
if field == "word":
|
|
expr = f"{field} == '{token}'"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
assert r[field] == token
|
|
# query single field for multi-word
|
|
for field in text_fields:
|
|
# match top 10 most common words
|
|
top_10_tokens = []
|
|
for word, count in wf_map[field].most_common(10):
|
|
top_10_tokens.append(word)
|
|
string_of_top_10_words = " ".join(top_10_tokens)
|
|
expr = f"text_match({field}, '{string_of_top_10_words}')"
|
|
log.info(f"expr {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
assert any([token in r[field] for token in top_10_tokens])
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
@pytest.mark.parametrize("enable_partition_key", [True, False])
|
|
@pytest.mark.parametrize("enable_inverted_index", [True, False])
|
|
@pytest.mark.parametrize("lang_type", ["chinese"])
|
|
def test_query_text_match_zh_normal(self, lang_type, enable_inverted_index, enable_partition_key):
|
|
"""
|
|
target: test text match normal
|
|
method: 1. enable text match and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: text match successfully and result is correct
|
|
"""
|
|
analyzer_params = {
|
|
"type": lang_type,
|
|
}
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
is_partition_key=enable_partition_key,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 3000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
if lang_type == "chinese":
|
|
language = "zh"
|
|
fake = fake_zh
|
|
else:
|
|
language = "en"
|
|
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"word": fake.word().lower(),
|
|
"sentence": fake.sentence().lower(),
|
|
"paragraph": fake.paragraph().lower(),
|
|
"text": fake.text().lower(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
# only if the collection is flushed, the inverted index ca be applied.
|
|
# growing segment may be not applied, although in strong consistency.
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
if enable_inverted_index:
|
|
collection_w.create_index("word", {"index_type": "INVERTED"})
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
|
|
# query with blank space and punctuation marks
|
|
for field in text_fields:
|
|
expr = f"text_match({field}, ' ') or text_match({field}, ',') or text_match({field}, '.')"
|
|
log.info(f"expr {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) == 0
|
|
|
|
# query single field for one token
|
|
for field in text_fields:
|
|
most_common_tokens = wf_map[field].most_common(10)
|
|
mid = len(most_common_tokens) // 2
|
|
idx = random.randint(0, max(0, mid - 1))
|
|
token = most_common_tokens[idx][0]
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
assert len(res) > 0
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
assert token in r[field]
|
|
|
|
# verify inverted index
|
|
if enable_inverted_index:
|
|
if field == "word":
|
|
expr = f"{field} == '{token}'"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
assert r[field] == token
|
|
# query single field for multi-word
|
|
for field in text_fields:
|
|
# match top 10 most common words
|
|
top_10_tokens = []
|
|
for word, count in wf_map[field].most_common(10):
|
|
top_10_tokens.append(word)
|
|
string_of_top_10_words = " ".join(top_10_tokens)
|
|
expr = f"text_match({field}, '{string_of_top_10_words}')"
|
|
log.info(f"expr {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
assert any([token in r[field] for token in top_10_tokens]), (
|
|
f"top 10 tokens {top_10_tokens} not in {r[field]}"
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
@pytest.mark.parametrize("enable_partition_key", [True, False])
|
|
@pytest.mark.parametrize("enable_inverted_index", [True, False])
|
|
@pytest.mark.parametrize("tokenizer", ["icu"])
|
|
def test_query_text_match_with_icu_tokenizer(self, tokenizer, enable_inverted_index, enable_partition_key):
|
|
"""
|
|
target: test text match with icu tokenizer
|
|
method: 1. enable text match and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: text match successfully and result is correct
|
|
"""
|
|
analyzer_params = {
|
|
"tokenizer": tokenizer,
|
|
}
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
is_partition_key=enable_partition_key,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 3000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = ICUTextGenerator()
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"word": fake.word().lower(),
|
|
"sentence": fake.sentence().lower(),
|
|
"paragraph": fake.paragraph().lower(),
|
|
"text": fake.text().lower(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
# only if the collection is flushed, the inverted index ca be applied.
|
|
# growing segment may be not applied, although in strong consistency.
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
if enable_inverted_index:
|
|
collection_w.create_index("word", {"index_type": "INVERTED"})
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents_with_analyzer_params(df[field].tolist(), analyzer_params)
|
|
# query single field for one token
|
|
for field in text_fields:
|
|
most_common_tokens = wf_map[field].most_common(10)
|
|
mid = len(most_common_tokens) // 2
|
|
idx = random.randint(0, max(0, mid - 1))
|
|
token = most_common_tokens[idx][0]
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
assert len(res) > 0
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
assert token in r[field]
|
|
|
|
# verify inverted index
|
|
if enable_inverted_index:
|
|
if field == "word":
|
|
expr = f"{field} == '{token}'"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
assert r[field] == token
|
|
# query single field for multi-word
|
|
for field in text_fields:
|
|
# match top 10 most common words
|
|
top_10_tokens = []
|
|
for word, count in wf_map[field].most_common(10):
|
|
top_10_tokens.append(word)
|
|
string_of_top_10_words = " ".join(top_10_tokens)
|
|
expr = f"text_match({field}, '{string_of_top_10_words}')"
|
|
log.info(f"expr {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
assert any([token in r[field] for token in top_10_tokens])
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
@pytest.mark.parametrize("enable_partition_key", [True])
|
|
@pytest.mark.parametrize("enable_inverted_index", [True])
|
|
@pytest.mark.parametrize("tokenizer", ["jieba", "standard"])
|
|
def test_query_text_match_with_growing_segment(self, tokenizer, enable_inverted_index, enable_partition_key):
|
|
"""
|
|
target: test text match normal
|
|
method: 1. enable text match and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: text match successfully and result is correct
|
|
"""
|
|
analyzer_params = {
|
|
"tokenizer": tokenizer,
|
|
}
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
is_partition_key=enable_partition_key,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 3000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
if tokenizer == "jieba":
|
|
language = "zh"
|
|
fake = fake_zh
|
|
else:
|
|
language = "en"
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
if enable_inverted_index:
|
|
collection_w.create_index("word", {"index_type": "INVERTED"})
|
|
collection_w.load()
|
|
# generate growing segment
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"word": fake.word().lower(),
|
|
"sentence": fake.sentence().lower(),
|
|
"paragraph": fake.paragraph().lower(),
|
|
"text": fake.text().lower(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
# Use server-side analyzer to get tokens consistent with what Milvus indexes
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents_with_analyzer_params(df[field].tolist(), analyzer_params)
|
|
|
|
# Wait for growing-segment BM25 index to be ready.
|
|
# Poll every 1s (up to 30s) using the top token from the "word" field.
|
|
probe_token = wf_map["word"].most_common(1)[0][0]
|
|
probe_expr = f"text_match(word, '{probe_token}')"
|
|
ready = False
|
|
for _ in range(30):
|
|
res, _ = collection_w.query(expr=probe_expr, output_fields=["id"])
|
|
if len(res) > 0:
|
|
ready = True
|
|
break
|
|
time.sleep(1)
|
|
assert ready, f"Growing segment BM25 index not ready within 30s (probe token: '{probe_token}')"
|
|
|
|
# query single field for one token
|
|
for field in text_fields:
|
|
most_common_tokens = wf_map[field].most_common(10)
|
|
mid = len(most_common_tokens) // 2
|
|
idx = random.randint(0, max(0, mid - 1))
|
|
token = most_common_tokens[idx][0]
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) > 0
|
|
|
|
# query single field for multi-word
|
|
for field in text_fields:
|
|
# match top 10 most common words
|
|
top_10_tokens = []
|
|
for word, count in wf_map[field].most_common(10):
|
|
top_10_tokens.append(word)
|
|
string_of_top_10_words = " ".join(top_10_tokens)
|
|
expr = f"text_match({field}, '{string_of_top_10_words}')"
|
|
log.info(f"expr {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) > 0
|
|
|
|
# flush and then query again
|
|
collection_w.flush()
|
|
for field in text_fields:
|
|
# match top 10 most common words
|
|
top_10_tokens = []
|
|
for word, count in wf_map[field].most_common(10):
|
|
top_10_tokens.append(word)
|
|
string_of_top_10_words = " ".join(top_10_tokens)
|
|
expr = f"text_match({field}, '{string_of_top_10_words}')"
|
|
log.info(f"expr {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) > 0
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
@pytest.mark.parametrize("enable_partition_key", [True, False])
|
|
@pytest.mark.parametrize("enable_inverted_index", [True, False])
|
|
@pytest.mark.parametrize("lang_type", ["chinese"])
|
|
def test_query_text_match_zh_en_mix(self, lang_type, enable_inverted_index, enable_partition_key):
|
|
"""
|
|
target: test text match normal
|
|
method: 1. enable text match and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: text match successfully and result is correct
|
|
"""
|
|
analyzer_params = {
|
|
"type": lang_type,
|
|
}
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
is_partition_key=enable_partition_key,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 3000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
if lang_type == "chinese":
|
|
language = "zh"
|
|
fake = fake_zh
|
|
else:
|
|
language = "en"
|
|
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"word": fake.word().lower() + " " + fake_en.word().lower(),
|
|
"sentence": fake.sentence().lower() + " " + fake_en.sentence().lower(),
|
|
"paragraph": fake.paragraph().lower() + " " + fake_en.paragraph().lower(),
|
|
"text": fake.text().lower() + " " + fake_en.text().lower(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
# only if the collection is flushed, the inverted index ca be applied.
|
|
# growing segment may be not applied, although in strong consistency.
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
if enable_inverted_index:
|
|
collection_w.create_index("word", {"index_type": "INVERTED"})
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
# query single field for one token
|
|
for field in text_fields:
|
|
most_common_tokens = wf_map[field].most_common(10)
|
|
mid = len(most_common_tokens) // 2
|
|
idx = random.randint(0, max(0, mid - 1))
|
|
token = most_common_tokens[idx][0]
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) > 0
|
|
for r in res:
|
|
assert token in r[field]
|
|
|
|
# verify inverted index
|
|
if enable_inverted_index:
|
|
if field == "word":
|
|
expr = f"{field} == '{token}'"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
assert r[field] == token
|
|
# query single field for multi-word
|
|
for field in text_fields:
|
|
# match top 10 most common words
|
|
top_10_tokens = []
|
|
for word, count in wf_map[field].most_common(10):
|
|
top_10_tokens.append(word)
|
|
string_of_top_10_words = " ".join(top_10_tokens)
|
|
expr = f"text_match({field}, '{string_of_top_10_words}')"
|
|
log.info(f"expr {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) > 0
|
|
for r in res:
|
|
assert any([token in r[field] for token in top_10_tokens]), (
|
|
f"top 10 tokens {top_10_tokens} not in {r[field]}"
|
|
)
|
|
|
|
# query single field for multi-word
|
|
for field in text_fields:
|
|
# match latest 10 most common english words
|
|
top_10_tokens = []
|
|
for word, count in cf.get_top_english_tokens(wf_map[field], 10):
|
|
top_10_tokens.append(word)
|
|
string_of_top_10_words = " ".join(top_10_tokens)
|
|
expr = f"text_match({field}, '{string_of_top_10_words}')"
|
|
log.info(f"expr {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) > 0
|
|
for r in res:
|
|
assert any([token in r[field] for token in top_10_tokens]), (
|
|
f"top 10 tokens {top_10_tokens} not in {r[field]}"
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_text_match_custom_analyzer_with_stop_words(self):
|
|
"""
|
|
target: test text match with custom analyzer
|
|
method: 1. enable text match, use custom analyzer and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: get the correct token, text match successfully and result is correct
|
|
"""
|
|
stops_words = ["in", "of"]
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
"filter": [
|
|
{
|
|
"type": "stop",
|
|
"stop_words": stops_words,
|
|
}
|
|
],
|
|
}
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"sentence": fake.sentence().lower() + " ".join(stops_words),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["sentence"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
# query single field for one word
|
|
for field in text_fields:
|
|
for token in stops_words:
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) == 0
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_text_match_custom_analyzer_with_lowercase(self):
|
|
"""
|
|
target: test text match with custom analyzer
|
|
method: 1. enable text match, use custom analyzer and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: get the correct token, text match successfully and result is correct
|
|
"""
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
"filter": ["lowercase"],
|
|
}
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"sentence": fake.sentence(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["sentence"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
# query single field for one word
|
|
for field in text_fields:
|
|
tokens = [item[0] for item in wf_map[field].most_common(1)]
|
|
for token in tokens:
|
|
# search with Capital case
|
|
token = token.capitalize()
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
capital_case_res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(capital_case_res)}")
|
|
# search with lower case
|
|
token = token.lower()
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
lower_case_res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(lower_case_res)}")
|
|
|
|
# search with upper case
|
|
token = token.upper()
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
upper_case_res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(upper_case_res)}")
|
|
assert len(capital_case_res) == len(lower_case_res) and len(capital_case_res) == len(upper_case_res)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_text_match_custom_analyzer_with_length_filter(self):
|
|
"""
|
|
target: test text match with custom analyzer
|
|
method: 1. enable text match, use custom analyzer and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: get the correct token, text match successfully and result is correct
|
|
"""
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
"filter": [
|
|
{
|
|
"type": "length", # Specifies the filter type as length
|
|
"max": 10, # Sets the maximum token length to 10 characters
|
|
}
|
|
],
|
|
}
|
|
|
|
long_word = "a" * 11
|
|
max_length_word = "a" * 10
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"sentence": fake.sentence() + " " + long_word + " " + max_length_word,
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["sentence"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
# query sentence field with long word
|
|
for field in text_fields:
|
|
tokens = [long_word]
|
|
for token in tokens:
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
assert len(res) == 0
|
|
# query sentence field with max length word
|
|
for field in text_fields:
|
|
tokens = [max_length_word]
|
|
for token in tokens:
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
assert len(res) == data_size
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_text_match_custom_analyzer_with_stemmer_filter(self):
|
|
"""
|
|
target: test text match with custom analyzer
|
|
method: 1. enable text match, use custom analyzer and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: get the correct token, text match successfully and result is correct
|
|
"""
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
"filter": [
|
|
{
|
|
"type": "stemmer", # Specifies the filter type as stemmer
|
|
"language": "english", # Sets the language for stemming to English
|
|
}
|
|
],
|
|
}
|
|
word_pairs = {
|
|
"play": ["play", "plays", "played", "playing"],
|
|
"book": ["book", "books", "booked", "booking"],
|
|
"study": ["study", "studies", "studied", "studying"],
|
|
}
|
|
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"sentence": fake.sentence() + " " + " ".join(word_pairs.keys()),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["sentence"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
# query sentence field with variant word
|
|
for field in text_fields:
|
|
for stem in word_pairs.keys():
|
|
tokens = word_pairs[stem]
|
|
for token in tokens:
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
pytest.assume(len(res) == data_size, f"stem {stem} token {token} not found in {res}")
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_text_match_custom_analyzer_with_ascii_folding_filter(self):
|
|
"""
|
|
target: test text match with custom analyzer
|
|
method: 1. enable text match, use custom analyzer and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: get the correct token, text match successfully and result is correct
|
|
"""
|
|
from unidecode import unidecode
|
|
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
"filter": ["asciifolding"],
|
|
}
|
|
|
|
origin_texts = [
|
|
"Café Möller serves crème brûlée",
|
|
"José works at Škoda in São Paulo",
|
|
"The œuvre of Łukasz includes æsthetic pieces",
|
|
"München's König Street has günstig prices",
|
|
"El niño está jugando en el jardín",
|
|
"Le système éducatif français",
|
|
]
|
|
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"sentence": fake.sentence() + " " + " ".join(origin_texts),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["sentence"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
# query sentence field with variant word
|
|
for field in text_fields:
|
|
for text in origin_texts:
|
|
ascii_folding_text = unidecode(text)
|
|
expr = f"""text_match({field}, "{ascii_folding_text}")"""
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
pytest.assume(
|
|
len(res) == data_size, f"origin {text} ascii_folding text {ascii_folding_text} not found in {res}"
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_text_match_custom_analyzer_with_decompounder_filter(self):
|
|
"""
|
|
target: test text match with custom analyzer
|
|
method: 1. enable text match, use custom analyzer and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: get the correct token, text match successfully and result is correct
|
|
"""
|
|
word_list = ["dampf", "schiff", "fahrt", "brot", "backen", "automat"]
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
"filter": [
|
|
"lowercase",
|
|
{
|
|
"type": "decompounder", # Specifies the filter type as decompounder
|
|
"word_list": word_list, # Sets the word list for decompounding
|
|
},
|
|
],
|
|
}
|
|
|
|
origin_texts = [
|
|
"Die tägliche Dampfschifffahrt von Hamburg nach Oslo startet um sechs Uhr morgens.",
|
|
"Unser altes Dampfschiff macht eine dreistündige Rundfahrt durch den Hafen.",
|
|
"Der erfahrene Dampfschifffahrtskapitän kennt jede Route auf dem Fluss.",
|
|
"Die internationale Dampfschifffahrtsgesellschaft erweitert ihre Flotte.",
|
|
"Während der Dampfschifffahrt können Sie die Küstenlandschaft bewundern.",
|
|
"Der neue Brotbackautomat produziert stündlich frische Brötchen.",
|
|
"Im Maschinenraum des Dampfschiffs steht ein moderner Brotbackautomat.",
|
|
"Die Brotbackautomatentechnologie wird ständig verbessert.",
|
|
"Unser Brotbackautomat arbeitet mit traditionellen Rezepten.",
|
|
"Der programmierbare Brotbackautomat bietet zwanzig verschiedene Programme.",
|
|
]
|
|
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"sentence": fake.sentence() + " " + " ".join(origin_texts),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["sentence"]
|
|
# query sentence field with word list
|
|
for field in text_fields:
|
|
match_text = " ".join(word_list)
|
|
expr = f"text_match({field}, '{match_text}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
pytest.assume(len(res) == data_size, f"res len {len(res)}, data size {data_size}")
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_text_match_custom_analyzer_with_alphanumonly_filter(self):
|
|
"""
|
|
target: test text match with custom analyzer
|
|
method: 1. enable text match, use custom analyzer and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: get the correct token, text match successfully and result is correct
|
|
"""
|
|
common_non_ascii = [
|
|
"é", # common in words like café, résumé
|
|
"©", # copyright
|
|
"™", # trademark
|
|
"®", # registered trademark
|
|
"°", # degrees, e.g. 20°C
|
|
"€", # euro currency
|
|
"£", # pound sterling
|
|
"±", # plus-minus sign
|
|
"→", # right arrow
|
|
"•", # bullet point
|
|
]
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
"filter": ["alphanumonly"],
|
|
}
|
|
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"sentence": fake.sentence() + " " + " ".join(common_non_ascii),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["sentence"]
|
|
# query sentence field with word list
|
|
for field in text_fields:
|
|
match_text = " ".join(common_non_ascii)
|
|
expr = f"text_match({field}, '{match_text}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
pytest.assume(len(res) == 0, f"res len {len(res)}, data size {data_size}")
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_text_match_custom_analyzer_with_cncharonly_filter(self):
|
|
"""
|
|
target: test text match with custom analyzer
|
|
method: 1. enable text match, use custom analyzer and insert data with varchar
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: get the correct token, text match successfully and result is correct
|
|
"""
|
|
non_zh_char_word_list = ["hello", "milvus", "vector", "database", "19530"]
|
|
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
"filter": ["cncharonly"],
|
|
}
|
|
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"sentence": fake.sentence() + " " + " ".join(non_zh_char_word_list),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["sentence"]
|
|
# query sentence field with word list
|
|
for field in text_fields:
|
|
match_text = " ".join(non_zh_char_word_list)
|
|
expr = f"text_match({field}, '{match_text}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
pytest.assume(len(res) == 0, f"res len {len(res)}, data size {data_size}")
|
|
|
|
@pytest.mark.parametrize("dict_kind", ["ipadic", "ko-dic", "cc-cedict"])
|
|
def test_query_text_match_with_Lindera_tokenizer(self, dict_kind):
|
|
"""
|
|
target: test text match with lindera tokenizer
|
|
method: 1. enable text match, use lindera tokenizer and insert data with varchar in different lang
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: get the correct token, text match successfully and result is correct
|
|
"""
|
|
analyzer_params = {"tokenizer": {"type": "lindera", "dict_kind": dict_kind}}
|
|
if dict_kind == "ipadic":
|
|
fake = fake_jp
|
|
elif dict_kind == "ko-dic":
|
|
fake = KoreanTextGenerator()
|
|
elif dict_kind == "cc-cedict":
|
|
fake = fake_zh
|
|
else:
|
|
fake = fake_en
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"sentence": fake.sentence(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup
|
|
text_fields = ["sentence"]
|
|
# query sentence field with word list
|
|
for field in text_fields:
|
|
match_text = df["sentence"].iloc[0]
|
|
expr = f"text_match({field}, '{match_text}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
assert len(res) > 0
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_text_match_with_combined_expression_for_single_field(self):
|
|
"""
|
|
target: test query text match with combined expression for single field
|
|
method: 1. enable text match, and insert data with varchar
|
|
2. get the most common words and form the combined expression with and operator
|
|
3. verify the result
|
|
expected: query successfully and result is correct
|
|
"""
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
}
|
|
# 1. initialize with data
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"word": fake.word().lower(),
|
|
"sentence": fake.sentence().lower(),
|
|
"paragraph": fake.paragraph().lower(),
|
|
"text": fake.text().lower(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup and get the tf-idf, then base on it to crate expr and ground truth
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
|
|
df_new = cf.split_dataframes(df, fields=text_fields)
|
|
log.info(f"df \n{df}")
|
|
log.info(f"new df \n{df_new}")
|
|
for field in text_fields:
|
|
expr_list = []
|
|
wf_counter = Counter(wf_map[field])
|
|
pd_tmp_res_list = []
|
|
for word, count in wf_counter.most_common(2):
|
|
tmp = f"text_match({field}, '{word}')"
|
|
log.info(f"tmp expr {tmp}")
|
|
expr_list.append(tmp)
|
|
tmp_res = cf.manual_check_text_match(df_new, word, field)
|
|
log.info(f"manual check result for {tmp} {len(tmp_res)}")
|
|
pd_tmp_res_list.append(tmp_res)
|
|
log.info(f"manual res {len(pd_tmp_res_list)}, {pd_tmp_res_list}")
|
|
final_res = set(pd_tmp_res_list[0])
|
|
for i in range(1, len(pd_tmp_res_list)):
|
|
final_res = final_res.intersection(set(pd_tmp_res_list[i]))
|
|
log.info(f"intersection res {len(final_res)}")
|
|
log.info(f"final res {final_res}")
|
|
and_expr = " and ".join(expr_list)
|
|
log.info(f"expr: {and_expr}")
|
|
res, _ = collection_w.query(expr=and_expr, output_fields=text_fields)
|
|
log.info(f"res len {len(res)}, final res {len(final_res)}")
|
|
assert len(res) == len(final_res)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_text_match_with_combined_expression_for_multi_field(self):
|
|
"""
|
|
target: test query text match with combined expression for multi field
|
|
method: 1. enable text match, and insert data with varchar
|
|
2. create the combined expression with `and`, `or` and `not` operator for multi field
|
|
3. verify the result
|
|
expected: query successfully and result is correct
|
|
"""
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
}
|
|
# 1. initialize with data
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"word": fake.word().lower(),
|
|
"sentence": fake.sentence().lower(),
|
|
"paragraph": fake.paragraph().lower(),
|
|
"text": fake.text().lower(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size)
|
|
]
|
|
df = pd.DataFrame(data)
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup and get the tf-idf, then base on it to crate expr and ground truth
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
|
|
df_new = cf.split_dataframes(df, fields=text_fields)
|
|
log.info(f"new df \n{df_new}")
|
|
for i in range(2):
|
|
query, text_match_expr, pandas_expr = cf.generate_random_query_from_freq_dict(
|
|
wf_map, min_freq=3, max_terms=5, p_not=0.2
|
|
)
|
|
log.info(f"expr: {text_match_expr}")
|
|
res, _ = collection_w.query(expr=text_match_expr, output_fields=text_fields)
|
|
onetime_res = res
|
|
log.info(f"res len {len(res)}")
|
|
step_by_step_results = []
|
|
for expr in query:
|
|
if isinstance(expr, dict):
|
|
if "not" in expr:
|
|
key = expr["not"]["field"]
|
|
else:
|
|
key = expr["field"]
|
|
|
|
tmp_expr = cf.generate_text_match_expr(expr)
|
|
res, _ = collection_w.query(expr=tmp_expr, output_fields=text_fields)
|
|
text_match_df = pd.DataFrame(res)
|
|
log.info(f"text match res {len(text_match_df)}\n{text_match_df[key]}")
|
|
log.info(f"tmp expr {tmp_expr} {len(res)}")
|
|
tmp_idx = [r["id"] for r in res]
|
|
step_by_step_results.append(tmp_idx)
|
|
pandas_filter_res = cf.generate_pandas_text_match_result(expr, df_new)
|
|
tmp_pd_idx = pandas_filter_res["id"].tolist()
|
|
diff_id = set(tmp_pd_idx).union(set(tmp_idx)) - set(tmp_pd_idx).intersection(set(tmp_idx))
|
|
log.info(f"diff between text match and manual check {diff_id}")
|
|
assert len(diff_id) == 0
|
|
for idx in diff_id:
|
|
log.info(df[df["id"] == idx][key].values)
|
|
log.info(f"pandas_filter_res {len(pandas_filter_res)} \n {pandas_filter_res}")
|
|
if isinstance(expr, str):
|
|
step_by_step_results.append(expr)
|
|
final_res = cf.evaluate_expression(step_by_step_results)
|
|
log.info(f"one time res {len(onetime_res)}, final res {len(final_res)}")
|
|
if len(onetime_res) != len(final_res):
|
|
log.info("res is not same")
|
|
assert False
|
|
|
|
@pytest.mark.tags(CaseLabel.L2)
|
|
def test_query_text_match_with_multi_lang(self):
|
|
"""
|
|
target: test text match with multi-language text data
|
|
method: 1. enable text match, and insert data with varchar in different language
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: get the correct token, text match successfully and result is correct
|
|
"""
|
|
|
|
# 1. initialize with data
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
}
|
|
# 1. initialize with data
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data_en = [
|
|
{
|
|
"id": i,
|
|
"word": fake.word().lower(),
|
|
"sentence": fake.sentence().lower(),
|
|
"paragraph": fake.paragraph().lower(),
|
|
"text": fake.text().lower(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size // 2)
|
|
]
|
|
fake = fake_de
|
|
data_de = [
|
|
{
|
|
"id": i,
|
|
"word": fake.word().lower(),
|
|
"sentence": fake.sentence().lower(),
|
|
"paragraph": fake.paragraph().lower(),
|
|
"text": fake.text().lower(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size // 2, data_size)
|
|
]
|
|
data = data_en + data_de
|
|
df = pd.DataFrame(data)
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup and get the tf-idf, then base on it to crate expr and ground truth
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
|
|
df_new = cf.split_dataframes(df, fields=text_fields)
|
|
log.info(f"new df \n{df_new}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# query single field for one word
|
|
for field in text_fields:
|
|
token = wf_map[field].most_common()[-1][0]
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) > 0
|
|
for r in res:
|
|
assert token in r[field]
|
|
|
|
# query single field for multi-word
|
|
for field in text_fields:
|
|
# match top 3 most common words
|
|
multi_words = []
|
|
for word, count in wf_map[field].most_common(3):
|
|
multi_words.append(word)
|
|
string_of_multi_words = " ".join(multi_words)
|
|
expr = f"text_match({field}, '{string_of_multi_words}')"
|
|
log.info(f"expr {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) > 0
|
|
for r in res:
|
|
assert any([token in r[field] for token in multi_words])
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_text_match_with_addition_inverted_index(self):
|
|
"""
|
|
target: test text match with addition inverted index
|
|
method: 1. enable text match, and insert data with varchar
|
|
2. create inverted index
|
|
3. get the most common words and query with text match
|
|
4. query with inverted index and verify the result
|
|
expected: get the correct token, text match successfully and result is correct
|
|
"""
|
|
# 1. initialize with data
|
|
fake_en = Faker("en_US")
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
}
|
|
dim = 128
|
|
default_fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
default_schema = CollectionSchema(fields=default_fields, description="test collection")
|
|
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=default_schema)
|
|
data = []
|
|
data_size = 10000
|
|
for i in range(data_size):
|
|
d = {
|
|
"id": i,
|
|
"word": fake_en.word().lower(),
|
|
"sentence": fake_en.sentence().lower(),
|
|
"paragraph": fake_en.paragraph().lower(),
|
|
"text": fake_en.text().lower(),
|
|
"emb": cf.gen_vectors(1, dim)[0],
|
|
}
|
|
data.append(d)
|
|
batch_size = 5000
|
|
for i in range(0, data_size, batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < data_size else data[i:data_size])
|
|
# only if the collection is flushed, the inverted index ca be applied.
|
|
# growing segment may be not applied, although in strong consistency.
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.create_index("word", {"index_type": "INVERTED"})
|
|
collection_w.load()
|
|
df = pd.DataFrame(data)
|
|
df_split = cf.split_dataframes(df, fields=["word", "sentence", "paragraph", "text"])
|
|
log.info(f"dataframe\n{df}")
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language="en")
|
|
# query single field for one word
|
|
for field in text_fields:
|
|
token = wf_map[field].most_common()[-1][0]
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
pandas_res = df_split[df_split.apply(lambda row: token in row[field], axis=1)]
|
|
log.info(f"res len {len(res)}, pandas res len {len(pandas_res)}")
|
|
log.info(f"pandas res\n{pandas_res}")
|
|
assert len(res) == len(pandas_res)
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
assert token in r[field]
|
|
if field == "word":
|
|
assert len(res) == wf_map[field].most_common()[-1][1]
|
|
expr = f"{field} == '{token}'"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) == wf_map[field].most_common()[-1][1]
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
@pytest.mark.parametrize("combine_op", ["and", "or"])
|
|
def test_query_text_match_with_non_varchar_fields_expr(self, combine_op):
|
|
"""
|
|
target: test text match with non-varchar fields expr
|
|
method: 1. enable text match for varchar field and add some non varchar fields
|
|
2. insert data, create index and load
|
|
3. query with text match expr and non-varchar fields expr
|
|
4. verify the result
|
|
expected: query result is correct
|
|
"""
|
|
# 1. initialize with data
|
|
fake_en = Faker("en_US")
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
}
|
|
dim = 128
|
|
default_fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="age",
|
|
dtype=DataType.INT64,
|
|
),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
default_schema = CollectionSchema(fields=default_fields, description="test collection")
|
|
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=default_schema)
|
|
data = []
|
|
data_size = 10000
|
|
for i in range(data_size):
|
|
d = {
|
|
"id": i,
|
|
"age": random.randint(1, 100),
|
|
"word": fake_en.word().lower(),
|
|
"sentence": fake_en.sentence().lower(),
|
|
"paragraph": fake_en.paragraph().lower(),
|
|
"text": fake_en.text().lower(),
|
|
"emb": cf.gen_vectors(1, dim)[0],
|
|
}
|
|
data.append(d)
|
|
batch_size = 5000
|
|
for i in range(0, data_size, batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < data_size else data[i:data_size])
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.create_index("word", {"index_type": "INVERTED"})
|
|
collection_w.load()
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language="en")
|
|
# query single field for one word
|
|
for field in text_fields:
|
|
most_common_tokens = wf_map[field].most_common(10)
|
|
mid = len(most_common_tokens) // 2
|
|
idx = random.randint(0, max(0, mid - 1))
|
|
token = most_common_tokens[idx][0]
|
|
tm_expr = f"text_match({field}, '{token}')"
|
|
int_expr = "age > 10"
|
|
combined_expr = f"{tm_expr} {combine_op} {int_expr}"
|
|
log.info(f"expr: {combined_expr}")
|
|
res, _ = collection_w.query(expr=combined_expr, output_fields=["id", field, "age"])
|
|
log.info(f"res len {len(res)}")
|
|
for r in res:
|
|
if combine_op == "and":
|
|
assert token in r[field] and r["age"] > 10
|
|
if combine_op == "or":
|
|
assert token in r[field] or r["age"] > 10
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_text_match_with_some_empty_string(self):
|
|
"""
|
|
target: test text match normal
|
|
method: 1. enable text match and insert data with varchar with some empty string
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: text match successfully and result is correct
|
|
"""
|
|
# 1. initialize with data
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
}
|
|
# 1. initialize with data
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data_en = [
|
|
{
|
|
"id": i,
|
|
"word": fake.word().lower(),
|
|
"sentence": fake.sentence().lower(),
|
|
"paragraph": fake.paragraph().lower(),
|
|
"text": fake.text().lower(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size // 2)
|
|
]
|
|
data_empty = [
|
|
{
|
|
"id": i,
|
|
"word": "",
|
|
"sentence": " ",
|
|
"paragraph": "",
|
|
"text": " ",
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(data_size // 2, data_size)
|
|
]
|
|
data = data_en + data_empty
|
|
df = pd.DataFrame(data)
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# analyze the croup and get the tf-idf, then base on it to crate expr and ground truth
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
|
|
df_new = cf.split_dataframes(df, fields=text_fields)
|
|
log.info(f"new df \n{df_new}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
# query single field for one word
|
|
for field in text_fields:
|
|
token = wf_map[field].most_common()[-1][0]
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) > 0
|
|
for r in res:
|
|
assert token in r[field]
|
|
# query single field for multi-word
|
|
for field in text_fields:
|
|
# match top 3 most common words
|
|
multi_words = []
|
|
for word, count in wf_map[field].most_common(3):
|
|
multi_words.append(word)
|
|
string_of_multi_words = " ".join(multi_words)
|
|
expr = f"text_match({field}, '{string_of_multi_words}')"
|
|
log.info(f"expr {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=["id", field])
|
|
log.info(f"res len {len(res)}")
|
|
assert len(res) > 0
|
|
for r in res:
|
|
assert any([token in r[field] for token in multi_words])
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_text_match_with_nullable(self):
|
|
"""
|
|
target: test text match with nullable
|
|
method: 1. enable text match and nullable, and insert data with varchar with some None value
|
|
2. get the most common words and query with text match
|
|
3. verify the result
|
|
expected: text match successfully and result is correct
|
|
"""
|
|
# 1. initialize with data
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
}
|
|
# 1. initialize with data
|
|
dim = 128
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="word",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
nullable=True,
|
|
),
|
|
FieldSchema(
|
|
name="sentence",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
nullable=True,
|
|
),
|
|
FieldSchema(
|
|
name="paragraph",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
nullable=True,
|
|
),
|
|
FieldSchema(
|
|
name="text",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
nullable=True,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
schema = CollectionSchema(fields=fields, description="test collection")
|
|
data_size = 5000
|
|
collection_w = self.init_collection_wrap(name=cf.gen_unique_str(prefix), schema=schema)
|
|
fake = fake_en
|
|
language = "en"
|
|
data_null = [
|
|
{
|
|
"id": i,
|
|
"word": None if random.random() < 0.9 else fake.word().lower(),
|
|
"sentence": None if random.random() < 0.9 else fake.sentence().lower(),
|
|
"paragraph": None if random.random() < 0.9 else fake.paragraph().lower(),
|
|
"text": None if random.random() < 0.9 else fake.paragraph().lower(),
|
|
"emb": [random.random() for _ in range(dim)],
|
|
}
|
|
for i in range(0, data_size)
|
|
]
|
|
data = data_null
|
|
df = pd.DataFrame(data)
|
|
log.info(f"dataframe\n{df}")
|
|
batch_size = 5000
|
|
for i in range(0, len(df), batch_size):
|
|
collection_w.insert(data[i : i + batch_size] if i + batch_size < len(df) else data[i : len(df)])
|
|
collection_w.flush()
|
|
collection_w.create_index(
|
|
"emb",
|
|
{"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}},
|
|
)
|
|
collection_w.load()
|
|
text_fields = ["word", "sentence", "paragraph", "text"]
|
|
wf_map = {}
|
|
for field in text_fields:
|
|
wf_map[field] = cf.analyze_documents(df[field].tolist(), language=language)
|
|
# query single field for one word
|
|
for field in text_fields:
|
|
token = wf_map[field].most_common()[-1][0]
|
|
expr = f"text_match({field}, '{token}')"
|
|
log.info(f"expr: {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=text_fields)
|
|
log.info(f"res len {len(res)}, \n{res}")
|
|
assert len(res) > 0
|
|
for r in res:
|
|
assert token in r[field]
|
|
# query single field for multi-word
|
|
for field in text_fields:
|
|
# match top 3 most common words
|
|
multi_words = []
|
|
for word, count in wf_map[field].most_common(3):
|
|
multi_words.append(word)
|
|
string_of_multi_words = " ".join(multi_words)
|
|
expr = f"text_match({field}, '{string_of_multi_words}')"
|
|
log.info(f"expr {expr}")
|
|
res, _ = collection_w.query(expr=expr, output_fields=text_fields)
|
|
log.info(f"res len {len(res)}, {res}")
|
|
assert len(res) > 0
|
|
for r in res:
|
|
assert any([token in r[field] for token in multi_words])
|
|
|
|
|
|
class TestQueryCompareTwoColumn(TestcaseBase):
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_compare_two_column_with_last_expr_skip(self):
|
|
"""
|
|
target: test query compare two column with last expr skip
|
|
method: 1. create collection and insert data
|
|
2. query with two column compare expr with last expr help to skip compare
|
|
expected: query successfully
|
|
"""
|
|
fields = [
|
|
FieldSchema(name="f1", dtype=DataType.INT64, is_primary=True, nullable=False),
|
|
FieldSchema(name="f2", dtype=DataType.SPARSE_FLOAT_VECTOR),
|
|
FieldSchema(name="f3", dtype=DataType.INT64, is_primary=False, nullable=False),
|
|
FieldSchema(name="f4", dtype=DataType.INT64, is_primary=False, nullable=False),
|
|
]
|
|
schema = CollectionSchema(fields=fields)
|
|
col_name = cf.gen_unique_str(prefix)
|
|
collection_w = self.init_collection_wrap(
|
|
name=col_name,
|
|
schema=schema,
|
|
)
|
|
collection_w.create_index("f3", {"index_type": "INVERTED"})
|
|
collection_w.create_index("f2", {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"})
|
|
for i in range(100):
|
|
collection_w.insert({"f1": i, "f2": {10: 0.07638, 3: 0.3925}, "f3": 4, "f4": 2})
|
|
collection_w.flush()
|
|
collection_w.load()
|
|
# f4 / 4 > 1 always false, so f3 < f4 will not be executed
|
|
res = collection_w.query(expr="f4 / 4 > 1 and f3 < f4", output_fields=["count(*)"])[0]
|
|
assert res[0]["count(*)"] == 0
|
|
collection_w.drop()
|
|
|
|
|
|
class TestQueryTextMatchNegative(TestcaseBase):
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
def test_query_text_match_with_unsupported_tokenizer(self):
|
|
"""
|
|
target: test query text match with unsupported tokenizer
|
|
method: 1. enable text match, and use unsupported tokenizer
|
|
2. create collection
|
|
expected: create collection failed and return error
|
|
"""
|
|
analyzer_params = {
|
|
"tokenizer": "Unsupported",
|
|
}
|
|
dim = 128
|
|
default_fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="title",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="overview",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="genres",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="producer",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="cast",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
default_schema = CollectionSchema(fields=default_fields, description="test collection")
|
|
error = {ct.err_code: 2001, ct.err_msg: "unsupported tokenizer"}
|
|
self.init_collection_wrap(
|
|
name=cf.gen_unique_str(prefix),
|
|
schema=default_schema,
|
|
check_task=CheckTasks.err_res,
|
|
check_items=error,
|
|
)
|
|
|
|
|
|
class TestQueryFunction(TestcaseBase):
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_function_calls(self):
|
|
"""
|
|
target: test query data
|
|
method: create collection and insert data
|
|
query with mix call expr in string field and int field
|
|
expected: query successfully
|
|
"""
|
|
collection_w, vectors = self.init_collection_general(
|
|
prefix, insert_data=True, primary_field=ct.default_string_field_name
|
|
)[0:2]
|
|
res = vectors[0].iloc[:, 1:3].to_dict("records")
|
|
output_fields = [default_float_field_name, default_string_field_name]
|
|
for mixed_call_expr in [
|
|
"not empty(varchar) && int64 >= 0",
|
|
# function call is case-insensitive
|
|
"not EmPty(varchar) && int64 >= 0",
|
|
"not EMPTY(varchar) && int64 >= 0",
|
|
"starts_with(varchar, varchar) && int64 >= 0",
|
|
]:
|
|
collection_w.query(
|
|
mixed_call_expr,
|
|
output_fields=output_fields,
|
|
check_task=CheckTasks.check_query_results,
|
|
check_items={exp_res: res, "pk_name": collection_w.primary_field.name},
|
|
)
|
|
|
|
@pytest.mark.tags(CaseLabel.L1)
|
|
def test_query_invalid(self):
|
|
"""
|
|
target: test query with invalid call expression
|
|
method: query with invalid call expr
|
|
expected: raise exception
|
|
"""
|
|
collection_w, entities = self.init_collection_general(prefix, insert_data=True, nb=10)[0:2]
|
|
test_cases = [
|
|
(
|
|
"A_FUNCTION_THAT_DOES_NOT_EXIST()".lower(),
|
|
"function A_FUNCTION_THAT_DOES_NOT_EXIST() not found".lower(),
|
|
),
|
|
# empty
|
|
("empty()", "function empty() not found"),
|
|
(f"empty({default_int_field_name})", "function empty(int64_t) not found"),
|
|
# starts_with
|
|
(f"starts_with({default_int_field_name})", "function starts_with(int64_t) not found"),
|
|
(
|
|
f"starts_with({default_int_field_name}, {default_int_field_name})",
|
|
"function starts_with(int64_t, int64_t) not found",
|
|
),
|
|
]
|
|
for call_expr, err_msg in test_cases:
|
|
error = {ct.err_code: 65535, ct.err_msg: err_msg}
|
|
collection_w.query(call_expr, check_task=CheckTasks.err_res, check_items=error)
|
|
|
|
@pytest.mark.tags(CaseLabel.L0)
|
|
@pytest.mark.xfail(reason="issue 36685")
|
|
def test_query_text_match_with_unsupported_fields(self):
|
|
"""
|
|
target: test enable text match with unsupported field
|
|
method: 1. enable text match in unsupported field
|
|
2. create collection
|
|
expected: create collection failed and return error
|
|
"""
|
|
analyzer_params = {
|
|
"tokenizer": "standard",
|
|
}
|
|
dim = 128
|
|
default_fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
|
|
FieldSchema(
|
|
name="title",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="overview",
|
|
dtype=DataType.VARCHAR,
|
|
max_length=65535,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(
|
|
name="age",
|
|
dtype=DataType.INT64,
|
|
enable_analyzer=True,
|
|
enable_match=True,
|
|
analyzer_params=analyzer_params,
|
|
),
|
|
FieldSchema(name="emb", dtype=DataType.FLOAT_VECTOR, dim=dim),
|
|
]
|
|
default_schema = CollectionSchema(fields=default_fields, description="test collection")
|
|
error = {ct.err_code: 2001, ct.err_msg: "field type is not supported"}
|
|
self.init_collection_wrap(
|
|
name=cf.gen_unique_str(prefix),
|
|
schema=default_schema,
|
|
check_task=CheckTasks.err_res,
|
|
check_items=error,
|
|
)
|