dqm_ml_job.outputwriter

Output writers module for DQM ML Job.

This module contains classes for writing pipeline results (features and metrics) to various storage backends.

Classes:

Name	Description
OutputWriter	Protocol for output writer implementations.
ParquetOutputWriter	Writer that saves data to Parquet files.

__all__ = ['OutputWriter', 'ParquetOutputWriter', 'dqml_outputs_registry'] module-attribute

dqml_outputs_registry = {'parquet': ParquetOutputWriter} module-attribute

OutputWriter

Bases: Protocol

Protocol for Output Writers.

Defines the interface for writing pipeline results (features or metrics) to storage.

Source code in packages/dqm-ml-job/src/dqm_ml_job/outputwriter/__init__.py

@runtime_checkable
class OutputWriter(Protocol):
    """
    Protocol for Output Writers.

    Defines the interface for writing pipeline results (features or metrics) to storage.
    """

    columns: list[str]
    name: str

    def write_metrics_dict(self, metrics_dict: dict[str, dict[str, Any]]) -> None:
        """ """

    def write_table(self, name: str, table: Any, part_index: int | None = None) -> None:
        """
        Write a table (features or metrics) to the output.

        Args:
            name: Name of the dataset or metric.
            table: The data to write (usually a pyarrow Table or dict of arrays).
            part_index: Index of the data part (for chunked writing).
        """

columns: list[str] instance-attribute

name: str instance-attribute

write_metrics_dict(metrics_dict: dict[str, dict[str, Any]]) -> None

Source code in packages/dqm-ml-job/src/dqm_ml_job/outputwriter/__init__.py

def write_metrics_dict(self, metrics_dict: dict[str, dict[str, Any]]) -> None:
    """ """

write_table(name: str, table: Any, part_index: int | None = None) -> None

Write a table (features or metrics) to the output.

Parameters:

Name	Type	Description	Default
name	str	Name of the dataset or metric.	required
table	Any	The data to write (usually a pyarrow Table or dict of arrays).	required
part_index	int | None	Index of the data part (for chunked writing).	None

Source code in packages/dqm-ml-job/src/dqm_ml_job/outputwriter/__init__.py

def write_table(self, name: str, table: Any, part_index: int | None = None) -> None:
    """
    Write a table (features or metrics) to the output.

    Args:
        name: Name of the dataset or metric.
        table: The data to write (usually a pyarrow Table or dict of arrays).
        part_index: Index of the data part (for chunked writing).
    """

ParquetOutputWriter

Output writer that saves processed features to a Parquet file.

Source code in packages/dqm-ml-job/src/dqm_ml_job/outputwriter/parquet.py

class ParquetOutputWriter:
    """
    Output writer that saves processed features to a Parquet file.
    """

    def __init__(self, name: str, config: dict[str, Any] | None = None):
        """
        Initialize a ParquetOutputWriter.

        Args:
            name: Unique name for this output writer.
            config: Configuration dictionary with keys:
                - path_pattern (str): Output file path format string.
                - columns (List[str]): Columns to save.

        Raises:
            ValueError: If required config keys are missing.
        """
        if not config or "path_pattern" not in config:
            raise ValueError(f"Configuration for ParquetOutputWriter '{name}' must contain 'path_pattern'")
        if "columns" not in config:
            raise ValueError(f"Configuration for ParquetOutputWriter '{name}' must contain 'columns'")

        self.path_pattern = config["path_pattern"]
        self.columns = config["columns"]
        self.name = name

    def write_metrics_dict(self, metrics_dict: dict[str, dict[str, Any]]) -> None:
        """Aggregate and write dataset-level metrics for all selections to a Parquet file.

        Args:
            metrics_dict: Map of selection names to their computed metric dictionaries.
        """
        if len(metrics_dict) > 0:
            logger.debug(f"Writing metrics for the {len(metrics_dict)} data selections")

            # We get all the selections computed
            keys = list(metrics_dict.keys())
            metric_names = list(metrics_dict[keys[0]].keys())
            metrics_table = {"selection": pa.array(keys)}

            for metric_name in metric_names:
                metrics_table[metric_name] = pa.array([metrics_dict[key][metric_name] for key in keys])
        self.write_table("", metrics_table)

    def write_table(self, path_pattern: str, features_array: dict[str, Any], part: int | None = None) -> None:
        """
        Write a table of features or metrics to a Parquet file on disk.

        Handles directory creation if the target path doesn't exist.

        Args:
            path_pattern: Identifier for the data detination (used in filename pattern).
            features_array: Map of column names to pyarrow Arrays.
            part: Optional partition index for chunked output.
        """

        for key in self.columns:
            if key not in features_array:
                logger.error(f"Missing {key} in features for output")

        table = pa.table(features_array)
        if part is None:
            filename = self.path_pattern.format(path_pattern, "")
        else:
            filename = self.path_pattern.format(path_pattern, part)

        output_dir = Path(filename).parent
        if not Path.exists(output_dir):
            logger.info(f"Creating output directory: {output_dir}")
            Path.mkdir(output_dir, parents=True, exist_ok=True)

        pq.write_table(table, filename)
        logger.info(f"Wrote output table to {filename}")

columns = config['columns'] instance-attribute

name = name instance-attribute

path_pattern = config['path_pattern'] instance-attribute

__init__(name: str, config: dict[str, Any] | None = None)

Initialize a ParquetOutputWriter.

Parameters:

Name	Type	Description	Default
name	str	Unique name for this output writer.	required
config	dict[str, Any] | None	Configuration dictionary with keys: - path_pattern (str): Output file path format string. - columns (List[str]): Columns to save.	None

Raises:

Type	Description
ValueError	If required config keys are missing.

Source code in packages/dqm-ml-job/src/dqm_ml_job/outputwriter/parquet.py

def __init__(self, name: str, config: dict[str, Any] | None = None):
    """
    Initialize a ParquetOutputWriter.

    Args:
        name: Unique name for this output writer.
        config: Configuration dictionary with keys:
            - path_pattern (str): Output file path format string.
            - columns (List[str]): Columns to save.

    Raises:
        ValueError: If required config keys are missing.
    """
    if not config or "path_pattern" not in config:
        raise ValueError(f"Configuration for ParquetOutputWriter '{name}' must contain 'path_pattern'")
    if "columns" not in config:
        raise ValueError(f"Configuration for ParquetOutputWriter '{name}' must contain 'columns'")

    self.path_pattern = config["path_pattern"]
    self.columns = config["columns"]
    self.name = name

write_metrics_dict(metrics_dict: dict[str, dict[str, Any]]) -> None

Aggregate and write dataset-level metrics for all selections to a Parquet file.

Parameters:

Name	Type	Description	Default
metrics_dict	dict[str, dict[str, Any]]	Map of selection names to their computed metric dictionaries.	required

Source code in packages/dqm-ml-job/src/dqm_ml_job/outputwriter/parquet.py

def write_metrics_dict(self, metrics_dict: dict[str, dict[str, Any]]) -> None:
    """Aggregate and write dataset-level metrics for all selections to a Parquet file.

    Args:
        metrics_dict: Map of selection names to their computed metric dictionaries.
    """
    if len(metrics_dict) > 0:
        logger.debug(f"Writing metrics for the {len(metrics_dict)} data selections")

        # We get all the selections computed
        keys = list(metrics_dict.keys())
        metric_names = list(metrics_dict[keys[0]].keys())
        metrics_table = {"selection": pa.array(keys)}

        for metric_name in metric_names:
            metrics_table[metric_name] = pa.array([metrics_dict[key][metric_name] for key in keys])
    self.write_table("", metrics_table)

write_table(path_pattern: str, features_array: dict[str, Any], part: int | None = None) -> None

Write a table of features or metrics to a Parquet file on disk.

Handles directory creation if the target path doesn't exist.

Parameters:

Name	Type	Description	Default
path_pattern	str	Identifier for the data detination (used in filename pattern).	required
features_array	dict[str, Any]	Map of column names to pyarrow Arrays.	required
part	int | None	Optional partition index for chunked output.	None

Source code in packages/dqm-ml-job/src/dqm_ml_job/outputwriter/parquet.py

def write_table(self, path_pattern: str, features_array: dict[str, Any], part: int | None = None) -> None:
    """
    Write a table of features or metrics to a Parquet file on disk.

    Handles directory creation if the target path doesn't exist.

    Args:
        path_pattern: Identifier for the data detination (used in filename pattern).
        features_array: Map of column names to pyarrow Arrays.
        part: Optional partition index for chunked output.
    """

    for key in self.columns:
        if key not in features_array:
            logger.error(f"Missing {key} in features for output")

    table = pa.table(features_array)
    if part is None:
        filename = self.path_pattern.format(path_pattern, "")
    else:
        filename = self.path_pattern.format(path_pattern, part)

    output_dir = Path(filename).parent
    if not Path.exists(output_dir):
        logger.info(f"Creating output directory: {output_dir}")
        Path.mkdir(output_dir, parents=True, exist_ok=True)

    pq.write_table(table, filename)
    logger.info(f"Wrote output table to {filename}")