| # 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. |
| |
| """Builder to build full-text search.""" |
| |
| from abc import ABC, abstractmethod |
| from typing import Optional |
| |
| from pypaimon.common.predicate_builder import PredicateBuilder |
| from pypaimon.globalindex.global_index_result import GlobalIndexResult |
| from pypaimon.table.source.full_text_read import FullTextRead, DataEvolutionFullTextRead |
| from pypaimon.table.source.full_text_scan import FullTextScan, DataEvolutionFullTextScan |
| from pypaimon.common.options.core_options import CoreOptions |
| from pypaimon.common.options.options import Options |
| from pypaimon.index.pk.primary_key_index_definition import PrimaryKeyIndexFamily |
| from pypaimon.index.pk.primary_key_index_definitions import PrimaryKeyIndexDefinitions |
| |
| |
| class FullTextSearchBuilder(ABC): |
| """Builder to build full-text search.""" |
| |
| @abstractmethod |
| def with_limit(self, limit: int) -> 'FullTextSearchBuilder': |
| """The top k results to return.""" |
| pass |
| |
| @abstractmethod |
| def with_query(self, field_name: str, query: str) -> 'FullTextSearchBuilder': |
| """The full-text query string to search against the given field.""" |
| pass |
| |
| @abstractmethod |
| def with_partition_filter(self, partition_filter) -> 'FullTextSearchBuilder': |
| """Partition predicate used to prune index manifest entries.""" |
| pass |
| |
| @abstractmethod |
| def new_full_text_scan(self) -> FullTextScan: |
| """Create full-text scan to scan index files.""" |
| pass |
| |
| @abstractmethod |
| def new_full_text_read(self) -> FullTextRead: |
| """Create full-text read to read index files.""" |
| pass |
| |
| def execute_local(self) -> GlobalIndexResult: |
| """Execute full-text index search in local.""" |
| return self.new_full_text_read().read_plan(self.new_full_text_scan().scan()) |
| |
| |
| class FullTextSearchBuilderImpl(FullTextSearchBuilder): |
| """Implementation for FullTextSearchBuilder.""" |
| |
| def __init__(self, table: 'FileStoreTable'): |
| self._table = table |
| self._limit: int = 0 |
| self._field_name: Optional[str] = None |
| self._query: Optional[str] = None |
| self._partition_filter = None |
| |
| def with_limit(self, limit: int) -> 'FullTextSearchBuilder': |
| self._limit = limit |
| return self |
| |
| def with_query(self, field_name: str, query: str) -> 'FullTextSearchBuilder': |
| self._field_name = field_name |
| self._query = query |
| return self |
| |
| def with_partition_filter(self, partition_filter) -> 'FullTextSearchBuilder': |
| if partition_filter is None: |
| self._partition_filter = None |
| return self |
| partition_keys = list(self._table.partition_keys or []) |
| if not partition_keys: |
| raise ValueError( |
| "with_partition_filter called on a non-partitioned table") |
| from pypaimon.read.push_down_utils import _get_all_fields |
| referenced = _get_all_fields(partition_filter) |
| extras = referenced - set(partition_keys) |
| if extras: |
| raise ValueError( |
| "Partition filter must reference only partition keys " |
| "(%s); got non-partition field(s): %s" |
| % (partition_keys, sorted(extras))) |
| rebuilt = self._rebuild_leaf_indices_by_name( |
| partition_filter, |
| {name: idx for idx, name in enumerate(partition_keys)}, |
| ) |
| if self._partition_filter is None: |
| self._partition_filter = rebuilt |
| else: |
| self._partition_filter = PredicateBuilder.and_predicates( |
| [self._partition_filter, rebuilt]) |
| return self |
| |
| @classmethod |
| def _rebuild_leaf_indices_by_name(cls, predicate, name_to_idx): |
| if predicate.method in ('and', 'or'): |
| return predicate.new_literals( |
| [cls._rebuild_leaf_indices_by_name(c, name_to_idx) |
| for c in (predicate.literals or [])]) |
| return predicate.new_index(name_to_idx[predicate.field]) |
| |
| def new_full_text_scan(self) -> FullTextScan: |
| definition = self._primary_key_full_text_definition() |
| if definition is not None: |
| from pypaimon.table.source.primary_key_full_text_scan import PrimaryKeyFullTextScan |
| return PrimaryKeyFullTextScan( |
| self._table, definition, partition_filter=self._partition_filter) |
| return DataEvolutionFullTextScan( |
| self._table, |
| self._text_columns(), |
| partition_filter=self._partition_filter, |
| ) |
| |
| def new_full_text_read(self) -> FullTextRead: |
| if self._limit <= 0: |
| raise ValueError("Limit must be positive, set via with_limit()") |
| definition = self._primary_key_full_text_definition() |
| if definition is not None: |
| from pypaimon.common.options.core_options import GlobalIndexSearchMode |
| mode = CoreOptions( |
| Options(dict(self._table.table_schema.options)) |
| ).full_text_index_search_mode() |
| if mode != GlobalIndexSearchMode.FAST: |
| raise NotImplementedError( |
| "Primary-key full-text search only supports the FAST " |
| "full-text index search mode; FULL and DETAIL require " |
| "merge-aware logical-row fallback.") |
| from pypaimon.table.source.primary_key_full_text_read import PrimaryKeyFullTextRead |
| return PrimaryKeyFullTextRead( |
| self._table, self._limit, self._text_columns(), self._query, |
| definition=definition, partition_filter=self._partition_filter) |
| return DataEvolutionFullTextRead( |
| self._table, |
| self._limit, |
| self._text_columns(), |
| self._query, |
| partition_filter=self._partition_filter, |
| ) |
| |
| def _text_columns(self): |
| if self._query is None: |
| raise ValueError("Query must be set via with_query()") |
| if self._field_name is None: |
| raise ValueError("Field name must be set via with_query()") |
| field_dict = {f.name: f for f in self._table.fields} |
| if self._field_name not in field_dict: |
| raise ValueError( |
| f"Text column '{self._field_name}' not found in table schema") |
| return [field_dict[self._field_name]] |
| |
| def _primary_key_full_text_definition(self): |
| text_column = self._text_columns()[0] |
| core = CoreOptions(Options(dict(self._table.table_schema.options))) |
| if (core.data_evolution_enabled() |
| or not core.primary_key_full_text_index_columns()): |
| return None |
| for definition in PrimaryKeyIndexDefinitions.create( |
| self._table.table_schema).definitions: |
| if (definition.family == PrimaryKeyIndexFamily.FULL_TEXT |
| and definition.field_id == text_column.id): |
| return definition |
| return None |
