This is an automated email from the ASF dual-hosted git repository. JingsongLi pushed a commit to branch codex/python-tantivy-fulltext-index in repository https://gitbox.apache.org/repos/asf/paimon.git
commit c053895b50e8138b5d4ac4c17fda98e527e74420 Author: JingsongLi <[email protected]> AuthorDate: Wed Jul 1 10:52:26 2026 +0800 [python] Support Tantivy full-text global index build --- .../multimodal-table/global-index/full-text.mdx | 116 ++++++- .../pypaimon/globalindex/create_global_index.py | 70 ++-- .../pypaimon/globalindex/tantivy/__init__.py | 6 + .../tantivy_full_text_global_index_reader.py | 327 +++++++++++++++--- .../tantivy/tantivy_full_text_index_writer.py | 335 ++++++++++++++++++ .../pypaimon/tests/global_index_build_test.py | 380 +++++++++++++++++++++ .../pypaimon/tests/vector_search_filter_test.py | 206 ++++++++++- 7 files changed, 1370 insertions(+), 70 deletions(-) diff --git a/docs/docs/multimodal-table/global-index/full-text.mdx b/docs/docs/multimodal-table/global-index/full-text.mdx index d51129d7c8..6665dcc3f8 100644 --- a/docs/docs/multimodal-table/global-index/full-text.mdx +++ b/docs/docs/multimodal-table/global-index/full-text.mdx @@ -34,6 +34,10 @@ the full-text index and can be consumed directly or combined with vector routes ## Build Full-Text Index +<Tabs groupId="fulltext-build"> + +<TabItem value="sql" label="SQL"> + ```sql -- Create full-text index on 'content' column CALL sys.create_global_index( @@ -77,6 +81,61 @@ CALL sys.create_global_index( ); ``` +</TabItem> + +<TabItem value="python-sdk" label="Python SDK"> + +```python +table = catalog.get_table("db.my_table") + +# Create full-text index on 'content' column. +added_files = table.create_global_index( + "content", + index_type="tantivy-fulltext", +) +print(added_files) +``` + +For other content where users often search by short character fragments, build +the index with Tantivy's `ngram` tokenizer: + +```python +added_files = table.create_global_index( + "content", + index_type="tantivy-fulltext", + options={ + "tantivy.tokenizer": "ngram", + "tantivy.ngram.min-gram": "2", + "tantivy.ngram.max-gram": "2", + }, +) +print(added_files) +``` + +To customize text analysis, choose a base tokenizer and compose token filters: + +```python +added_files = table.create_global_index( + "content", + index_type="tantivy-fulltext", + options={ + "tantivy.tokenizer": "simple", + "tantivy.stem": "true", + "tantivy.remove-stop-words": "true", + }, +) +print(added_files) +``` + +PyPaimon does not build indexes with `tantivy.tokenizer=jieba` because +tantivy-py does not expose the same `tantivy_jieba` tokenizer used by the Java +implementation. Use SQL or Java to build `jieba` full-text indexes, then +PyPaimon can query them when the Python `jieba` package is installed. + +</TabItem> + +</Tabs> + Supported tokenizer options: | Option | Default | Description | @@ -98,7 +157,9 @@ Tokenizer settings are persisted in global index metadata. Existing index files tokenizer they were built with, even if later index builds use different options. Paimon does not load arbitrary Rust tokenizer plugins from configuration; custom analysis is provided by composing the supported tokenizer and filter options above. PyPaimon can query -`jieba` indexes when the Python `jieba` package is installed. +`jieba` indexes when the Python `jieba` package is installed. PyPaimon full-text index +builds require the Python `tantivy` package. PyPaimon does not build +`tantivy.tokenizer=jieba` indexes; use Java or SQL for that tokenizer. Choose tokenizer settings based on the query pattern: @@ -106,11 +167,56 @@ Choose tokenizer settings based on the query pattern: |---|---| | Natural-language English text | Use the default tokenizer, or enable `tantivy.stem=true` and `tantivy.remove-stop-words=true`. | | Short fragments or substring-like lookup | Use `tantivy.tokenizer=ngram` and tune `tantivy.ngram.min-gram` / `tantivy.ngram.max-gram`. | -| Chinese text | Use `tantivy.tokenizer=jieba`. | +| Chinese text | Use SQL or Java with `tantivy.tokenizer=jieba`; PyPaimon can query these indexes but does not build them. | | Exact token matching | Use `tantivy.tokenizer=raw` or `tantivy.tokenizer=whitespace`, depending on how the field is written. | Set `tantivy.with-position=false` to reduce index size when phrase queries are not needed. +## Drop Full-Text Index + +<Tabs groupId="fulltext-drop"> + +<TabItem value="sql" label="SQL"> + +```sql +CALL sys.drop_global_index( + table => 'db.my_table', + index_column => 'content', + index_type => 'tantivy-fulltext' +); +``` + +</TabItem> + +<TabItem value="python-sdk" label="Python SDK"> + +```python +table = catalog.get_table("db.my_table") + +dropped_files = table.drop_global_index( + "content", + index_type="tantivy-fulltext", +) +print(dropped_files) +``` + +You can also restrict the drop to selected partitions, or count matched files +without committing: + +```python +matched_files = table.drop_global_index( + "content", + index_type="tantivy-fulltext", + partitions=[{"dt": "2026-06-18"}, {"dt": "2026-06-19"}], + dry_run=True, +) +print(matched_files) +``` + +</TabItem> + +</Tabs> + ## Full-Text Search Full-text search accepts a JSON query DSL. The root JSON object should contain one query type. @@ -406,7 +512,9 @@ Supported `occur` values are `should`, `must`, and `must_not`. - `fuzziness` must be `0`, `1`, `2`, `null`, or `auto`. - `max_expansions` and `prefix_length` are part of the DSL for LanceDB compatibility. The current Tantivy backend accepts the default values only: `max_expansions=50` and `prefix_length=0`. -- PyPaimon can parse the same DSL, but its local Tantivy reader does not support every advanced - scoring option yet. Unsupported options fail fast instead of changing query semantics silently. +- PyPaimon parses the same DSL and fails fast for unsupported Tantivy options instead of + changing query semantics silently. +- PyPaimon can query Java/SQL-built `jieba` indexes for `match` queries. `match_phrase` + on `jieba` indexes is not supported by the local Python reader. - Multi-column queries such as `multi_match` require full-text indexes for every referenced column. diff --git a/paimon-python/pypaimon/globalindex/create_global_index.py b/paimon-python/pypaimon/globalindex/create_global_index.py index db184ddad9..6fe5ae82e3 100644 --- a/paimon-python/pypaimon/globalindex/create_global_index.py +++ b/paimon-python/pypaimon/globalindex/create_global_index.py @@ -35,6 +35,12 @@ from pypaimon.globalindex.bitmap.bitmap_index_writer import ( ) from pypaimon.globalindex.global_index_meta import GlobalIndexMeta from pypaimon.globalindex.key_serializer import create_serializer +from pypaimon.globalindex.tantivy.tantivy_full_text_global_index_reader import ( + TANTIVY_FULLTEXT_IDENTIFIER, +) +from pypaimon.globalindex.tantivy.tantivy_full_text_index_writer import ( + TantivyFullTextIndexWriter, +) from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import ( VINDEX_IDENTIFIERS, ) @@ -86,6 +92,9 @@ def create_global_index( _SORTED_INDEX_IDENTIFIERS = (BTREE_IDENTIFIER, BITMAP_IDENTIFIER) +_GENERIC_INDEX_IDENTIFIERS = tuple(VINDEX_IDENTIFIERS) + ( + TANTIVY_FULLTEXT_IDENTIFIER, +) _SORTED_INDEX_RECORDS_PER_RANGE_FLOATING = 1.2 @@ -111,11 +120,15 @@ class GlobalIndexBuilder: if ( self._index_type not in _SORTED_INDEX_IDENTIFIERS - and self._index_type not in VINDEX_IDENTIFIERS + and self._index_type not in _GENERIC_INDEX_IDENTIFIERS ): raise ValueError( "Python global index build currently supports %s and %s, got '%s'." - % (_SORTED_INDEX_IDENTIFIERS, VINDEX_IDENTIFIERS, index_type) + % ( + _SORTED_INDEX_IDENTIFIERS, + _GENERIC_INDEX_IDENTIFIERS, + index_type, + ) ) if len(self._index_columns) != 1: raise ValueError( @@ -133,8 +146,8 @@ class GlobalIndexBuilder: "Column '%s' does not exist in table '%s'." % (column, self._table.identifier) ) - if self._index_type in VINDEX_IDENTIFIERS: - self._validate_vindex_table() + if self._index_type in _GENERIC_INDEX_IDENTIFIERS: + self._validate_generic_index_table() def build(self) -> List[CommitMessage]: read_builder = self._table.new_read_builder() @@ -148,7 +161,7 @@ class GlobalIndexBuilder: return [] index_field = self._table.field_dict[self._index_columns[0]] - if self._index_type in VINDEX_IDENTIFIERS: + if self._index_type in _GENERIC_INDEX_IDENTIFIERS: splits = _filter_non_indexable_splits( self._table, splits, self._index_columns) if not splits: @@ -168,7 +181,8 @@ class GlobalIndexBuilder: if self._index_type in _SORTED_INDEX_IDENTIFIERS: return self._build_sorted_index( splits, index_field, table_read, index_path) - return self._build_vindex(splits, index_field, table_read, index_path) + return self._build_generic_index( + splits, index_field, table_read, index_path) def _build_sorted_index( self, splits, index_field, table_read, index_path: str @@ -244,7 +258,7 @@ class GlobalIndexBuilder: ) raise ValueError("Unsupported sorted global index type: %s" % self._index_type) - def _build_vindex( + def _build_generic_index( self, splits, index_field, table_read, index_path: str ) -> List[CommitMessage]: rows_per_shard = self._core_options.global_index_row_count_per_shard() @@ -261,22 +275,15 @@ class GlobalIndexBuilder: if table is None or table.num_rows == 0: continue - writer = VindexVectorIndexWriter( - self._table.file_io, - index_path, - index_field.type, - self._index_type, - self._options.to_map(), - index_field.name, - ) + writer = self._create_generic_index_writer(index_path, index_field) try: - for vector, row_id in _extract_vector_rows( + for value, row_id in _extract_index_rows( table, self._index_columns[0], SpecialFields.ROW_ID.name, row_range, ): - writer.write(vector, row_id - row_range.from_) + writer.write(value, row_id - row_range.from_) index_adds = _to_index_manifest_entries( self._table, @@ -299,7 +306,26 @@ class GlobalIndexBuilder: ) return messages - def _validate_vindex_table(self) -> None: + def _create_generic_index_writer(self, index_path: str, index_field): + if self._index_type in VINDEX_IDENTIFIERS: + return VindexVectorIndexWriter( + self._table.file_io, + index_path, + index_field.type, + self._index_type, + self._options.to_map(), + index_field.name, + ) + if self._index_type == TANTIVY_FULLTEXT_IDENTIFIER: + return TantivyFullTextIndexWriter( + self._table.file_io, + index_path, + index_field.type, + self._options.to_map(), + ) + raise ValueError("Unsupported generic global index type: %s" % self._index_type) + + def _validate_generic_index_table(self) -> None: bucket = self._core_options.bucket() if bucket != -1: raise ValueError( @@ -607,22 +633,22 @@ def _extract_sorted_rows( return sorted(rows, key=cmp_to_key(compare)) -def _extract_vector_rows( +def _extract_index_rows( table: pa.Table, index_column: str, row_id_column: str, row_range: Optional[Range] = None, ): - vectors = table.column(index_column).to_pylist() + values = table.column(index_column).to_pylist() row_ids = table.column(row_id_column).to_pylist() rows = [] - for vector, row_id in zip(vectors, row_ids): + for value, row_id in zip(values, row_ids): if row_id is None: raise ValueError("Cannot build global index because _ROW_ID is null.") row_id = int(row_id) if row_range is not None and not row_range.contains(row_id): continue - rows.append((vector, row_id)) + rows.append((value, row_id)) return rows diff --git a/paimon-python/pypaimon/globalindex/tantivy/__init__.py b/paimon-python/pypaimon/globalindex/tantivy/__init__.py index 1c7579cb78..d972022d3e 100644 --- a/paimon-python/pypaimon/globalindex/tantivy/__init__.py +++ b/paimon-python/pypaimon/globalindex/tantivy/__init__.py @@ -17,10 +17,16 @@ from pypaimon.globalindex.tantivy.tantivy_full_text_global_index_reader import ( TantivyFullTextGlobalIndexReader, + TantivyFullTextIndexOptions, TANTIVY_FULLTEXT_IDENTIFIER, ) +from pypaimon.globalindex.tantivy.tantivy_full_text_index_writer import ( + TantivyFullTextIndexWriter, +) __all__ = [ 'TantivyFullTextGlobalIndexReader', + 'TantivyFullTextIndexOptions', + 'TantivyFullTextIndexWriter', 'TANTIVY_FULLTEXT_IDENTIFIER', ] diff --git a/paimon-python/pypaimon/globalindex/tantivy/tantivy_full_text_global_index_reader.py b/paimon-python/pypaimon/globalindex/tantivy/tantivy_full_text_global_index_reader.py index 2da9ebe3bd..425c6e927e 100644 --- a/paimon-python/pypaimon/globalindex/tantivy/tantivy_full_text_global_index_reader.py +++ b/paimon-python/pypaimon/globalindex/tantivy/tantivy_full_text_global_index_reader.py @@ -27,7 +27,7 @@ import os import struct import threading from dataclasses import dataclass -from typing import Dict, List +from typing import Dict, List, Mapping from pypaimon.globalindex.global_index_reader import GlobalIndexReader, FieldRef, _completed_future from pypaimon.globalindex.vector_search_result import ( @@ -61,6 +61,21 @@ _SUPPORTED_LANGUAGES = { "tamil", "turkish", } +_TANTIVY_OPTION_PREFIX = "tantivy." +_TANTIVY_OPTION_KEYS = { + "tokenizer", + "ngram.min-gram", + "ngram.max-gram", + "ngram.prefix-only", + "lower-case", + "max-token-length", + "ascii-folding", + "stem", + "language", + "remove-stop-words", + "stop-words", + "with-position", +} @dataclass(frozen=True) @@ -86,27 +101,81 @@ class TantivyFullTextIndexOptions: return TantivyFullTextIndexOptions() return TantivyFullTextIndexOptions._deserialize_json(data) + @staticmethod + def from_options(options: Mapping[str, object]): + """Create options from table/global-index build options. + + Java's Tantivy factory strips the ``tantivy.`` prefix before creating + TantivyFullTextIndexOptions, while users configure PyPaimon with the + public prefixed keys. Accept both forms to keep metadata compatible. + """ + + config = {} + for key, value in options.items(): + key = str(key) + if key.startswith(_TANTIVY_OPTION_PREFIX): + key = key[len(_TANTIVY_OPTION_PREFIX):] + if key in _TANTIVY_OPTION_KEYS: + config[key] = value + return TantivyFullTextIndexOptions._from_config(config) + @staticmethod def _deserialize_json(data): config = json.loads(data.decode("utf-8")) + return TantivyFullTextIndexOptions._from_config(config) + + @staticmethod + def _from_config(config): stop_words = config.get("stop-words", []) if isinstance(stop_words, list): stop_words = ";".join( - word for word in stop_words if word is not None) + str(word) for word in stop_words if word is not None) return TantivyFullTextIndexOptions( - tokenizer=config.get("tokenizer", "default"), - ngram_min_gram=config.get("ngram.min-gram", 2), - ngram_max_gram=config.get("ngram.max-gram", 2), - ngram_prefix_only=config.get("ngram.prefix-only", False), - lower_case=config.get("lower-case", True), - max_token_length=config.get("max-token-length", 40), - ascii_folding=config.get("ascii-folding", False), - stem=config.get("stem", False), - language=config.get("language", "english"), - remove_stop_words=config.get("remove-stop-words", False), + tokenizer=_get_string(config, "tokenizer", "default"), + ngram_min_gram=_get_int(config, "ngram.min-gram", 2), + ngram_max_gram=_get_int(config, "ngram.max-gram", 2), + ngram_prefix_only=_get_bool(config, "ngram.prefix-only", False), + lower_case=_get_bool(config, "lower-case", True), + max_token_length=_get_int(config, "max-token-length", 40), + ascii_folding=_get_bool(config, "ascii-folding", False), + stem=_get_bool(config, "stem", False), + language=_get_string(config, "language", "english"), + remove_stop_words=_get_bool(config, "remove-stop-words", False), stop_words=stop_words, - with_position=config.get("with-position", True)) + with_position=_get_bool(config, "with-position", True)) + + def serialize(self) -> bytes: + return self.to_native_config_json().encode("utf-8") + + def to_native_config_json(self) -> str: + config = {} + if self.tokenizer != "default": + config["tokenizer"] = self.tokenizer + if self.ngram_min_gram != 2: + config["ngram.min-gram"] = self.ngram_min_gram + if self.ngram_max_gram != 2: + config["ngram.max-gram"] = self.ngram_max_gram + if self.ngram_prefix_only: + config["ngram.prefix-only"] = self.ngram_prefix_only + if not self.lower_case: + config["lower-case"] = self.lower_case + if self.max_token_length != 40: + config["max-token-length"] = self.max_token_length + if self.ascii_folding: + config["ascii-folding"] = self.ascii_folding + if self.stem: + config["stem"] = self.stem + if self.language != "english": + config["language"] = self.language + if self.remove_stop_words: + config["remove-stop-words"] = self.remove_stop_words + stop_words = self.stop_word_list() + if stop_words: + config["stop-words"] = stop_words + if not self.with_position: + config["with-position"] = self.with_position + return json.dumps(config, separators=(",", ":")) def __post_init__(self): tokenizer = "" if self.tokenizer is None else self.tokenizer.strip().lower() @@ -275,24 +344,88 @@ class TantivyFullTextGlobalIndexReader(GlobalIndexReader): limit = full_text_search.limit - searcher = self._searcher import tantivy - query = self._parse_query(tantivy, full_text_search) + id_to_scores = self._search_full_text_query( + tantivy, full_text_search.query, limit) + return _completed_future( + DictBasedScoredIndexResult(id_to_scores).top_k(limit)) + + def _search_full_text_query(self, tantivy, query, limit): + from pypaimon.globalindex.full_text_query import ( + BooleanQuery, + BoostQuery, + MultiMatchQuery, + ) + + if isinstance(query, BoostQuery): + child_limit = self._child_query_limit(limit) + positive = self._search_full_text_query( + tantivy, query.positive, child_limit) + negative = self._search_full_text_query( + tantivy, query.negative, child_limit) + return _demote_scores(positive, negative, query.negative_boost) + if isinstance(query, BooleanQuery) and _contains_boost_query(query): + return self._search_boolean_query( + tantivy, query, self._child_query_limit(limit)) + if isinstance(query, MultiMatchQuery): + raise ValueError( + "multi_match is not supported by single-column Tantivy full-text indexes" + ) + + return self._search_tantivy_query( + self._parse_structured_query(tantivy, query), limit) + + def _child_query_limit(self, limit): + num_docs = getattr(self._searcher, "num_docs", None) + if num_docs is None: + return limit + return max(limit, int(num_docs)) + + def _search_boolean_query(self, tantivy, query, limit): + from pypaimon.globalindex.full_text_query import Occur + + result = None + should_results = [] + must_not_results = [] + for occur, child in query.queries: + child_scores = self._search_full_text_query(tantivy, child, limit) + if occur == Occur.MUST: + result = ( + child_scores if result is None + else _intersect_scores(result, child_scores) + ) + elif occur == Occur.SHOULD: + should_results.append(child_scores) + else: + must_not_results.append(child_scores) + + if should_results: + should_result = _union_scores(should_results) + result = ( + should_result if result is None + else _and_with_bonus_scores(result, should_result) + ) + if result is None: + return {} + for must_not in must_not_results: + result = _remove_scores(result, must_not) + return result + + def _search_tantivy_query(self, query, limit): + results = self._searcher.search(query, limit) - results = searcher.search(query, limit) if not results.hits: - return _completed_future(DictBasedScoredIndexResult({})) + return {} doc_addresses = [addr for score, addr in results.hits] scores = [score for score, addr in results.hits] - row_ids = searcher.fast_field_values("row_id", doc_addresses) + row_ids = self._searcher.fast_field_values("row_id", doc_addresses) id_to_scores: Dict[int, float] = {} for row_id, score in zip(row_ids, scores): id_to_scores[row_id] = score - - return _completed_future(DictBasedScoredIndexResult(id_to_scores)) + return id_to_scores def _ensure_loaded(self): if self._searcher is not None: @@ -444,7 +577,7 @@ class TantivyFullTextGlobalIndexReader(GlobalIndexReader): return tantivy.Query.boolean_query(subqueries) if isinstance(query, BoostQuery): raise ValueError( - "boost query is not supported by PyPaimon Tantivy full-text reader" + "boost query must be evaluated through search, not parsed as a Tantivy query" ) if isinstance(query, MultiMatchQuery): raise ValueError( @@ -453,44 +586,86 @@ class TantivyFullTextGlobalIndexReader(GlobalIndexReader): raise ValueError("Unsupported full-text query type: %s" % type(query).__name__) def _parse_match_query(self, tantivy, query): - if query.boost != 1.0: - raise ValueError( - "match query boost is not supported by PyPaimon Tantivy full-text reader" - ) - if query.fuzziness not in (None, 0): - raise ValueError( - "match query fuzziness is not supported by PyPaimon Tantivy full-text reader" - ) if query.max_expansions != 50: raise ValueError( - "match query max_expansions is not supported by PyPaimon Tantivy full-text reader" + "match query max_expansions is not supported by Tantivy 0.22" ) if query.prefix_length != 0: raise ValueError( - "match query prefix_length is not supported by PyPaimon Tantivy full-text reader" + "match query prefix_length is not supported by Tantivy 0.22" ) conjunction_by_default = query.operator.value == "AND" if self._index_options.tokenizer != "jieba": + parse_kwargs = {} if conjunction_by_default: - return self._index.parse_query( - query.query, ["text"], conjunction_by_default=True) - return self._index.parse_query(query.query, ["text"]) + parse_kwargs["conjunction_by_default"] = True + if query.fuzziness not in (None, 0): + parse_kwargs["fuzzy_fields"] = { + "text": (False, int(query.fuzziness), True) + } + parsed = self._parse_index_query( + tantivy, query.query, parse_kwargs) + return self._boost_query(tantivy, parsed, query.boost) tokens = self._jieba_query_tokens(query.query) if not tokens: return tantivy.Query.empty_query() - term_queries = [ - tantivy.Query.term_query(self._schema, "text", token) - for token in tokens - ] + term_queries = [] + for token in tokens: + if query.fuzziness not in (None, 0): + fuzzy_query = getattr(tantivy.Query, "fuzzy_term_query", None) + if fuzzy_query is None: + raise RuntimeError( + "PyPaimon Tantivy full-text search requires a tantivy-py " + "version with Query.fuzzy_term_query support for " + "match query fuzziness." + ) + term_queries.append( + fuzzy_query( + self._schema, + "text", + token, + distance=int(query.fuzziness), + transposition_cost_one=True, + prefix=False, + ) + ) + else: + term_queries.append( + tantivy.Query.term_query(self._schema, "text", token)) if len(term_queries) == 1: - return term_queries[0] + return self._boost_query(tantivy, term_queries[0], query.boost) occur = tantivy.Occur.Must if conjunction_by_default else tantivy.Occur.Should - return tantivy.Query.boolean_query([ + parsed = tantivy.Query.boolean_query([ (occur, query) for query in term_queries ]) + return self._boost_query(tantivy, parsed, query.boost) + + def _parse_index_query(self, tantivy, query_text, parse_kwargs): + try: + return self._index.parse_query(query_text, ["text"], **parse_kwargs) + except TypeError as e: + if "fuzzy_fields" in parse_kwargs: + raise RuntimeError( + "PyPaimon Tantivy full-text search requires a tantivy-py " + "version with Index.parse_query fuzzy_fields support for " + "match query fuzziness." + ) from e + raise + + @staticmethod + def _boost_query(tantivy, parsed_query, boost): + if boost == 1.0: + return parsed_query + boost_query = getattr(tantivy.Query, "boost_query", None) + if boost_query is None: + raise RuntimeError( + "PyPaimon Tantivy full-text search requires a tantivy-py " + "version with Query.boost_query support for match query boost." + ) + return boost_query(parsed_query, float(boost)) def _jieba_query_tokens(self, query_text): try: @@ -685,17 +860,85 @@ def _read_fully(stream, length: int) -> bytes: return bytes(buf) +def _get_string(config, key, default): + value = config.get(key) + return default if value is None else str(value) + + +def _get_int(config, key, default): + value = config.get(key) + if value is None: + return default + return int(value) + + +def _get_bool(config, key, default): + value = config.get(key) + if value is None: + return default + if isinstance(value, bool): + return value + return str(value).strip().lower() == "true" + + def _normalize_stop_words(stop_words): if stop_words is None: return "" if isinstance(stop_words, list): return ";".join( - word.strip() + str(word).strip() for word in stop_words - if word is not None and word.strip() + if word is not None and str(word).strip() ) return ";".join( word.strip() for word in stop_words.split(";") if word.strip() ) + + +def _contains_boost_query(query): + from pypaimon.globalindex.full_text_query import BooleanQuery, BoostQuery + + if isinstance(query, BoostQuery): + return True + if isinstance(query, BooleanQuery): + return any(_contains_boost_query(child) for _, child in query.queries) + return False + + +def _demote_scores(positive, negative, negative_boost): + return { + row_id: score * (negative_boost if row_id in negative else 1.0) + for row_id, score in positive.items() + } + + +def _intersect_scores(left, right): + return { + row_id: left[row_id] + right[row_id] + for row_id in left.keys() & right.keys() + } + + +def _and_with_bonus_scores(base, bonus): + return { + row_id: score + bonus.get(row_id, 0.0) + for row_id, score in base.items() + } + + +def _union_scores(results): + merged = {} + for result in results: + for row_id, score in result.items(): + merged[row_id] = merged.get(row_id, 0.0) + score + return merged + + +def _remove_scores(left, right): + return { + row_id: score + for row_id, score in left.items() + if row_id not in right + } diff --git a/paimon-python/pypaimon/globalindex/tantivy/tantivy_full_text_index_writer.py b/paimon-python/pypaimon/globalindex/tantivy/tantivy_full_text_index_writer.py new file mode 100644 index 0000000000..f338dc59df --- /dev/null +++ b/paimon-python/pypaimon/globalindex/tantivy/tantivy_full_text_index_writer.py @@ -0,0 +1,335 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Tantivy full-text global index writer compatible with Java's layout.""" + +import os +import shutil +import struct +import tempfile +from typing import List, Mapping + +from pypaimon.globalindex.index_file_utils import new_global_index_file_name +from pypaimon.globalindex.result_entry import ResultEntry +from pypaimon.globalindex.tantivy.tantivy_full_text_global_index_reader import ( + TantivyFullTextIndexOptions, +) +from pypaimon.schema.data_types import AtomicType, DataType + + +_FILE_NAME_PREFIX = "tantivy" + + +class TantivyFullTextIndexWriter: + """Writer for one Tantivy full-text global index file. + + Row IDs are local to the manifest row range, matching the Java generic + global index writer contract. Null text values increase row_count but are + not added to the Tantivy index. + """ + + def __init__( + self, + file_io, + index_path: str, + data_type: DataType, + options: Mapping[str, object], + ): + validate_text_type(data_type) + self.file_name = new_global_index_file_name(_FILE_NAME_PREFIX) + self._file_io = file_io + self._index_path = index_path.rstrip("/") + self._index_options = TantivyFullTextIndexOptions.from_options(options) + if self._index_options.tokenizer == "jieba": + raise ValueError( + "PyPaimon Tantivy full-text index build does not support " + "tantivy.tokenizer=jieba because tantivy-py does not expose " + "the Java/Rust tantivy_jieba tokenizer. Build jieba full-text " + "indexes with Java or SQL." + ) + self._row_count = 0 + self._temp_index_dir = None + self._tantivy = None + self._index = None + self._writer = None + self._schema = None + self._closed = False + + @property + def index_options(self) -> TantivyFullTextIndexOptions: + return self._index_options + + def write(self, text, relative_row_id: int) -> None: + if self._closed: + raise RuntimeError("TantivyFullTextIndexWriter is already closed.") + + self._row_count += 1 + if text is None: + return + + text = _materialize_text(text) + self._ensure_writer() + self._add_document(int(relative_row_id), text) + + def finish(self) -> List[ResultEntry]: + if self._closed: + raise RuntimeError("TantivyFullTextIndexWriter is already closed.") + self._closed = True + + file_path = self._file_path() + try: + if self._row_count == 0: + return [] + + self._ensure_writer() + self._writer.commit() + if hasattr(self._writer, "wait_merging_threads"): + self._writer.wait_merging_threads() + + self._file_io.check_or_mkdirs(self._index_path) + with self._file_io.new_output_stream(file_path) as output_stream: + _pack_index_directory(self._temp_index_dir, output_stream) + except Exception: + self._file_io.delete_quietly(file_path) + raise + finally: + self._writer = None + self._index = None + self._schema = None + self._delete_temp_dir() + + return [ + ResultEntry( + self.file_name, + self._row_count, + self._index_options.serialize(), + ) + ] + + def close(self) -> None: + if not self._closed: + self._closed = True + self._writer = None + self._index = None + self._schema = None + self._delete_temp_dir() + + def _file_path(self) -> str: + return "%s/%s" % (self._index_path, self.file_name) + + def _ensure_writer(self) -> None: + if self._writer is not None: + return + + try: + import tantivy + except ImportError as e: + raise ImportError( + "tantivy is required to build Tantivy full-text indexes. " + "Install tantivy or pypaimon with the appropriate full-text " + "search dependencies." + ) from e + + self._verify_tantivy_writer_api(tantivy) + self._tantivy = tantivy + self._temp_index_dir = tempfile.mkdtemp(prefix="tantivy-index-") + self._schema = _build_schema(tantivy, self._index_options) + self._index = _create_index(tantivy, self._schema, self._temp_index_dir) + _register_tokenizer(tantivy, self._index, self._index_options) + self._writer = self._index.writer() + + def _add_document(self, row_id: int, text: str) -> None: + document = self._tantivy.Document() + document.add_unsigned("row_id", row_id) + document.add_text("text", text) + self._writer.add_document(document) + + def _delete_temp_dir(self) -> None: + if self._temp_index_dir is not None: + shutil.rmtree(self._temp_index_dir, ignore_errors=True) + self._temp_index_dir = None + + def _verify_tantivy_writer_api(self, tantivy) -> None: + missing = [] + for name in ("Document", "Index", "SchemaBuilder"): + if not hasattr(tantivy, name): + missing.append(name) + + if self._index_options.tokenizer_name() != "default": + for name in ("TextAnalyzerBuilder", "Tokenizer"): + if not hasattr(tantivy, name): + missing.append(name) + tokenizer = getattr(tantivy, "Tokenizer", None) + tokenizer_apis = { + "default": "simple", + "ngram": "ngram", + "simple": "simple", + "whitespace": "whitespace", + "raw": "raw", + } + tokenizer_api = tokenizer_apis.get(self._index_options.tokenizer) + if ( + tokenizer_api is not None + and tokenizer is not None + and not hasattr(tokenizer, tokenizer_api) + ): + missing.append("Tokenizer.%s" % tokenizer_api) + + filter_checks = [] + if self._index_options.tokenizer_name() != "default": + filter_checks.append(("remove_long", "Filter.remove_long")) + if self._index_options.lower_case: + filter_checks.append(("lowercase", "Filter.lowercase")) + if self._index_options.ascii_folding: + filter_checks.append(("ascii_fold", "Filter.ascii_fold")) + if self._index_options.stem: + filter_checks.append(("stemmer", "Filter.stemmer")) + if self._index_options.remove_stop_words: + filter_checks.append(("stopword", "Filter.stopword")) + if self._index_options.stop_word_list(): + filter_checks.append(("custom_stopword", "Filter.custom_stopword")) + if self._index_options.tokenizer_name() != "default" and filter_checks: + filter_ = getattr(tantivy, "Filter", None) + if filter_ is None: + missing.append("Filter") + else: + for attr, api_name in filter_checks: + if not hasattr(filter_, attr): + missing.append(api_name) + if missing: + raise RuntimeError( + "PyPaimon Tantivy full-text index build requires a tantivy-py " + "version with writer support. Missing API(s): %s" + % ", ".join(missing) + ) + + +def validate_text_type(data_type: DataType) -> None: + if not isinstance(data_type, AtomicType): + raise ValueError( + "Tantivy full-text index requires string type, but got: %s" + % data_type + ) + + type_name = data_type.type.upper() + if ( + type_name == "STRING" + or type_name.startswith("CHAR") + or type_name.startswith("VARCHAR") + ): + return + raise ValueError( + "Tantivy full-text index requires string type, but got: %s" + % data_type + ) + + +def _create_index(tantivy, schema, path): + try: + return tantivy.Index(schema, path=path) + except TypeError: + return tantivy.Index(schema, path) + + +def _build_schema(tantivy, index_options: TantivyFullTextIndexOptions): + schema_builder = tantivy.SchemaBuilder() + schema_builder.add_unsigned_field( + "row_id", stored=False, indexed=True, fast=True, + ) + tokenizer_name = index_options.tokenizer_name() + field_kwargs = {} + if not index_options.with_position: + field_kwargs["index_option"] = "freq" + if tokenizer_name == "default": + schema_builder.add_text_field( + "text", stored=False, **field_kwargs, + ) + else: + schema_builder.add_text_field( + "text", stored=False, + tokenizer_name=tokenizer_name, **field_kwargs, + ) + return schema_builder.build() + + +def _register_tokenizer(tantivy, index, index_options: TantivyFullTextIndexOptions): + if ( + index_options.tokenizer == "default" + and index_options.tokenizer_name() == "default" + ): + return + + if index_options.tokenizer == "ngram": + tokenizer = tantivy.Tokenizer.ngram( + min_gram=index_options.ngram_min_gram, + max_gram=index_options.ngram_max_gram, + prefix_only=index_options.ngram_prefix_only) + elif index_options.tokenizer in ("default", "simple"): + tokenizer = tantivy.Tokenizer.simple() + elif index_options.tokenizer == "whitespace": + tokenizer = tantivy.Tokenizer.whitespace() + elif index_options.tokenizer == "raw": + tokenizer = tantivy.Tokenizer.raw() + else: + raise ValueError("Unsupported Tantivy tokenizer: %s" % index_options.tokenizer) + + analyzer_builder = tantivy.TextAnalyzerBuilder(tokenizer) + analyzer_builder = analyzer_builder.filter( + tantivy.Filter.remove_long(index_options.max_token_length)) + if index_options.lower_case: + analyzer_builder = analyzer_builder.filter(tantivy.Filter.lowercase()) + if index_options.ascii_folding: + analyzer_builder = analyzer_builder.filter(tantivy.Filter.ascii_fold()) + if index_options.stem: + analyzer_builder = analyzer_builder.filter( + tantivy.Filter.stemmer(index_options.language)) + if index_options.remove_stop_words: + analyzer_builder = analyzer_builder.filter( + tantivy.Filter.stopword(index_options.language)) + stop_words = index_options.stop_word_list() + if stop_words: + analyzer_builder = analyzer_builder.filter( + tantivy.Filter.custom_stopword(stop_words)) + index.register_tokenizer(index_options.tokenizer_name(), analyzer_builder.build()) + + +def _materialize_text(value) -> str: + if hasattr(value, "as_py"): + value = value.as_py() + if not isinstance(value, str): + raise ValueError("Unsupported field type: %s" % type(value).__name__) + return value + + +def _pack_index_directory(directory: str, output_stream) -> None: + files = [ + path for path in sorted(os.listdir(directory)) + if os.path.isfile(os.path.join(directory, path)) + ] + output_stream.write(struct.pack(">i", len(files))) + for file_name in files: + file_path = os.path.join(directory, file_name) + name_bytes = file_name.encode("utf-8") + output_stream.write(struct.pack(">i", len(name_bytes))) + output_stream.write(name_bytes) + output_stream.write(struct.pack(">q", os.path.getsize(file_path))) + with open(file_path, "rb") as input_stream: + while True: + chunk = input_stream.read(8192) + if not chunk: + break + output_stream.write(chunk) diff --git a/paimon-python/pypaimon/tests/global_index_build_test.py b/paimon-python/pypaimon/tests/global_index_build_test.py index 957f6583be..6006a8ae83 100644 --- a/paimon-python/pypaimon/tests/global_index_build_test.py +++ b/paimon-python/pypaimon/tests/global_index_build_test.py @@ -19,6 +19,7 @@ import unittest from datetime import date, datetime from decimal import Decimal import os +import struct import sys import types @@ -30,6 +31,15 @@ from pypaimon.globalindex.create_global_index import ( _split_one_by_contiguous_row_range, ) from pypaimon.globalindex.key_serializer import create_serializer +from pypaimon.globalindex.tantivy.tantivy_full_text_global_index_reader import ( + TANTIVY_FULLTEXT_IDENTIFIER, + TANTIVY_JIEBA_TOKENIZER, + TANTIVY_NGRAM_TOKENIZER, + TantivyFullTextIndexOptions, +) +from pypaimon.globalindex.tantivy.tantivy_full_text_index_writer import ( + TantivyFullTextIndexWriter, +) from pypaimon.globalindex.vindex.vindex_vector_index_writer import ( VindexVectorIndexWriter, native_options, @@ -101,6 +111,193 @@ class _FakeVectorIndexWriter: return False +class _FakeTantivyDocument: + + def __init__(self): + self.values = {} + + def add_unsigned(self, field_name, value): + self.values[field_name] = value + + def add_text(self, field_name, text): + self.values[field_name] = text + + +class _FakeTantivyIndexWriter: + + def __init__(self, index): + self.index = index + self.documents = [] + self.committed = False + self.waited = False + + def add_document(self, document): + self.documents.append(dict(document.values)) + + def commit(self): + self.committed = True + with open(os.path.join(self.index.path, "meta.json"), "wb") as f: + f.write(b"{}") + with open(os.path.join(self.index.path, "docs.bin"), "wb") as f: + for document in self.documents: + row_id = document["row_id"] + text = document["text"].encode("utf-8") + f.write(struct.pack(">q", row_id)) + f.write(struct.pack(">i", len(text))) + f.write(text) + + def wait_merging_threads(self): + self.waited = True + + +class _FakeTantivySchemaBuilder: + + def __init__(self, parent): + self._parent = parent + self.fields = {} + + def add_unsigned_field(self, name, stored=False, indexed=False, fast=False): + self.fields[name] = { + "stored": stored, + "indexed": indexed, + "fast": fast, + } + + def add_text_field(self, name, stored=False, tokenizer_name=None, **kwargs): + self.fields[name] = { + "stored": stored, + "tokenizer_name": tokenizer_name or "default", + } + if "index_option" in kwargs: + self.fields[name]["index_option"] = kwargs["index_option"] + + def build(self): + schema = types.SimpleNamespace(fields=self.fields) + self._parent.last_schema = schema + return schema + + +class _FakeTantivyTokenizer: + + @staticmethod + def ngram(min_gram=2, max_gram=3, prefix_only=False): + return ("ngram", min_gram, max_gram, prefix_only) + + @staticmethod + def simple(): + return ("simple",) + + @staticmethod + def whitespace(): + return ("whitespace",) + + @staticmethod + def raw(): + return ("raw",) + + +class _FakeTantivyFilter: + + @staticmethod + def lowercase(): + return "lowercase" + + @staticmethod + def remove_long(length_limit): + return ("remove_long", length_limit) + + @staticmethod + def ascii_fold(): + return "ascii_fold" + + @staticmethod + def stemmer(language): + return ("stemmer", language) + + @staticmethod + def stopword(language): + return ("stopword", language) + + @staticmethod + def custom_stopword(stopwords): + return ("custom_stopword", tuple(stopwords)) + + +class _FakeTantivyTextAnalyzerBuilder: + + def __init__(self, tokenizer): + self._tokenizer = tokenizer + self._filters = [] + + def filter(self, filter_): + result = _FakeTantivyTextAnalyzerBuilder(self._tokenizer) + result._filters = self._filters + [filter_] + return result + + def build(self): + return self._tokenizer + (tuple(self._filters),) + + +class _FakeTantivyIndex: + + def __init__(self, parent, schema, path=None): + self.parent = parent + self.schema = schema + self.path = path + self.registered_tokenizers = [] + self.writer_instance = _FakeTantivyIndexWriter(self) + parent.indexes.append(self) + parent.last_index = self + + def register_tokenizer(self, name, analyzer): + self.registered_tokenizers.append((name, analyzer)) + + def writer(self): + return self.writer_instance + + +class _FakeTantivyForBuild(types.SimpleNamespace): + + def __init__(self): + super().__init__() + self.Document = _FakeTantivyDocument + self.Tokenizer = _FakeTantivyTokenizer + self.Filter = _FakeTantivyFilter + self.TextAnalyzerBuilder = _FakeTantivyTextAnalyzerBuilder + self.indexes = [] + self.last_schema = None + self.last_index = None + parent = self + + class SchemaBuilder(_FakeTantivySchemaBuilder): + + def __init__(self_inner): + super().__init__(parent) + + class Index(_FakeTantivyIndex): + + def __init__(self_inner, schema, path=None): + super().__init__(parent, schema, path=path) + + self.SchemaBuilder = SchemaBuilder + self.Index = Index + + +def _archive_file_names(file_io, file_path): + stream = file_io.new_input_stream(file_path) + try: + file_count = struct.unpack(">i", stream.read(4))[0] + names = [] + for _ in range(file_count): + name_len = struct.unpack(">i", stream.read(4))[0] + names.append(stream.read(name_len).decode("utf-8")) + data_len = struct.unpack(">q", stream.read(8))[0] + stream.seek(stream.tell() + data_len) + return names + finally: + stream.close() + + class _FakeSchemaManager: def __init__(self, fields_by_schema_id): @@ -511,6 +708,189 @@ class GlobalIndexBuildTest( ], ) + def test_create_tantivy_fulltext_global_index_from_python(self): + schema = pa.schema([ + ('id', pa.int32()), + ('content', pa.string()), + ]) + table = self._create_table(pa_schema=schema, options=self.table_options) + self._write_arrow(table, pa.table( + { + 'id': [1, 2, 3], + 'content': [ + 'Apache Paimon full text', + None, + 'Tantivy full text search', + ], + }, + schema=schema, + )) + + tantivy = _FakeTantivyForBuild() + old_tantivy = sys.modules.get("tantivy") + sys.modules["tantivy"] = tantivy + try: + added = table.create_global_index( + 'content', + index_type=TANTIVY_FULLTEXT_IDENTIFIER, + options={ + 'global-index.row-count-per-shard': '2', + 'tantivy.tokenizer': 'ngram', + 'tantivy.ngram.min-gram': '2', + 'tantivy.ngram.max-gram': '3', + 'tantivy.ngram.prefix-only': 'true', + 'tantivy.with-position': 'false', + }, + ) + finally: + if old_tantivy is None: + sys.modules.pop("tantivy", None) + else: + sys.modules["tantivy"] = old_tantivy + + self.assertEqual(2, added) + self.assertEqual(2, len(tantivy.indexes)) + self.assertEqual( + {"row_id": {"stored": False, "indexed": True, "fast": True}, + "text": {"stored": False, + "tokenizer_name": TANTIVY_NGRAM_TOKENIZER, + "index_option": "freq"}}, + tantivy.indexes[0].schema.fields, + ) + self.assertEqual( + [(TANTIVY_NGRAM_TOKENIZER, + ("ngram", 2, 3, True, (("remove_long", 40), "lowercase")))], + tantivy.indexes[0].registered_tokenizers, + ) + self.assertEqual( + [{'row_id': 0, 'text': 'Apache Paimon full text'}], + tantivy.indexes[0].writer_instance.documents, + ) + self.assertEqual( + [{'row_id': 0, 'text': 'Tantivy full text search'}], + tantivy.indexes[1].writer_instance.documents, + ) + + snapshot = table.snapshot_manager().get_latest_snapshot() + entries = sorted( + IndexFileHandler(table).scan(snapshot), + key=lambda entry: entry.index_file.global_index_meta.row_range_start, + ) + self.assertEqual( + [(0, 1, 2), (2, 2, 1)], + [ + ( + entry.index_file.global_index_meta.row_range_start, + entry.index_file.global_index_meta.row_range_end, + entry.index_file.row_count, + ) + for entry in entries + ], + ) + self.assertEqual( + {TANTIVY_FULLTEXT_IDENTIFIER}, + {entry.index_file.index_type for entry in entries}, + ) + for entry in entries: + index_options = TantivyFullTextIndexOptions.deserialize( + entry.index_file.global_index_meta.index_meta) + self.assertEqual("ngram", index_options.tokenizer) + self.assertEqual(2, index_options.ngram_min_gram) + self.assertEqual(3, index_options.ngram_max_gram) + self.assertTrue(index_options.ngram_prefix_only) + self.assertFalse(index_options.with_position) + file_path = table.path_factory().global_index_path_factory().to_path( + entry.index_file.file_name) + self.assertEqual( + ['docs.bin', 'meta.json'], + _archive_file_names(table.file_io, file_path), + ) + + def test_tantivy_fulltext_writer_rejects_jieba_tokenizer(self): + table = self._create_table() + index_path = ( + table.path_factory() + .global_index_path_factory() + .global_index_root_path() + ) + with self.assertRaisesRegex(ValueError, "tantivy.tokenizer=jieba"): + writer = TantivyFullTextIndexWriter( + table.file_io, + index_path, + AtomicType('STRING'), + {'tantivy.tokenizer': 'jieba'}, + ) + writer.close() + + def test_tantivy_fulltext_options_serialize_matches_java_sparse_json(self): + ngram = TantivyFullTextIndexOptions.from_options({ + 'tantivy.tokenizer': ' NGRAM ', + 'tantivy.ngram.min-gram': '2', + 'tantivy.ngram.max-gram': '3', + 'tantivy.ngram.prefix-only': 'true', + 'tantivy.lower-case': 'false', + }) + self.assertEqual( + b'{"tokenizer":"ngram","ngram.max-gram":3,' + b'"ngram.prefix-only":true,"lower-case":false}', + ngram.serialize(), + ) + + analyzer = TantivyFullTextIndexOptions.from_options({ + 'tantivy.tokenizer': 'whitespace', + 'tantivy.max-token-length': '12', + 'tantivy.ascii-folding': 'true', + 'tantivy.stem': 'true', + 'tantivy.remove-stop-words': 'true', + 'tantivy.stop-words': 'paimon;lake', + 'tantivy.with-position': 'false', + }) + self.assertEqual( + b'{"tokenizer":"whitespace","max-token-length":12,' + b'"ascii-folding":true,"stem":true,' + b'"remove-stop-words":true,"stop-words":["paimon","lake"],' + b'"with-position":false}', + analyzer.serialize(), + ) + + def test_tantivy_fulltext_options_validate_like_java(self): + with self.assertRaisesRegex(ValueError, "Unsupported Tantivy tokenizer"): + TantivyFullTextIndexOptions.from_options({ + 'tantivy.tokenizer': 'ik', + }) + + with self.assertRaisesRegex( + ValueError, "ngram min gram must not be greater than max gram"): + TantivyFullTextIndexOptions.from_options({ + 'tantivy.tokenizer': 'ngram', + 'tantivy.ngram.min-gram': '3', + 'tantivy.ngram.max-gram': '2', + }) + + with self.assertRaisesRegex(ValueError, "Unsupported Tantivy language"): + TantivyFullTextIndexOptions.from_options({ + 'tantivy.stem': 'true', + 'tantivy.language': 'klingon', + }) + + def test_create_tantivy_fulltext_global_index_rejects_non_string_column(self): + table = self._create_table() + self._write_arrow(table, pa.table( + { + 'id': [1], + 'name': ['a'], + 'age': [10], + 'city': ['x'], + }, + schema=self.pa_schema, + )) + + with self.assertRaisesRegex(ValueError, 'requires string type'): + table.create_global_index( + 'id', + index_type=TANTIVY_FULLTEXT_IDENTIFIER, + ) + def test_create_vindex_global_index_rejects_generic_unsupported_tables(self): schema = pa.schema([ ('id', pa.int32()), diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py b/paimon-python/pypaimon/tests/vector_search_filter_test.py index 19b9bca2d4..f189ec471c 100644 --- a/paimon-python/pypaimon/tests/vector_search_filter_test.py +++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py @@ -344,6 +344,24 @@ class _FakeQuery: def empty_query(): return ("empty",) + @staticmethod + def boost_query(query, boost): + return ("boost", query, boost) + + @staticmethod + def fuzzy_term_query( + schema, field_name, text, distance=1, + transposition_cost_one=True, prefix=False): + return ( + "fuzzy", + schema, + field_name, + text, + distance, + transposition_cost_one, + prefix, + ) + @staticmethod def term_query(schema, field_name, field_value, index_option="position"): return ("term", schema, field_name, field_value, index_option) @@ -360,19 +378,39 @@ class _FakeOccur: class _FakeSearchResults: - hits = [(2.0, "addr")] + def __init__(self, hits=None): + self.hits = hits if hits is not None else [(2.0, "addr")] class _FakeSearcher: def __init__(self): self.query = None + self.queries = [] def search(self, query, limit): self.query = query + self.queries.append(query) + query_text = _fake_query_text(query) + if query_text == "positive": + return _FakeSearchResults([(10.0, ("addr", 1)), (5.0, ("addr", 2))]) + if query_text == "negative": + return _FakeSearchResults([(7.0, ("addr", 2))]) return _FakeSearchResults() def fast_field_values(self, name, addresses): - return [7] + return [ + address[1] if isinstance(address, tuple) else 7 + for address in addresses + ] + + +def _fake_query_text(query): + if isinstance(query, tuple) and query: + if query[0] == "boost": + return _fake_query_text(query[1]) + if isinstance(query[0], str): + return query[0] + return None class _FakeIndex: @@ -774,6 +812,170 @@ class TantivyFullTextIndexOptionsTest(unittest.TestCase): ("stopword", "english"), ("custom_stopword", ("paimon", "lake"))))), tantivy.last_index.registered_tokenizer) + def test_match_reader_aligns_java_boost_and_fuzziness(self): + from pypaimon.globalindex.full_text_search import FullTextSearch + from pypaimon.globalindex.tantivy.tantivy_full_text_global_index_reader import ( + TantivyFullTextGlobalIndexReader, + ) + + tantivy = _FakeTantivy() + old_tantivy = sys.modules.get("tantivy") + sys.modules["tantivy"] = tantivy + try: + reader = TantivyFullTextGlobalIndexReader( + _FakeFileIO(), + "/unused", + [GlobalIndexIOMeta(file_name="ft.index", file_size=1)]) + try: + reader.visit_full_text_search( + FullTextSearch( + MatchQuery( + "paimon", + "content", + boost=2.0, + fuzziness=1, + operator="and", + ), + 10, + ) + ).result() + finally: + reader.close() + finally: + if old_tantivy is None: + sys.modules.pop("tantivy", None) + else: + sys.modules["tantivy"] = old_tantivy + + self.assertEqual( + ("boost", + ("paimon", + ("text",), + {"conjunction_by_default": True, + "fuzzy_fields": {"text": (False, 1, True)}}), + 2.0), + tantivy.last_index.searcher_instance.query, + ) + + def test_boost_reader_demotes_negative_matches_like_java(self): + from pypaimon.globalindex.full_text_query import BoostQuery + from pypaimon.globalindex.full_text_search import FullTextSearch + from pypaimon.globalindex.tantivy.tantivy_full_text_global_index_reader import ( + TantivyFullTextGlobalIndexReader, + ) + + tantivy = _FakeTantivy() + old_tantivy = sys.modules.get("tantivy") + sys.modules["tantivy"] = tantivy + try: + reader = TantivyFullTextGlobalIndexReader( + _FakeFileIO(), + "/unused", + [GlobalIndexIOMeta(file_name="ft.index", file_size=1)]) + try: + result = reader.visit_full_text_search( + FullTextSearch( + BoostQuery( + MatchQuery("positive", "content"), + MatchQuery("negative", "content"), + negative_boost=0.2, + ), + 10, + ) + ).result() + finally: + reader.close() + finally: + if old_tantivy is None: + sys.modules.pop("tantivy", None) + else: + sys.modules["tantivy"] = old_tantivy + + self.assertEqual([1, 2], sorted(list(result.results()))) + score_getter = result.score_getter() + self.assertEqual(10.0, score_getter(1)) + self.assertEqual(1.0, score_getter(2)) + self.assertEqual( + ["positive", "negative"], + [_fake_query_text(q) for q in tantivy.last_index.searcher_instance.queries], + ) + + def test_reader_rejects_java_unsupported_match_options(self): + from pypaimon.globalindex.full_text_search import FullTextSearch + from pypaimon.globalindex.tantivy.tantivy_full_text_global_index_reader import ( + TantivyFullTextGlobalIndexReader, + ) + + tantivy = _FakeTantivy() + old_tantivy = sys.modules.get("tantivy") + sys.modules["tantivy"] = tantivy + try: + reader = TantivyFullTextGlobalIndexReader( + _FakeFileIO(), + "/unused", + [GlobalIndexIOMeta(file_name="ft.index", file_size=1)]) + try: + with self.assertRaisesRegex(ValueError, "max_expansions"): + reader.visit_full_text_search( + FullTextSearch( + MatchQuery( + "paimon", + "content", + max_expansions=10, + ), + 10, + ) + ).result() + with self.assertRaisesRegex(ValueError, "prefix_length"): + reader.visit_full_text_search( + FullTextSearch( + MatchQuery( + "paimon", + "content", + prefix_length=1, + ), + 10, + ) + ).result() + finally: + reader.close() + finally: + if old_tantivy is None: + sys.modules.pop("tantivy", None) + else: + sys.modules["tantivy"] = old_tantivy + + def test_reader_rejects_single_column_multi_match_like_java(self): + from pypaimon.globalindex.full_text_query import MultiMatchQuery + from pypaimon.globalindex.full_text_search import FullTextSearch + from pypaimon.globalindex.tantivy.tantivy_full_text_global_index_reader import ( + TantivyFullTextGlobalIndexReader, + ) + + tantivy = _FakeTantivy() + old_tantivy = sys.modules.get("tantivy") + sys.modules["tantivy"] = tantivy + try: + reader = TantivyFullTextGlobalIndexReader( + _FakeFileIO(), + "/unused", + [GlobalIndexIOMeta(file_name="ft.index", file_size=1)]) + try: + with self.assertRaisesRegex(ValueError, "multi_match"): + reader.visit_full_text_search( + FullTextSearch( + MultiMatchQuery("paimon", ["title", "content"]), + 10, + ) + ).result() + finally: + reader.close() + finally: + if old_tantivy is None: + sys.modules.pop("tantivy", None) + else: + sys.modules["tantivy"] = old_tantivy + def test_ngram_reader_requires_custom_tokenizer_api(self): from pypaimon.globalindex.full_text_search import FullTextSearch from pypaimon.globalindex.tantivy.tantivy_full_text_global_index_reader import (
