data_handlers ¶

def load_csv(self, csv_path: str) -> pl.DataFrame:
    """Load the CSV file into a polars DataFrame.

    Args:
        csv_path (str): Path to the CSV file to load.

    Returns:
        pl.DataFrame: Polars DataFrame containing the loaded CSV data.
    """
    return pl.read_csv(csv_path)

def read_csv_header(self, csv_path: str) -> list:
    """Get the column names from the header of the CSV file.

    Args:
        csv_path (str): Path to the CSV file to read headers from.

    Returns:
        list: List of column names from the CSV header.
    """
    with open(csv_path) as f:
        return f.readline().strip().split(",")

def save(self, path: str) -> None:
    """Saves the data to a csv file."""
    self.data.write_csv(path)

def select_columns(self, columns: list) -> dict:
    """Select specific columns from the DataFrame and return as a dictionary.

    Args:
        columns (list): List of column names to select.

    Returns:
        dict: A dictionary where keys are column names and values are lists containing the column data.

    Example:
        >>> handler = DatasetHandler(...)
        >>> data_dict = handler.select_columns(["col1", "col2"])
        >>> # Returns {'col1': [1, 2, 3], 'col2': [4, 5, 6]}
    """
    df = self.data.select(columns)
    return {col: df[col].to_list() for col in columns}

def get_all_items(self) -> tuple[dict, dict, dict]:
    """Get the full dataset as three separate dictionaries for inputs, labels and metadata.

    Returns:
        tuple[dict, dict, dict]: Three dictionaries containing:
            - Input dictionary mapping input column names to encoded input data
            - Label dictionary mapping label column names to encoded label data
            - Meta dictionary mapping meta column names to meta data

    Example:
        >>> handler = DatasetHandler(...)
        >>> input_dict, label_dict, meta_dict = handler.get_dataset()
        >>> print(input_dict.keys())
        dict_keys(['age', 'fare'])
        >>> print(label_dict.keys())
        dict_keys(['survived'])
        >>> print(meta_dict.keys())
        dict_keys(['passenger_id'])
    """
    input_columns, label_columns, meta_columns = (
        self.dataset_manager.column_categories["input"],
        self.dataset_manager.column_categories["label"],
        self.dataset_manager.column_categories["meta"],
    )
    input_data = self.encoder_manager.encode_dataframe(self.data[input_columns])
    label_data = self.encoder_manager.encode_dataframe(self.data[label_columns])
    meta_data = {key: self.data[key].to_list() for key in meta_columns}
    return input_data, label_data, meta_data

def get_all_items_and_length(self) -> tuple[tuple[dict, dict, dict], int]:
    """Get the full dataset as three separate dictionaries for inputs, labels and metadata, and the length of the data."""
    return self.get_all_items(), len(self.data)

def load_csv(self, csv_path: str) -> pl.DataFrame:
    """Load the CSV file into a polars DataFrame.

    Args:
        csv_path (str): Path to the CSV file to load.

    Returns:
        pl.DataFrame: Polars DataFrame containing the loaded CSV data.
    """
    return pl.read_csv(csv_path)

def load_csv_per_split(self, csv_path: str, split: int) -> pl.DataFrame:
    """Load the part of csv file that has the specified split value.

    Split is a number that for 0 is train, 1 is validation, 2 is test.
    This is accessed through the column with category `split`. Example column name could be `split:split:int`.

    NOTE that the aim of having this function is that depending on the training, validation and test scenarios,
    we are gonna load only the relevant data for it.
    """
    if "split" not in self.columns:
        raise ValueError("The category split is not present in the csv file")
    if split not in [0, 1, 2]:
        raise ValueError(f"The split value should be 0, 1 or 2. The specified split value is {split}")
    return pl.scan_csv(csv_path).filter(pl.col("split") == split).collect()

def read_csv_header(self, csv_path: str) -> list:
    """Get the column names from the header of the CSV file.

    Args:
        csv_path (str): Path to the CSV file to read headers from.

    Returns:
        list: List of column names from the CSV header.
    """
    with open(csv_path) as f:
        return f.readline().strip().split(",")

def save(self, path: str) -> None:
    """Saves the data to a csv file."""
    self.data.write_csv(path)

def select_columns(self, columns: list) -> dict:
    """Select specific columns from the DataFrame and return as a dictionary.

    Args:
        columns (list): List of column names to select.

    Returns:
        dict: A dictionary where keys are column names and values are lists containing the column data.

    Example:
        >>> handler = DatasetHandler(...)
        >>> data_dict = handler.select_columns(["col1", "col2"])
        >>> # Returns {'col1': [1, 2, 3], 'col2': [4, 5, 6]}
    """
    df = self.data.select(columns)
    return {col: df[col].to_list() for col in columns}

def categorize_columns_by_type(self) -> dict:
    """Organizes columns from config into input, label, and meta categories.

    Reads the column definitions from the config and sorts them into categories
    based on their column_type field.

    Returns:
        dict: Dictionary containing lists of column names for each category:
            {
                "input": ["col1", "col2"],  # Input columns
                "label": ["target"],        # Label/output columns
                "meta": ["id"]     # Metadata columns
            }

    Example:
        >>> manager = DatasetManager("config.yaml")
        >>> categories = manager.categorize_columns_by_type()
        >>> print(categories)
        {
            'input': ['hello', 'bonjour'],
            'label': ['ciao'],
            'meta': ["id"]
        }
    """
    input_columns = []
    label_columns = []
    meta_columns = []
    for column in self.config.columns:
        if column.column_type == "input":
            input_columns.append(column.column_name)
        elif column.column_type == "label":
            label_columns.append(column.column_name)
        elif column.column_type == "meta":
            meta_columns.append(column.column_name)

    return {"input": input_columns, "label": label_columns, "meta": meta_columns}

def get_split_columns(self) -> list[str]:
    """Get the columns that are used for splitting."""
    return self.config.split.split_input_columns

def get_transform_logic(self) -> dict:
    """Get the transformation logic.

    Returns a dictionary in the following structure :
    {
        "transformation_name": str,
        "transformations": list[tuple[str, str, dict]]
    }
    """
    transformation_logic = {
        "transformation_name": self.config.transforms.transformation_name,
        "transformations": [],
    }
    for column in self.config.transforms.columns:
        for transformation in column.transformations:
            transformation_logic["transformations"].append(
                (column.column_name, transformation.name, transformation.params),
            )
    return transformation_logic

def add_split(self, split_manager: SplitManager, *, force: bool = False) -> None:
    """Add a column specifying the train, validation, test splits of the data.

    An error exception is raised if the split column is already present in the csv file. This behaviour can be overriden by setting force=True.

    Args:
        split_manager (SplitManager): Manager for handling dataset splitting
        force (bool): If True, the split column present in the csv file will be overwritten.
    """
    if ("split" in self.columns) and (not force):
        raise ValueError(
            "The category split is already present in the csv file. If you want to still use this function, set force=True",
        )
    # get relevant split columns from the dataset_manager
    split_columns = self.dataset_manager.get_split_columns()
    split_input_data = self.select_columns(split_columns)

    # get the split indices
    train, validation, test = split_manager.get_split_indices(split_input_data)

    # add the split column to the data
    split_column = np.full(len(self.data), -1).astype(int)
    split_column[train] = 0
    split_column[validation] = 1
    split_column[test] = 2
    self.data = self.data.with_columns(pl.Series("split", split_column))

    if "split" not in self.columns:
        self.columns.append("split")

def apply_transformation_group(self, transform_manager: TransformManager) -> None:
    """Apply the transformation group to the data."""
    for column_name, transform_name, _params in self.dataset_manager.get_transform_logic()["transformations"]:
        transformed_data, add_row = transform_manager.transform_column(
            column_name,
            transform_name,
            self.data[column_name],
        )
        if add_row:
            new_rows = self.data.with_columns(pl.Series(column_name, transformed_data))
            self.data = pl.vstack(self.data, new_rows)
        else:
            self.data = self.data.with_columns(pl.Series(column_name, transformed_data))

def load_csv(self, csv_path: str) -> pl.DataFrame:
    """Load the CSV file into a polars DataFrame.

    Args:
        csv_path (str): Path to the CSV file to load.

    Returns:
        pl.DataFrame: Polars DataFrame containing the loaded CSV data.
    """
    return pl.read_csv(csv_path)

def read_csv_header(self, csv_path: str) -> list:
    """Get the column names from the header of the CSV file.

    Args:
        csv_path (str): Path to the CSV file to read headers from.

    Returns:
        list: List of column names from the CSV header.
    """
    with open(csv_path) as f:
        return f.readline().strip().split(",")

def save(self, path: str) -> None:
    """Saves the data to a csv file."""
    self.data.write_csv(path)

def select_columns(self, columns: list) -> dict:
    """Select specific columns from the DataFrame and return as a dictionary.

    Args:
        columns (list): List of column names to select.

    Returns:
        dict: A dictionary where keys are column names and values are lists containing the column data.

    Example:
        >>> handler = DatasetHandler(...)
        >>> data_dict = handler.select_columns(["col1", "col2"])
        >>> # Returns {'col1': [1, 2, 3], 'col2': [4, 5, 6]}
    """
    df = self.data.select(columns)
    return {col: df[col].to_list() for col in columns}

def shuffle_labels(self, seed: Optional[float] = None) -> None:
    """Shuffles the labels in the data."""
    # set the np seed
    np.random.seed(seed)

    label_keys = self.dataset_manager.column_categories["label"]
    for key in label_keys:
        self.data = self.data.with_columns(pl.Series(key, np.random.permutation(list(self.data[key]))))

def encode_column(self, column_name: str, column_data: list) -> torch.Tensor:
    """Encodes a column of data using the configured encoder.

    Gets the appropriate encoder for the column from the encoder_loader and uses it
    to encode all the data in the column.

    Args:
        column_name: Name of the column to encode.
        column_data: List of data values from the column to encode.

    Returns:
        Encoded data as a torch.Tensor. The exact shape depends on the encoder used.

    Example:
        >>> data = ["ACGT", "TGCA"]
        >>> encoded = encode_manager.encode_column("dna_seq", data)
        >>> print(encoded.shape)
        torch.Size([2, 4, 4])  # 2 sequences, length 4, one-hot encoded
    """
    encode_all_function = self.encoder_loader.get_function_encode_all(column_name)
    return encode_all_function(column_data)

def encode_columns(self, column_data: dict) -> dict:
    """Encodes multiple columns of data using the configured encoders.

    Gets the appropriate encoder for each column from the encoder_loader and encodes
    all data values in those columns.

    Args:
        column_data: Dict mapping column names to lists of data values to encode.

    Returns:
        Dict mapping column names to their encoded tensors. The exact shape of each
        tensor depends on the encoder used for that column.

    Example:
        >>> data = {"dna_seq": ["ACGT", "TGCA"], "labels": ["1", "2"]}
        >>> encoded = encode_manager.encode_columns(data)
        >>> print(encoded["dna_seq"].shape)
        torch.Size([2, 4, 4])  # 2 sequences, length 4, one-hot encoded
    """
    return {col: self.encode_column(col, values) for col, values in column_data.items()}

def encode_dataframe(self, dataframe: pl.DataFrame) -> dict[str, torch.Tensor]:
    """Encode the dataframe using the encoders."""
    return {col: self.encode_column(col, dataframe[col].to_list()) for col in dataframe.columns}

def get_split_indices(self, data: dict) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Get the indices for train, validation, and test splits."""
    return self.split_loader.get_function_split()(data)

def transform_column(self, column_name: str, transform_name: str, column_data: list) -> tuple[list, bool]:
    """Transform a column of data using the specified transformation.

    Args:
        column_name (str): The name of the column to transform.
        transform_name (str): The name of the transformation to use.
        column_data (list): The data to transform.

    Returns:
        list: The transformed data.
        bool: Whether the transformation added new rows to the data.
    """
    transformer = self.transform_loader.__getattribute__(column_name)[transform_name]
    return transformer.transform_all(column_data), transformer.add_row