Function preprocess_data

catboost/python-package/ut/medium/test_with_polars.py:328–361

:return: (X, Xpl)

(
    X: List[object],
    Xpl: Union[pl.DataFrame, pl.Series],
    offset: int,
    multiple_chunks: bool
)

def preprocess_data(
    X: List[object],
    Xpl: Union[pl.DataFrame, pl.Series],
    offset: int,
    multiple_chunks: bool
) -> Tuple[List[object], Union[pl.DataFrame, pl.Series]]:
    """
    :return: (X, Xpl)
    """

    if offset:
        X = X[offset:]
        Xpl = Xpl.slice(offset=offset)
        if has_pyarrow:
            if isinstance(Xpl, pl.DataFrame):
                for column in Xpl:
                    assert column.to_arrow().offset == offset
            elif isinstance(Xpl, pl.Series):
                assert Xpl.to_arrow().offset == offset

    n_samples_after_offset = len(X)

    if multiple_chunks:
        cut_pos = n_samples_after_offset // 3
        Xpl1 = Xpl.slice(offset=0, length=cut_pos)
        Xpl2 = Xpl.slice(offset=cut_pos)
        Xpl = pl.concat([Xpl1, Xpl2])
        if isinstance(Xpl, pl.DataFrame):
            for column in Xpl:
                assert column.n_chunks() == 2
        elif isinstance(Xpl, pl.Series):
            assert Xpl.n_chunks() == 2

    return X, Xpl


def generate_group_sizes(n_samples: int, min_group_size: int, max_group_size: int, seed: int = 42) -> List[int]: