Dev/spark epic refactor and fix hypex functionality

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 = [

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),

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

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

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:

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.
 

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

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.