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 = DEFAULT_RANDOM_STATE,
        shuffle: bool = True,
        stratify_col: str | None = None,
    ):
        if not 0 < test_size < 1:
            raise ValueError(f"test_size must be between 0 and 1 (exclusive), got {test_size!r}.")
        if not 0 <= validation_size < 1:
            raise ValueError(
                f"validation_size must be between 0 (inclusive) and 1 (exclusive), "
                f"got {validation_size!r}."
            )
        if test_size + validation_size >= 1:
            raise ValueError(
                f"test_size ({test_size!r}) + validation_size ({validation_size!r}) must be "
                f"less than 1, otherwise there is no data left for training."
            )
        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: pd.DataFrame | SkyulfDataFrame, y: pd.Series | Any) -> SplitDataset:
        if is_polars(X):
            return self._split_xy_polars(cast(Any, X), y)

        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: pd.DataFrame | SkyulfDataFrame) -> SplitDataset:
        if is_polars(df):
            return self._split_polars(cast(Any, df))

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

    # ---- polars-native paths (index split + gather, no frame conversion) ----

    def _split_indices(self, n: int, stratify: Any) -> tuple[Any, Any]:
        """Split row positions ``0..n-1``; same partitioning as splitting rows."""
        return train_test_split(
            np.arange(n),
            test_size=self.test_size,
            random_state=self.random_state,
            shuffle=self.shuffle,
            stratify=stratify,
        )

    def _split_xy_polars(self, X: Any, y: Any) -> SplitDataset:
        stratify = _safe_stratify_polars(y, "Stratified split") if self.stratify_col else None

        tv_idx, test_idx = self._split_indices(X.height, stratify)

        validation = None
        train_idx = tv_idx
        if self.validation_size > 0:
            relative_val_size = self.validation_size / (1 - self.test_size)
            stratify_val = (
                _safe_stratify_polars(y.gather(tv_idx), "Stratified validation split")
                if stratify is not None and y is not None
                else None
            )
            train_idx, val_idx = train_test_split(
                tv_idx,
                test_size=relative_val_size,
                random_state=self.random_state,
                shuffle=self.shuffle,
                stratify=stratify_val,
            )
            validation = (X.gather(val_idx), y.gather(val_idx) if y is not None else None)

        return SplitDataset(
            train=(X.gather(train_idx), y.gather(train_idx) if y is not None else None),
            test=(X.gather(test_idx), y.gather(test_idx) if y is not None else None),
            validation=validation,
        )

    def _split_polars(self, df: Any) -> SplitDataset:
        stratify = self._frame_stratify_polars(df, label="Stratified split")

        tv_idx, test_idx = self._split_indices(df.height, stratify)

        validation = None
        train_idx = tv_idx
        if self.validation_size > 0:
            relative_val_size = self.validation_size / (1 - self.test_size)
            stratify_val = (
                _safe_stratify_polars(
                    df.get_column(self.stratify_col).gather(tv_idx), "Stratified validation split"
                )
                if stratify is not None
                else None
            )
            train_idx, val_idx = train_test_split(
                tv_idx,
                test_size=relative_val_size,
                random_state=self.random_state,
                shuffle=self.shuffle,
                stratify=stratify_val,
            )
            validation = df.gather(val_idx)

        return SplitDataset(
            train=df.gather(train_idx),
            test=df.gather(test_idx),
            validation=validation,
        )

    def _frame_stratify_polars(self, df: Any, label: str) -> Any:
        """Polars counterpart of :meth:`_frame_stratify`."""
        if not (self.stratify_col and self.stratify_col in df.columns):
            if self.stratify_col:
                logger.warning(
                    "%s requested but no target_column is configured for this "
                    "plain-DataFrame input, so there is no column to stratify on. "
                    "Stratification will be disabled.",
                    label,
                )
            return None
        return _safe_stratify_polars(df.get_column(self.stratify_col), label)

    # ---- 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):
            if self.stratify_col:
                logger.warning(
                    "%s requested but no target_column is configured for this "
                    "plain-DataFrame input, so there is no column to stratify on. "
                    "Stratification will be disabled.",
                    label,
                )
            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