API: preprocessing.split

skyulf.preprocessing.split

Train/test/val splitter and feature/target splitter nodes.

These nodes return a :class:SplitDataset (not the canonical (X, y) shape), so they cannot use :func:apply_dual_engine — the dispatcher expects appliers that return (X_out, y_out). Instead, we centralise the engine handling in two small helpers:

• :func:_to_pandas_remember_engine converts to pandas and records whether a conversion happened.
• :func:_back_to_engine converts the result back to polars when needed.

This removes the inline if engine.name == EngineName.POLARS branches from both :class:DataSplitter methods and :class:FeatureTargetSplitApplier.

DataSplitter

Split a DataFrame (or X/y pair) into Train, Test, and optional Validation.

Source code in skyulf-core/skyulf/preprocessing/split.py

class DataSplitter:
    """Split a DataFrame (or X/y pair) into Train, Test, and optional Validation."""

    def __init__(
        self,
        test_size: float = 0.2,
        validation_size: float = 0.0,
        random_state: int = 42,
        shuffle: bool = True,
        stratify_col: Optional[str] = None,
    ):
        self.test_size = test_size
        self.validation_size = validation_size
        self.random_state = random_state
        self.shuffle = shuffle
        self.stratify_col = stratify_col

    # ---- public API ---------------------------------------------------------

    def split_xy(
        self, X: Union[pd.DataFrame, SkyulfDataFrame], y: Union[pd.Series, Any]
    ) -> SplitDataset:
        X_pd, was_polars = _to_pandas_remember_engine(X)
        y_pd, _ = _to_pandas_remember_engine(y)

        stratify = _safe_stratify(y_pd, "Stratified split") if self.stratify_col else None

        X_tv, X_test, y_tv, y_test = train_test_split(
            X_pd,
            y_pd,
            test_size=self.test_size,
            random_state=self.random_state,
            shuffle=self.shuffle,
            stratify=stratify,
        )

        validation, X_train, y_train = self._maybe_split_validation_xy(X_tv, y_tv)

        train = (_back_to_engine(X_train, was_polars), _back_to_engine(y_train, was_polars))
        test = (_back_to_engine(X_test, was_polars), _back_to_engine(y_test, was_polars))
        if validation is not None:
            validation = (
                _back_to_engine(validation[0], was_polars),
                _back_to_engine(validation[1], was_polars),
            )
        return SplitDataset(train=train, test=test, validation=validation)

    def split(self, df: Union[pd.DataFrame, SkyulfDataFrame]) -> SplitDataset:
        df_pd, was_polars = _to_pandas_remember_engine(df)
        stratify = self._frame_stratify(df_pd, label="Stratified split")

        train_val, test = train_test_split(
            df_pd,
            test_size=self.test_size,
            random_state=self.random_state,
            shuffle=self.shuffle,
            stratify=stratify,
        )

        validation, train = self._maybe_split_validation_frame(train_val)
        return SplitDataset(
            train=_back_to_engine(train, was_polars),
            test=_back_to_engine(test, was_polars),
            validation=_back_to_engine(validation, was_polars),
        )

    # ---- private helpers ----------------------------------------------------

    def _frame_stratify(self, df_pd: Any, label: str) -> Any:
        """Pick + sanity-check the stratify column on a frame split."""
        if not (self.stratify_col and self.stratify_col in df_pd.columns):
            return None
        return _safe_stratify(df_pd[self.stratify_col], label)

    def _maybe_split_validation_xy(self, X_tv: Any, y_tv: Any) -> Tuple[Any, Any, Any]:
        """Carve a validation set off of (X_tv, y_tv); returns (val, X_train, y_train)."""
        if self.validation_size <= 0:
            return None, X_tv, y_tv

        relative_val_size = self.validation_size / (1 - self.test_size)
        stratify_val = (
            _safe_stratify(y_tv, "Stratified validation split") if self.stratify_col else None
        )
        X_train, X_val, y_train, y_val = train_test_split(
            X_tv,
            y_tv,
            test_size=relative_val_size,
            random_state=self.random_state,
            shuffle=self.shuffle,
            stratify=stratify_val,
        )
        return (X_val, y_val), X_train, y_train

    def _maybe_split_validation_frame(self, train_val: Any) -> Tuple[Any, Any]:
        """Carve a validation set off of ``train_val`` (frame mode); returns (val, train)."""
        if self.validation_size <= 0:
            return None, train_val

        relative_val_size = self.validation_size / (1 - self.test_size)
        stratify_val = self._frame_stratify(train_val, label="Stratified validation split")
        train, val = train_test_split(
            train_val,
            test_size=relative_val_size,
            random_state=self.random_state,
            shuffle=self.shuffle,
            stratify=stratify_val,
        )
        return val, train