diff --git a/hypex/ab.py b/hypex/ab.py index be27ff63..9a14d8b3 100644 --- a/hypex/ab.py +++ b/hypex/ab.py @@ -76,14 +76,13 @@ def _make_experiment( additional_tests = ( [ABTestTypesEnum.t_test] if additional_tests is None else additional_tests ) - multitest_method = ( - ABNTestMethodsEnum(multitest_method) - if ( - multitest_method is not None - and multitest_method in ABNTestMethodsEnum.__members__.values() - ) - else ABNTestMethodsEnum.holm - ) + if ( + multitest_method is not None + and multitest_method in ABNTestMethodsEnum._value2member_map_ + ): + multitest_method = ABNTestMethodsEnum(multitest_method) + else: + multitest_method = ABNTestMethodsEnum.holm if additional_tests: if isinstance(additional_tests, list): additional_tests = [ diff --git a/hypex/analyzers/aa.py b/hypex/analyzers/aa.py index 33297afb..70cb02a3 100644 --- a/hypex/analyzers/aa.py +++ b/hypex/analyzers/aa.py @@ -83,7 +83,7 @@ class OneAAStatAnalyzer(Executor): score to evaluate the overall consistency of data splitting configurations. """ #: Registered test classes whose results are aggregated by OneAAStatAnalyzer. - ANALYSIS_TEST_CLASSES: ClassVar[tuple[type]] = tuple([ + ANALYSIS_TEST_CLASSES: ClassVar[tuple[type, ...]] = tuple([ GroupTTest, GroupKSTest, GroupChi2Test, @@ -95,7 +95,7 @@ class OneAAStatAnalyzer(Executor): #: (preferred_class, fallback_class, weight) for composite score computation. #: Preferred = Spark-backed (Stats*), fallback = Pandas-backed (Group*). - _SCORE_RULES: ClassVar[tuple[tuple[str, str, int]]] = tuple([ + _SCORE_RULES: ClassVar[tuple[tuple[str, str, int], ...]] = tuple([ ("StatsTTest", "GroupTTest", 1), ("StatsKSTest", "GroupKSTest", 2), ("StatsChi2Test", "GroupChi2Test", 2), diff --git a/hypex/analyzers/ab.py b/hypex/analyzers/ab.py index 9152a6ea..81aead29 100644 --- a/hypex/analyzers/ab.py +++ b/hypex/analyzers/ab.py @@ -92,24 +92,34 @@ def _set_value(self, data: ExperimentData, value, key=None) -> ExperimentData: def execute_multitest(self, data: ExperimentData, p_values: Dataset, **kwargs): """Applies multiple testing correction to aggregated p-values. - Retrieves treatment and target fields from the experiment data, then - applies the specified correction method if more than two groups exist. - For standard methods, uses ``MultiTest`` from statsmodels. For the - ``quantile`` method, uses simulation-based ``MultitestQuantile``. + Retrieves treatment and target fields from the experiment data and calculates + the total number of statistical comparisons being made. The correction is + applied if the total number of comparisons (calculated as + ``(num_groups - 1) * num_target_fields``) is strictly greater than 1. + + For standard correction methods (e.g., Bonferroni, Holm), it uses the + ``MultiTest`` extension wrapping ``statsmodels``. For the ``quantile`` method, + it uses the simulation-based ``MultitestQuantile`` extension. Args: - data: The experiment data container with group information. - p_values: Dataset containing raw p-values from statistical tests. - **kwargs: Additional arguments passed to the correction method. + data: The experiment data container holding dataset roles, groups, + and metadata. + p_values: A dataset containing the raw, uncorrected p-values to be + adjusted. + **kwargs: Extra keyword arguments forwarded to the underlying + multitest extensions. Returns: - Updated ``ExperimentData`` with corrected p-values stored under - the ``"MultiTest"`` key, or the original ``data`` if correction - is not applicable. + ExperimentData: The updated experiment data instance with the + multitest correction results stored in the analysis tables. """ group_field = data.ds.search_columns(TreatmentRole())[0] target_fields = data.ds.search_columns(TargetRole(), search_types=[int, float]) - if self.multitest_method and len(data.groups[group_field]) > 2: + + num_groups = len(data.groups[group_field]) + num_comparisons = (num_groups - 1) * len(target_fields) + + if self.multitest_method and num_comparisons > 1: if self.multitest_method != ABNTestMethodsEnum.quantile: multitest_result = MultiTest(self.multitest_method).calc( p_values, **kwargs diff --git a/hypex/comparators/abstract.py b/hypex/comparators/abstract.py index 068d86da..c9b094dc 100644 --- a/hypex/comparators/abstract.py +++ b/hypex/comparators/abstract.py @@ -739,18 +739,82 @@ def calc( group_col_stats: dict[str, dict[str, dict[str, Any]]] | None = None, **kwargs, ) -> dict: - """ - Stateless entry point mirroring :meth:`GroupsComparator.calc`, so the - comparator can be run outside the experiment pipeline. + """Stateless entry point mirroring :meth:`GroupsComparator.calc`. + + Runs the two-phase stats comparator outside the experiment pipeline. + + **Phase 1 – Aggregate.** When ``group_col_stats`` is not supplied, + this method builds a column projection from ``target_fields_data`` + (merging in ``group_field_data`` when the grouping column is absent) + and delegates to :meth:`_compute_stats`, which issues a **single** + backend aggregation job (one ``groupBy().agg()`` on Spark) for all + target columns simultaneously. + + **Phase 2 – Compare.** The pre-aggregated statistics are fed + pairwise (baseline vs. each compared group) into + :meth:`_inner_function`, which returns the test result + (``p-value``, ``statistic``, ``pass``) for every + ``(group, column)`` pair. - Pass either pre-aggregated ``group_col_stats`` (as produced by - :meth:`_compute_stats`) or the raw ``target_fields_data`` and - ``group_field_data`` to have the statistics aggregated here. ``stats`` - defaults to the comparator's ``REQUIRED_STATS``, so callers normally - don't need to supply it. + Two invocation modes are supported: - Returns ``{f"{group}{NAME_BORDER_SYMBOL}{col}": Dataset}`` pairwise test - results, comparing every non-baseline group against the first group. + * **Pre-aggregated** – pass ``group_col_stats`` directly (e.g. the + output of a previous :meth:`_compute_stats` call). No data + preparation is performed. + * **Raw data** – pass ``target_fields_data`` **and** + ``group_field_data``. Statistics are aggregated internally. + + Args: + target_fields_data: Dataset containing the target metric + columns to compare. Required when ``group_col_stats`` + is ``None``. + group_field_data: Single-column Dataset that defines group + membership (e.g. treatment assignment). Required when + ``group_col_stats`` is ``None``. + baseline_fields_data: Dataset with match-index columns. + Used only when ``compare_by="matched_pairs"``; ignored + for ``compare_by="groups"``. + stats: Statistic names to compute in Phase 1 + (e.g. ``["mean", "std", "count"]``). When ``None``, + defaults to the class-level ``REQUIRED_STATS``. + compare_by: Comparison mode. Supported values: + + * ``"groups"`` – standard multi-group comparison. + The first (alphabetically smallest) group is treated + as the baseline. + * ``"matched_pairs"`` – each observation is compared + against its matched counterpart. Requires + ``baseline_fields_data``. + + group_col_stats: Pre-aggregated statistics in the nested-dict + format produced by :meth:`_compute_stats`:: + + {group_key: {column: {stat_name: value}}} + + When provided, ``target_fields_data`` and + ``group_field_data`` are ignored. + **kwargs: Additional keyword arguments forwarded to + :meth:`_inner_function` (e.g. ``reliability``). + + Returns: + dict[str, Dataset]: Pairwise test results keyed by + ``f"{group}{NAME_BORDER_SYMBOL}{col}"``. Every non-baseline + group is compared against the first (baseline) group. + + Raises: + ValueError: If neither ``group_col_stats`` nor both + ``target_fields_data`` and ``group_field_data`` are + provided, or if ``compare_by`` is not one of the + supported modes. + + Example: + >>> from hypex.comparators.stats_hypothesis_testing import StatsTTest + >>> result = StatsTTest.calc( + ... target_fields_data=ds_metrics, + ... group_field_data=ds_treat, + ... ) + >>> for key, ds in result.items(): + ... print(key, ds.get_values(row="p-value", column="p-value")) """ if group_col_stats is None: if target_fields_data is None or group_field_data is None: @@ -759,15 +823,67 @@ def calc( "target_fields_data and group_field_data." ) - grouped = cls._prepare_data(compare_by, target_fields_data, group_field_data, baseline_fields_data) - group_col_stats = cls._compute_stats( - grouped, list(target_fields_data.columns), stats or cls.REQUIRED_STATS - ) + group_col = group_field_data.columns[0] + target_cols = list(target_fields_data.columns) + + if compare_by == "groups": + if group_col in target_fields_data.columns: + agg_data = target_fields_data + else: + agg_data = target_fields_data.merge( + group_field_data, left_index=True, right_index=True + ) + + group_col_stats = cls._compute_stats( + data=agg_data, + group_cols=[group_col], + target_columns=target_cols, + stats=stats or cls.REQUIRED_STATS, + ) + + elif compare_by == "matched_pairs": + best_match_col = baseline_fields_data.columns[0] + baseline_fields = baseline_fields_data[best_match_col] + + tmp_data = group_field_data.merge( + right=target_fields_data, left_index=True, right_index=True + ) + tmp_data = tmp_data.merge( + right=baseline_fields, right_index=True, left_index=True + ) + tmp_data = tmp_data.merge( + right=tmp_data, right_index=True, left_on=best_match_col, + suffixes=("", "_matched"), + ) + prepared_data = tmp_data.drop( + columns=[ + best_match_col, + best_match_col + "_matched", + group_col + "_matched", + ] + ) + + matched_target_cols = target_cols + [ + f"{c}_matched" for c in target_cols + ] + + group_col_stats = cls._compute_stats( + data=prepared_data, + group_cols=[group_col], + target_columns=matched_target_cols, + stats=stats or cls.REQUIRED_STATS, + ) + + else: + raise ValueError( + f"StatsComparator supports 'groups' and 'matched_pairs' only, " + f"got compare_by={compare_by!r}" + ) return cls._execute_inner_function( group_col_stats=group_col_stats, compare_by=compare_by, - **kwargs + **kwargs, ) @classmethod diff --git a/hypex/dataset/abstract.py b/hypex/dataset/abstract.py index 8b6af725..e26c4f4a 100644 --- a/hypex/dataset/abstract.py +++ b/hypex/dataset/abstract.py @@ -1,43 +1,38 @@ from __future__ import annotations -import warnings -import copy -from copy import deepcopy import json -from abc import ABC +import warnings from collections.abc import Iterable as IterableABC +from copy import deepcopy from dataclasses import dataclass -from typing import Any, Iterable, Callable, Hashable, Literal, Optional, Sequence +from typing import Any, Callable, Hashable, Iterable, Literal, Sequence try: from typing import Self # Python >= 3.11 except ImportError: - from typing_extensions import Self # Python < 3.11 + from typing_extensions import Self # type: ignore # Python < 3.11 import pandas as pd # type: ignore +import pyspark.pandas as ps # type: ignore +import pyspark.sql as spark # type: ignore from numpy import ndarray -import pyspark.sql as spark -import pyspark.pandas as ps - -from .backends import PandasDataset, SparkDataset -from .groupby_dataset import GroupedDataset +from ..config import DatasetConfig from ..utils import ( BackendsEnum, BackendTypeError, ConcatBackendError, ConcatDataError, DataTypeError, + GenericManager, RoleColumnError, ScalarType, SourceDataTypes, - GenericManager ) -from ..config import DatasetConfig from ..utils.adapter import Adapter -from .groupby_dataset import GroupedDataset from .backends import PandasDataset, SparkDataset +from .groupby_dataset import GroupedDataset from .roles import ( ABCRole, DefaultRole, @@ -180,20 +175,14 @@ def _set_all_roles(self, roles: dict[str, ABCRole]) -> dict[str, ABCRole]: return roles def _set_empty_types(self, roles): - colunms_dtypes = self._backend_data.get_column_type(self._backend_data.columns) + columns_dtypes = self._backend_data.get_column_type(self._backend_data.columns) new_types = {} for column, role in roles.items(): if role.data_type is None: - role.data_type = colunms_dtypes[column] - elif role.data_type != colunms_dtypes[column]: + role.data_type = columns_dtypes[column] + elif role.data_type != columns_dtypes[column]: new_types[column] = role.data_type - - if new_types: - for c in new_types: - try: - na = int(self._backend_data.data[c].isna().sum()) - except Exception: - na = -1 + self._backend_data = self._backend_data.update_column_type(new_types) def __init__( @@ -202,7 +191,7 @@ def __init__( data: spark.DataFrame | pd.DataFrame | str | Self | None = None, backend: BackendsEnum | None = None, default_role: ABCRole | None = None, - session: Optional[spark.SparkSession] = None, + session: spark.SparkSession | None = None, data_compression: Literal[ "downcasting", "encoding", "auto", "disable" ] = "auto", @@ -214,7 +203,7 @@ def __init__( ) elif data is not None: if isinstance(data, DatasetBase): - self._backend_data = deepcopy(data._backend) + self._backend_data = deepcopy(data._backend_data) elif isinstance(data, (PandasDataset, SparkDataset)): self._backend_data = data elif any( @@ -234,7 +223,7 @@ def __init__( self._backend_data = PandasDataset(data) self.default_role = default_role - if roles is None and data.hasattr("roles") and data.roles is not None: + if roles is None and hasattr(data, "roles") and data.roles is not None: roles = data.roles elif roles is None: roles = {} @@ -638,7 +627,7 @@ def __round__(self, ndigits: int = 0) -> Self: session=self.session, ) - def __bool__(self) -> Self: + def __bool__(self) -> bool: return not self._backend_data.is_empty() # Binary math operators: diff --git a/hypex/extensions/faiss.py b/hypex/extensions/faiss.py index b6836492..6de38e31 100644 --- a/hypex/extensions/faiss.py +++ b/hypex/extensions/faiss.py @@ -235,7 +235,7 @@ def _predict( data: Dataset, test_data: Dataset, X: np.ndarray - ) -> pd.Series: + ) -> Dataset: """ Perform the FAISS search on the query vectors. diff --git a/hypex/extensions/scipy_stats.py b/hypex/extensions/scipy_stats.py index 3657528d..4ff1ffe6 100644 --- a/hypex/extensions/scipy_stats.py +++ b/hypex/extensions/scipy_stats.py @@ -56,7 +56,7 @@ def check_data(self, data: Dataset, other: Dataset | None) -> Dataset: return other - def _extract_arrays(self, data: Dataset, other: Dataset) -> tuple[Sequence]: + def _extract_arrays(self, data: Dataset, other: Dataset) -> tuple[Sequence, ...]: raise NotImplementedError("This method should be relized using backend-dependent mixin.") @staticmethod diff --git a/hypex/ml/faiss.py b/hypex/ml/faiss.py index 296b68de..7452da7e 100644 --- a/hypex/ml/faiss.py +++ b/hypex/ml/faiss.py @@ -73,7 +73,7 @@ def __init__( @classmethod def _set_global_match_indexes( cls, local_indexes: Dataset, data: tuple[str, Dataset] - ) -> list[int, list[int]]: + ) -> list[int] | list[list[int]]: """ Map local group indexes to global dataset indexes. diff --git a/hypex/reporters/__init__.py b/hypex/reporters/__init__.py index 5e6fd544..f3f210c2 100644 --- a/hypex/reporters/__init__.py +++ b/hypex/reporters/__init__.py @@ -1,20 +1,40 @@ -from .abstract import Reporter, DictReporter, TestDictReporter, DatasetReporter, ResultKey, REPORTABLE_METRICS -from .aa import AATestReporter, AAPassedReporter, AABestSplitReporter +from .aa import AABestSplitReporter, AAPassedReporter, AATestReporter from .ab import ABTestReporter, CupacReporter +from .abstract import ( + REPORTABLE_METRICS, + DatasetReporter, + DictReporter, + Reporter, + ResultKey, + TestDictReporter, +) from .homo import HomogeneityReporter -from .matching import MatchingReporter, MatchingQualityReporter +from .matching import MatchingQualityReporter, MatchingReporter __all__ = [ - "Reporter", "DictReporter", "TestDictReporter", "DatasetReporter", - "ResultKey", "REPORTABLE_METRICS", - "AATestReporter", "AAPassedReporter", "AABestSplitReporter", - "ABTestReporter", "CupacReporter", + "REPORTABLE_METRICS", + "AABestSplitReporter", + "AADatasetReporter", + "AAPassedReporter", + "AATestReporter", + "ABDatasetReporter", + "ABDictReporter", + "ABTestReporter", + "CupacReporter", + "DatasetReporter", + "DictReporter", + "HomoDatasetReporter", + "HomoDictReporter", "HomogeneityReporter", - "MatchingReporter", "MatchingQualityReporter", + "MatchingDatasetReporter", + "MatchingDictReporter", + "MatchingQualityDatasetReporter", + "MatchingQualityDictReporter", + "MatchingQualityReporter", + "MatchingReporter", # Backwards compat - "OneAADictReporter", "AADatasetReporter", - "ABDictReporter", "ABDatasetReporter", - "HomoDictReporter", "HomoDatasetReporter", - "MatchingDictReporter", "MatchingQualityDictReporter", - "MatchingDatasetReporter", "MatchingQualityDatasetReporter" + "OneAADictReporter", + "Reporter", + "ResultKey", + "TestDictReporter" ] \ No newline at end of file diff --git a/hypex/reporters/ab.py b/hypex/reporters/ab.py index 961b3da5..5ce514dc 100644 --- a/hypex/reporters/ab.py +++ b/hypex/reporters/ab.py @@ -20,23 +20,6 @@ class ABTestReporter(DatasetReporter): tests: ClassVar[list[type[BaseComparator]]] = [GroupTTest, GroupUTest, GroupChi2Test, StatsTTest, StatsChi2Test] - def _report(self, data: ExperimentData) -> dict[str, Any]: - """Construct the internal dictionary report for A/B tests. - - Args: - data: The experiment data container. - - Returns: - A dictionary containing group sizes, differences, test results, - and analyzer metrics. - """ - result = {} - result.update(extract_group_sizes(data, self.front)) - result.update(extract_group_difference(data, self.front)) - result.update(extract_tests(data, self.tests, self.front)) - result.update(extract_analyzer_data(data, ABAnalyzer)) - return result - def _report(self, data: ExperimentData) -> dict[str, Any]: """Generate the final A/B test report. diff --git a/hypex/reporters/matching.py b/hypex/reporters/matching.py index 42619437..b0697fb8 100644 --- a/hypex/reporters/matching.py +++ b/hypex/reporters/matching.py @@ -4,7 +4,15 @@ from typing import Any, ClassVar, Literal from ..analyzers.matching import MatchingAnalyzer -from ..comparators import GroupChi2Test, GroupKSTest, GroupTTest +from ..comparators import ( + BaseComparator, + GroupChi2Test, + GroupKSTest, + GroupTTest, + StatsChi2Test, + StatsKSTest, + StatsTTest, +) from ..dataset import Dataset, ExperimentData from ..ml import FaissNearestNeighbors from ..utils import ( @@ -106,7 +114,10 @@ def _extract_indexes(self, data: ExperimentData) -> dict[str, str]: class MatchingQualityReporter(DatasetReporter): """Reporter for matching quality tests (T-Test, KS-Test, Chi2-Test).""" - tests: ClassVar[list] = [GroupTTest, GroupKSTest, GroupChi2Test] + tests: ClassVar[list[type[BaseComparator]]] = [ + GroupTTest, GroupKSTest, GroupChi2Test, + StatsTTest, StatsKSTest, StatsChi2Test + ] def _report(self, data: ExperimentData) -> dict: """Extract quality test outcomes. diff --git a/hypex/ui/base.py b/hypex/ui/base.py index f6804d4c..e123e325 100644 --- a/hypex/ui/base.py +++ b/hypex/ui/base.py @@ -5,7 +5,7 @@ from ..dataset import Dataset, ExperimentData from ..experiments.base import Experiment from ..reporters import Reporter -from ..utils import ID_SPLIT_SYMBOL +from ..utils import ID_SPLIT_SYMBOL, BackendsEnum from ..utils.enums import RenameEnum @@ -148,45 +148,86 @@ def __init__( experiment: Experiment, output: Output, experiment_params: dict[str, Any] | None = None, + auto_persist: bool = True, ): if experiment_params: experiment.set_params(experiment_params) self._out = output self._experiment = experiment + self.auto_persist = auto_persist - @property - def experiment(self): - """Gets the configured experiment instance. + def execute(self, data: Dataset | ExperimentData) -> Output: + """Execute the experiment pipeline on the provided data. - Returns: - Experiment: The experiment configuration object. - """ - return self._experiment + Orchestrates the full experiment lifecycle: data preparation, optional + caching for Spark backends, pipeline execution, and result extraction. - def execute(self, data: Dataset | ExperimentData) -> Output: - """Executes the experiment on the provided data. + **Auto-persist behaviour (Spark only):** + When ``auto_persist`` is enabled (default) and the input dataset uses + the Spark backend, the method automatically persists the dataset with + ``MEMORY_AND_DISK`` storage level before the pipeline starts. This + avoids costly recomputation of the source DataFrame across multiple + stages (splitters, comparators, analyzers). After the pipeline + completes, the dataset is unpersisted **only** if it was persisted by + this method — datasets that the user cached manually are left untouched. - Runs the configured experiment on the input data and formats the results - using the configured output handler. + For the Pandas backend, ``persist`` / ``unpersist`` are no-ops, so + the method behaves identically regardless of backend. Args: - data (Union[Dataset, ExperimentData]): Input data for the experiment. - Can be either a Dataset or ExperimentData instance. + data: Input data for the experiment. Accepts either a raw + :class:`~hypex.dataset.Dataset` (which will be wrapped in an + :class:`~hypex.dataset.ExperimentData` container) or an + already-prepared :class:`~hypex.dataset.ExperimentData` instance. Returns: - Output: Formatted experiment results through the configured output handler. + Output: The experiment output object containing the formatted + results (resume, multitest table, quality reports, etc.), + populated by the configured :class:`Output` handler. - Examples - -------- - .. code-block:: python + Example: + .. code-block:: python + + ab_test = ABTest(multitest_method="bonferroni") + result = ab_test.execute(spark_dataset) + print(result.resume) + print(result.multitest) - shell = ExperimentShell(experiment, output) - dataset = Dataset(...) # Your input data - results = shell.execute(dataset) - print(results.resume) # Access formatted results + See Also: + :meth:`Dataset.persist`: Manual caching control. + :class:`ExperimentShell`: Constructor accepting ``auto_persist`` flag. """ if isinstance(data, Dataset): data = ExperimentData(data) + + # ── Auto-persist for Spark backend ────────────────────────── + original_ds = data.ds + persisted_by_us = False + if ( + self.auto_persist + and original_ds.backend_type == BackendsEnum.spark + and not original_ds.is_persisted + ): + original_ds.persist( + storage_level="MEMORY_AND_DISK", action="count" + ) + persisted_by_us = True + # ───────────────────────────────────────────────────────────── + result_experiment_data = self._experiment.execute(data) self._out.extract(result_experiment_data) + + # Unpersist only if WE persisted it (not the user) + if persisted_by_us and original_ds.is_persisted: + original_ds.unpersist() + return self._out + + @property + def experiment(self): + """Gets the configured experiment instance. + + Returns: + Experiment: The experiment configuration object. + """ + return self._experiment diff --git a/hypex/ui/matching.py b/hypex/ui/matching.py index 30f11acb..68b9ecd1 100644 --- a/hypex/ui/matching.py +++ b/hypex/ui/matching.py @@ -42,18 +42,17 @@ def __init__(self, searching_class: type = MatchingAnalyzer): def _extract_full_data(self, experiment_data: ExperimentData, indexes: Dataset): """Build the full matched dataset from original data and matched indexes. - Materialises the dataset index as a plain Python list so that it can be safely assigned to Pandas-backed SmallDataset objects regardless of the backend used by ``experiment_data.ds`` (Pandas or Spark). - Args: experiment_data: The experiment data container. indexes: Dataset containing matched neighbor indexes per group. """ # ── Convert to list to avoid PySpark Index → Pandas assignment error ── ds_index = experiment_data.ds.index.to_numpy().tolist() - self.indexes = Dataset(roles={}, data=experiment_data.ds.index) + + self.indexes = SmallDataset.create_empty(roles={}) for i in range(len(indexes.columns)): t_indexes = indexes.iloc[:, i] @@ -61,10 +60,8 @@ def _extract_full_data(self, experiment_data: ExperimentData, indexes: Dataset): filtered_field = indexes.drop( indexes[indexes[t_indexes.columns[0]] == -1], axis=0 ) - lookup_vals = list(map(lambda x: x[0], filtered_field.get_values())) filtered_index_list = filtered_field.index.to_numpy().tolist() - if experiment_data.ds.backend_type == BackendsEnum.spark: mapping_ds = Dataset( roles={ @@ -81,7 +78,6 @@ def _extract_full_data(self, experiment_data: ExperimentData, indexes: Dataset): orig_cols = set(experiment_data.ds.columns) ds_reset = experiment_data.ds.reset_index() idx_col = next(c for c in ds_reset.columns if c not in orig_cols) - matched_data = mapping_ds.merge( ds_reset, left_on="_hypex_lookup", @@ -93,11 +89,9 @@ def _extract_full_data(self, experiment_data: ExperimentData, indexes: Dataset): else: matched_data = experiment_data.ds.loc[lookup_vals] matched_data.index = filtered_index_list - matched_data = matched_data.rename( {col: f"{col}_matched_{i}" for col in matched_data.columns} ) - reindexed_matched = experiment_data.ds.merge( matched_data, left_index=True, @@ -108,11 +102,13 @@ def _extract_full_data(self, experiment_data: ExperimentData, indexes: Dataset): columns=list(experiment_data.ds.columns) ) - self.indexes = ( - t_indexes - if self.indexes.is_empty() - else self.indexes.add_column(t_indexes) - ) + if self.indexes.is_empty(): + self.indexes = t_indexes + else: + self.indexes.add_column( + data=t_indexes.data, + role={col: t_indexes.roles.get(col, InfoRole()) for col in t_indexes.columns} + ) if hasattr(self, "full_data") and self.full_data is not None: self.full_data = self.full_data.merge( diff --git a/hypex/utils/index_utils.py b/hypex/utils/index_utils.py index bd47958d..e41b4d26 100644 --- a/hypex/utils/index_utils.py +++ b/hypex/utils/index_utils.py @@ -1,3 +1,15 @@ +"""Utilities for FAISS index storage and caching in distributed Spark environments. + +This module provides two core components: + +- :class:`FaissIndexStorage`: Manages serialization, persistence, and distribution + of FAISS indexes across Spark executors using the local file system and + ``SparkFiles``. +- :class:`CachingIndex`: A thread-safe LRU cache that prevents redundant + deserialization of FAISS indexes when processing multiple query batches + on the same executor. +""" + from __future__ import annotations import gc @@ -6,149 +18,202 @@ import threading import uuid from collections import OrderedDict -from typing import Optional +from typing import Any, ClassVar -import faiss -from pyspark import RDD, SparkFiles -from pyspark.sql import SparkSession +import faiss # pyright: ignore[reportMissingImports] +from pyspark import RDD, SparkFiles # pyright: ignore[reportMissingImports] +from pyspark.sql import SparkSession # pyright: ignore[reportMissingImports] class FaissIndexStorage: - """ - Faiss file manager. + """Manages FAISS index files for distributed Spark processing. - Simplified version: only the local file system is used. - and the Spark file distribution mechanism (SparkFiles.addFile). + Uses only the local file system combined with Spark's file distribution + mechanism (``SparkFiles.addFile``) to avoid connectivity issues with + distributed file systems (HDFS/viewfs/WebHDFS) on corporate clusters. - This avoids problems with connecting to distributed FS - (HDFS/viewfs/WebHDFS) on corporate clusters. + Attributes: + DISTRIBUTED_DIRS: Retained for backward compatibility. Always empty. + LOCAL_DIRS: List of temporary local directories created by instances + of this class. Used by :meth:`cleanup` for resource release. + + Args: + sp_s: Active Spark session used for file distribution. + base_dir: Retained for backward compatibility. Not used. + + Example: + >>> storage = FaissIndexStorage(spark_session) + >>> refs = storage.collect_and_register(sharded_rdd) + >>> index = storage.load_index(refs[0]) # on executor """ - # DISTRIBUTED_DIRS оставлен для совместимости - DISTRIBUTED_DIRS = [] - LOCAL_DIRS = [] + + DISTRIBUTED_DIRS: ClassVar[list[str]] = [] + LOCAL_DIRS: ClassVar[list[str]] = [] def __init__( self, sp_s: SparkSession, - base_dir: str | None = None # сохранён для совместимости, не используется - ): - self.sp_s = sp_s - # для совместимости с внешним кодом - self._distributed = False - self._distributed_dir = None - - dir_id = uuid.uuid1().hex[:8] - self._local_tmp_dir = f"__partition_indexes_{dir_id}" + base_dir: str | None = None, + ) -> None: + self.sp_s: SparkSession = sp_s + # Retained for backward compatibility with external code. + self._distributed: bool = False + self._distributed_dir: str | None = None + + dir_id: str = uuid.uuid1().hex[:8] + self._local_tmp_dir: str = f"__partition_indexes_{dir_id}" os.makedirs(self._local_tmp_dir, exist_ok=True) FaissIndexStorage.LOCAL_DIRS.append(self._local_tmp_dir) - def __getstate__(self): - state = self.__dict__.copy() - if "sp_s" in state: - del state["sp_s"] + def __getstate__(self) -> dict[str, Any]: + """Exclude the non-serializable SparkSession before pickling. + + Returns: + Instance state dictionary without the ``sp_s`` key. + """ + state: dict[str, Any] = self.__dict__.copy() + state.pop("sp_s", None) return state def save_index(self, index: faiss.Index) -> bytes: - """ - Serialize index to bytes for sending to executors. - Raisses on executors. - - Args - ---- - index: faiss.Index - partition faiss index. - Returns - ------- - serialized index in bytes. + """Serialize a FAISS index to bytes for transmission to executors. + + This method is invoked on executors during the distributed fit phase. + + Args: + index: The partition-level FAISS index to serialize. + + Returns: + Serialized index as a bytes object. """ return faiss.serialize_index(index) - def collect_and_register(self, rdd: RDD) -> list: - """ - Collects serialized indexes from executors, - saves them to temporary files and distributes them via SparkFiles. - It is called on the driver. - - Args - ---- - rdd: RDD - ... - - Return - ------ - list of index references. + def collect_and_register(self, rdd: RDD) -> list[str]: + """Collect serialized indexes from executors and distribute via SparkFiles. + + Iterates over the RDD produced by the distributed fit phase, + deserializes each partition index, writes it to a temporary local + file, and registers the file with ``SparkFiles`` so that executors + can retrieve it during the predict phase. + + This method must be called on the driver. + + Args: + rdd: RDD containing serialized FAISS indexes (one per partition). + + Returns: + List of index file reference names suitable for + :meth:`load_index`. """ - index_refs = [] + index_refs: list[str] = [] for shard in rdd.toLocalIterator(): - partition_indexes = faiss.deserialize_index(shard) - run_id = uuid.uuid1().hex[:8] - index_file_name = f"__partition_index_{run_id}.index" + partition_index: faiss.Index = faiss.deserialize_index(shard) + run_id: str = uuid.uuid1().hex[:8] + index_file_name: str = f"__partition_index_{run_id}.index" faiss.write_index( - partition_indexes, + partition_index, + f"{self._local_tmp_dir}/{index_file_name}", + ) + self.sp_s.sparkContext.addFile( f"{self._local_tmp_dir}/{index_file_name}" ) - self.sp_s.sparkContext.addFile(f"{self._local_tmp_dir}/{index_file_name}") index_refs.append(index_file_name) - del partition_indexes # явное освобождение памяти + del partition_index gc.collect() return index_refs def load_index(self, link: str) -> faiss.Index: - """ - Loads the index distributed through SparkFiles. - It is invoked on executors. + """Load a FAISS index distributed through SparkFiles. - Args - ---- - link: str - link to index file. + This method is invoked on executors during the predict phase. - Return - ------ - faiss index loaded from file. + Args: + link: File reference name returned by :meth:`collect_and_register`. + + Returns: + The deserialized FAISS index. """ return faiss.read_index(SparkFiles.get(link)) @staticmethod - def cleanup(): - """Deletes all temporary local directories.""" - for dir in FaissIndexStorage.LOCAL_DIRS: - if os.path.exists(dir): + def cleanup() -> None: + """Delete all temporary local directories created by FaissIndexStorage. + + Iterates over :attr:`LOCAL_DIRS`, removes each directory if it + exists, and resets the list. Exceptions during removal are + silently suppressed. + """ + for directory in FaissIndexStorage.LOCAL_DIRS: + if os.path.exists(directory): try: - shutil.rmtree(dir) + shutil.rmtree(directory) except Exception: pass FaissIndexStorage.LOCAL_DIRS = [] class CachingIndex: + """Thread-safe LRU cache for FAISS indexes on Spark executors. + + Prevents repeated deserialization of the same index file when + processing multiple query batches within a single executor process. + The cache is stored as a module-level singleton via + :func:`hypex.extensions.faiss.get_executor_cache`. + + Args: + max_index: Maximum number of indexes to hold in cache. + When the limit is reached, the least-recently-used entry + is evicted. If ``None``, the cache grows without bound. + + Example: + >>> cache = CachingIndex(max_index=4) + >>> index = cache.get("partition_index_0.index", storage, nprobe=8) """ - LRU is the cache of FAISS indexes in the executor's memory. - Prevents repeated loading of the same index - when processing multiple batches. - """ - def __init__(self, max_index: Optional[int] = None): - self._max = max_index - self._cache: OrderedDict = OrderedDict() - self._lock = threading.Lock() + + def __init__(self, max_index: int | None = None) -> None: + self._max: int | None = max_index + self._cache: OrderedDict[str, faiss.Index] = OrderedDict() + self._lock: threading.Lock = threading.Lock() def get( self, reference: str, storage: FaissIndexStorage, nprobe: int, - ): + ) -> faiss.Index: + """Retrieve a FAISS index from cache or load it from storage. + + If the index identified by ``reference`` is already cached, it is + moved to the most-recently-used position and returned immediately. + Otherwise, the index is loaded via ``storage.load_index``, configured + with the given ``nprobe``, and inserted into the cache. If the cache + has reached its capacity limit, the least-recently-used entry is + evicted first. + + Args: + reference: File reference name for the index. + storage: The :class:`FaissIndexStorage` instance used to load + the index if it is not cached. + nprobe: Number of IVF clusters to probe during search. + Applied to the inner index if it supports the attribute. + + Returns: + The loaded or cached FAISS index ready for search. + """ with self._lock: if reference in self._cache: self._cache.move_to_end(key=reference) return self._cache[reference] - if self._max and len(self._cache) >= self._max: + + if self._max is not None and len(self._cache) >= self._max: _, evicted = self._cache.popitem(last=False) del evicted gc.collect() - tmp_index = storage.load_index(reference) - inner = faiss.downcast_index(tmp_index) + + tmp_index: faiss.Index = storage.load_index(reference) + inner: faiss.Index = faiss.downcast_index(tmp_index) if hasattr(inner, "nprobe"): inner.nprobe = nprobe + self._cache[reference] = tmp_index - return tmp_index \ No newline at end of file + return tmp_index