| # 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 hybrid search.""" |
| |
| import heapq |
| import math |
| from abc import ABC, abstractmethod |
| from dataclasses import dataclass, field |
| from typing import Dict, List, Optional |
| |
| from pypaimon.common.predicate_builder import PredicateBuilder |
| from pypaimon.globalindex.global_index_result import GlobalIndexResult |
| from pypaimon.globalindex.vector_search_result import ( |
| DictBasedScoredIndexResult, |
| ScoredGlobalIndexResult, |
| ) |
| |
| RRF_RANKER = "rrf" |
| WEIGHTED_SCORE_RANKER = "weighted_score" |
| MRR_RANKER = "mrr" |
| _RRF_K = 60.0 |
| |
| |
| def _check_full_text_options(options: Dict[str, str]): |
| if options: |
| raise ValueError( |
| "Full-text hybrid route options are not supported yet.") |
| |
| |
| def _normalize_ranker(ranker: Optional[str]) -> str: |
| if ranker is None or not ranker.strip(): |
| return RRF_RANKER |
| normalized = ranker.strip().lower() |
| if normalized not in (RRF_RANKER, WEIGHTED_SCORE_RANKER, MRR_RANKER): |
| raise ValueError("Unsupported hybrid ranker: %s" % ranker) |
| return normalized |
| |
| |
| @dataclass |
| class HybridSearchRoute: |
| """A single route in a hybrid search.""" |
| |
| route_type: str |
| field_name: str |
| limit: int |
| weight: float = 1.0 |
| vector: Optional[List[float]] = None |
| full_text_query: Optional[str] = None |
| options: Dict[str, str] = field(default_factory=dict) |
| |
| VECTOR = "vector" |
| FULL_TEXT = "full_text" |
| |
| def __post_init__(self): |
| if self.route_type not in (self.VECTOR, self.FULL_TEXT): |
| raise ValueError("Unsupported hybrid route type: %s" % self.route_type) |
| if not self.field_name: |
| raise ValueError("Field name cannot be None or empty") |
| if self.route_type == self.VECTOR and self.vector is None: |
| raise ValueError("Search vector cannot be None for vector route") |
| if self.route_type == self.FULL_TEXT and self.full_text_query is None: |
| raise ValueError( |
| "Query cannot be None for full-text route") |
| if self.limit <= 0: |
| raise ValueError("Limit must be positive, got: %s" % self.limit) |
| if not math.isfinite(self.weight) or self.weight <= 0: |
| raise ValueError( |
| "Weight must be finite and positive, got: %s" % self.weight) |
| self.options = dict(self.options or {}) |
| if self.route_type == self.FULL_TEXT: |
| _check_full_text_options(self.options) |
| |
| @classmethod |
| def vector_route( |
| cls, |
| field_name: str, |
| vector: List[float], |
| limit: int, |
| weight: float = 1.0, |
| options: Optional[Dict[str, str]] = None) -> 'HybridSearchRoute': |
| return cls( |
| route_type=cls.VECTOR, |
| field_name=field_name, |
| vector=vector, |
| limit=limit, |
| weight=weight, |
| options=dict(options or {}), |
| ) |
| |
| @classmethod |
| def full_text_route( |
| cls, |
| field_name: str, |
| query: str, |
| limit: int, |
| weight: float = 1.0, |
| options: Optional[Dict[str, str]] = None) -> 'HybridSearchRoute': |
| _check_full_text_options(options or {}) |
| return cls( |
| route_type=cls.FULL_TEXT, |
| field_name=field_name, |
| full_text_query=query, |
| limit=limit, |
| weight=weight, |
| options=dict(options or {}), |
| ) |
| |
| def is_vector(self) -> bool: |
| return self.route_type == self.VECTOR |
| |
| def is_full_text(self) -> bool: |
| return self.route_type == self.FULL_TEXT |
| |
| def columns(self) -> List[str]: |
| return [self.field_name] |
| |
| |
| @dataclass |
| class HybridSearchRouteResult: |
| route: HybridSearchRoute |
| result: ScoredGlobalIndexResult |
| |
| |
| class HybridSearchRouteBuilder: |
| """A route plus its configured search builder.""" |
| |
| def __init__(self, route, search_builder): |
| self.route = route |
| self.search_builder = search_builder |
| |
| def execute_local(self): |
| return self.search_builder.execute_local() |
| |
| |
| class HybridSearchBuilder(ABC): |
| """Builder to build hybrid search.""" |
| |
| @abstractmethod |
| def with_limit(self, limit: int) -> 'HybridSearchBuilder': |
| """The final top k ranked results to return.""" |
| pass |
| |
| @abstractmethod |
| def with_ranker(self, ranker: str) -> 'HybridSearchBuilder': |
| """Ranker for combining route results.""" |
| pass |
| |
| def with_rrf_ranker(self) -> 'HybridSearchBuilder': |
| """Use reciprocal rank fusion to combine route results.""" |
| return self.with_ranker(RRF_RANKER) |
| |
| def with_weighted_score_ranker(self) -> 'HybridSearchBuilder': |
| """Use weighted score to combine route results.""" |
| return self.with_ranker(WEIGHTED_SCORE_RANKER) |
| |
| @abstractmethod |
| def with_filter(self, predicate) -> 'HybridSearchBuilder': |
| """Scalar predicate used to pre-filter vector routes.""" |
| pass |
| |
| @abstractmethod |
| def with_partition_filter(self, partition_filter) -> 'HybridSearchBuilder': |
| """Partition predicate used to prune index manifest entries.""" |
| pass |
| |
| @abstractmethod |
| def add_route(self, route: HybridSearchRoute) -> 'HybridSearchBuilder': |
| """Add a hybrid-search route.""" |
| pass |
| |
| def add_vector_route( |
| self, |
| vector_column: str, |
| vector: List[float], |
| limit: int, |
| weight: float = 1.0, |
| options: Optional[Dict[str, str]] = None) -> 'HybridSearchBuilder': |
| """Add a vector-search route.""" |
| return self.add_route( |
| HybridSearchRoute.vector_route( |
| vector_column, vector, limit, weight, options)) |
| |
| def add_full_text_route( |
| self, |
| field_name: str, |
| query: str, |
| limit: int, |
| weight: float = 1.0, |
| options: Optional[Dict[str, str]] = None) -> 'HybridSearchBuilder': |
| """Add a full-text-search route.""" |
| return self.add_route( |
| HybridSearchRoute.full_text_route( |
| field_name, query, limit, weight, options)) |
| |
| @abstractmethod |
| def route_builders(self) -> List[HybridSearchRouteBuilder]: |
| """Create one search builder for each route.""" |
| pass |
| |
| @abstractmethod |
| def to_route_result( |
| self, |
| route_builder: HybridSearchRouteBuilder, |
| result: GlobalIndexResult) -> HybridSearchRouteResult: |
| """Convert a search result into a weighted route result.""" |
| pass |
| |
| @abstractmethod |
| def rank( |
| self, |
| route_results: List[HybridSearchRouteResult]) -> ScoredGlobalIndexResult: |
| """Rank route results.""" |
| pass |
| |
| def execute_local(self) -> ScoredGlobalIndexResult: |
| """Execute hybrid index search locally.""" |
| route_builders = self.route_builders() |
| route_results = [] |
| for route_builder in route_builders: |
| route_results.append( |
| self.to_route_result( |
| route_builder, route_builder.execute_local())) |
| return self.rank(route_results) |
| |
| |
| class HybridSearchBuilderImpl(HybridSearchBuilder): |
| """Implementation for HybridSearchBuilder.""" |
| |
| def __init__(self, table): |
| self._table = table |
| self._routes = [] |
| self._limit = 0 |
| self._ranker = RRF_RANKER |
| self._filter = None |
| self._partition_filter = None |
| |
| def with_limit(self, limit: int) -> 'HybridSearchBuilder': |
| self._limit = limit |
| return self |
| |
| def with_ranker(self, ranker: str) -> 'HybridSearchBuilder': |
| self._ranker = _normalize_ranker(ranker) |
| return self |
| |
| def with_filter(self, predicate) -> 'HybridSearchBuilder': |
| if predicate is None: |
| return self |
| partition_filter, data_filter = self._split_partition_filter(predicate) |
| if partition_filter is not None: |
| self._add_partition_filter(partition_filter) |
| if data_filter is not None: |
| if self._filter is None: |
| self._filter = data_filter |
| else: |
| self._filter = PredicateBuilder.and_predicates( |
| [self._filter, data_filter]) |
| return self |
| |
| def with_partition_filter(self, partition_filter) -> 'HybridSearchBuilder': |
| if partition_filter is not None: |
| self._add_partition_filter( |
| self._validate_and_rebuild_partition_filter(partition_filter)) |
| return self |
| |
| def add_route(self, route: HybridSearchRoute) -> 'HybridSearchBuilder': |
| if route is None: |
| raise ValueError("Route cannot be None") |
| self._routes.append(route) |
| return self |
| |
| def route_builders(self) -> List[HybridSearchRouteBuilder]: |
| self._validate_search() |
| builders = [] |
| for route in self._routes: |
| if route.is_vector(): |
| builders.append( |
| HybridSearchRouteBuilder( |
| route, self._new_vector_search_builder(route))) |
| else: |
| builders.append( |
| HybridSearchRouteBuilder( |
| route, self._new_full_text_search_builder(route))) |
| return builders |
| |
| def to_route_result( |
| self, |
| route_builder: HybridSearchRouteBuilder, |
| result: GlobalIndexResult) -> HybridSearchRouteResult: |
| if isinstance(result, ScoredGlobalIndexResult): |
| scored = result |
| elif result.results().is_empty(): |
| scored = ScoredGlobalIndexResult.create_empty() |
| else: |
| raise ValueError( |
| "Hybrid search requires scored index results, but got: %s" |
| % result.__class__.__name__) |
| return HybridSearchRouteResult(route_builder.route, scored) |
| |
| def rank( |
| self, |
| route_results: List[HybridSearchRouteResult]) -> ScoredGlobalIndexResult: |
| self._validate_search() |
| non_empty = [ |
| route_result for route_result in route_results |
| if not route_result.result.results().is_empty() |
| ] |
| if self._ranker == WEIGHTED_SCORE_RANKER: |
| return self._weighted_score(non_empty) |
| if self._ranker == MRR_RANKER: |
| return self._mrr(non_empty) |
| return self._rrf(non_empty) |
| |
| def _validate_search(self): |
| if not self._routes: |
| raise ValueError("Routes cannot be empty") |
| if self._limit <= 0: |
| raise ValueError("Limit must be positive, got: %s" % self._limit) |
| if self._filter is not None: |
| for route in self._routes: |
| if route.is_full_text(): |
| raise ValueError( |
| "Hybrid search with full-text routes does not support " |
| "non-partition filters because full-text indexes cannot " |
| "apply row-id pre-filters before top-k ranking.") |
| |
| def _new_vector_search_builder(self, route): |
| builder = ( |
| self._table.new_vector_search_builder() |
| .with_vector_column(route.field_name) |
| .with_query_vector(route.vector) |
| .with_limit(route.limit) |
| .with_options(route.options) |
| ) |
| if self._partition_filter is not None: |
| builder.with_partition_filter(self._partition_filter) |
| if self._filter is not None: |
| builder.with_filter(self._filter) |
| return builder |
| |
| def _new_full_text_search_builder(self, route): |
| builder = ( |
| self._table.new_full_text_search_builder() |
| .with_query(route.field_name, route.full_text_query) |
| .with_limit(route.limit) |
| ) |
| if self._partition_filter is not None: |
| builder.with_partition_filter(self._partition_filter) |
| return builder |
| |
| def _rrf(self, route_results): |
| scores = {} |
| for route_result in route_results: |
| for rank, row_id in enumerate(_ranked_row_ids(route_result.result)): |
| contribution = route_result.route.weight / (_RRF_K + rank + 1.0) |
| scores[row_id] = scores.get(row_id, 0.0) + contribution |
| return _top_k(scores, self._limit) |
| |
| def _mrr(self, route_results): |
| scores = {} |
| for route_result in route_results: |
| for rank, row_id in enumerate(_ranked_row_ids(route_result.result)): |
| contribution = route_result.route.weight / (rank + 1.0) |
| scores[row_id] = scores.get(row_id, 0.0) + contribution |
| return _top_k(scores, self._limit) |
| |
| def _weighted_score(self, route_results): |
| scores = {} |
| for route_result in route_results: |
| result = route_result.result |
| weight = route_result.route.weight |
| score_getter = result.score_getter() |
| |
| # Route score scales are heterogeneous (e.g. bounded vector similarity |
| # vs unbounded BM25), so raw scores are not comparable across routes. |
| # Min-max normalize each route into [0, 1] before weighting, so that |
| # weights -- not a route's numeric magnitude -- decide its influence on |
| # the fused score. This mirrors the Java HybridSearchRanker. |
| route_scores = { |
| row_id: (score_getter(row_id) or 0.0) |
| for row_id in result.results() |
| } |
| if not route_scores: |
| continue |
| min_score = min(route_scores.values()) |
| score_range = max(route_scores.values()) - min_score |
| for row_id, raw in route_scores.items(): |
| # No spread within the route (single hit or all ties) carries no |
| # relative signal, so every hit maps to 1.0 rather than zeroed out. |
| normalized = ( |
| (raw - min_score) / score_range if score_range > 0.0 else 1.0) |
| scores[row_id] = scores.get(row_id, 0.0) + weight * normalized |
| return _top_k(scores, self._limit) |
| |
| def _split_partition_filter(self, predicate): |
| partition_keys = list(self._table.partition_keys or []) |
| if not partition_keys: |
| return None, predicate |
| from pypaimon.read.push_down_utils import _get_all_fields, _split_and |
| partition_key_set = set(partition_keys) |
| pk_to_idx = {name: idx for idx, name in enumerate(partition_keys)} |
| partition_parts = [] |
| data_parts = [] |
| for part in _split_and(predicate): |
| if _get_all_fields(part).issubset(partition_key_set): |
| partition_parts.append( |
| self._rebuild_leaf_indices_by_name(part, pk_to_idx)) |
| else: |
| data_parts.append(part) |
| return ( |
| PredicateBuilder.and_predicates(partition_parts), |
| PredicateBuilder.and_predicates(data_parts), |
| ) |
| |
| def _validate_and_rebuild_partition_filter(self, partition_filter): |
| 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))) |
| return self._rebuild_leaf_indices_by_name( |
| partition_filter, |
| {name: idx for idx, name in enumerate(partition_keys)}) |
| |
| def _add_partition_filter(self, partition_filter): |
| if self._partition_filter is None: |
| self._partition_filter = partition_filter |
| else: |
| self._partition_filter = PredicateBuilder.and_predicates( |
| [self._partition_filter, partition_filter]) |
| |
| @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 _ranked_row_ids(result): |
| score_getter = result.score_getter() |
| return sorted( |
| result.results(), |
| key=lambda row_id: (-(score_getter(row_id) or 0.0), row_id)) |
| |
| |
| def _top_k(scores, limit): |
| if not scores: |
| return ScoredGlobalIndexResult.create_empty() |
| heap = [] |
| for row_id, score in scores.items(): |
| item = (score, -row_id) |
| if len(heap) < limit: |
| heapq.heappush(heap, item) |
| elif item > heap[0]: |
| heapq.heapreplace(heap, item) |
| top_scores = {-neg_row_id: score for score, neg_row_id in heap} |
| return DictBasedScoredIndexResult(top_scores) |
