Skip to content

dqm_ml_job.outputwriter.parquet

Parquet output writer for persisting pipeline results.

This module contains the ParquetOutputWriter class that writes metrics and features to Parquet files.

logger = logging.getLogger(__name__) module-attribute

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 
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
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.
                - storage (bool or dict, optional): Storage configuration.
                  If dict with type "s3", can contain access_key, secret_key, and endpoint_override.

        Raises:
            ValueError: If required config keys are missing.
        """
        cfg = ParquetOutputConfig.model_validate(config or {})

        self.path_pattern = cfg.path_pattern
        self.columns = list(cfg.columns)
        self.exclude = list(cfg.exclude)
        self.name = name
        self.s3_filesystem = None

        self._accumulate = "{}" not in self.path_pattern
        self._accumulated_features: dict[str, list[pa.Array]] = {}
        storage_cfg = cfg.storage
        if storage_cfg:
            storage_config = StorageConfig.model_validate(storage_cfg)
            if storage_config.type == "s3":
                self.s3_filesystem = get_s3_filesystem(storage_config)

    @staticmethod
    def _collect_metric(metric_name: str, metrics_dict: dict[str, dict[str, Any]], keys: list[str]) -> pa.Array | None:
        values = []
        for key in keys:
            val = metrics_dict[key][metric_name]
            if isinstance(val, pa.FixedSizeListArray):
                return None
            if isinstance(val, pa.Array):
                values.extend(val.to_pylist())
            else:
                values.append(val)
        return pa.array(values)

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

        Args:
            metrics_dict: Map of selection names to their computed
                metric dictionaries.
        """
        if len(metrics_dict) <= 0:
            return
        logger.debug(f"Writing metrics for the {len(metrics_dict)} data selections")
        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:
            if metric_name.startswith("__") and metric_name.endswith("__"):
                continue
            col = self._collect_metric(metric_name, metrics_dict, keys)
            if col is not None:
                metrics_table[metric_name] = col
        self.write_table("", metrics_table)

    def _format_filename(self, path_pattern: str, part: int | None = None) -> str:
        """Format the output filename from the path pattern and optional part."""
        if part is None:
            return self.path_pattern.format(path_pattern, "")
        return self.path_pattern.format(path_pattern, part)

    def _write_local(self, table: pa.Table, filename: str) -> None:
        """Write a table to a local file, creating the parent directory if needed."""
        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}")

    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.

        Handles directory creation if the target path doesn't exist
        (for local writes).

        Args:
            path_pattern: Identifier for the data destination.
            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")

        # Accumulate mode: buffer features for a single flush at the end
        if self._accumulate:
            for k, v in features_array.items():
                if isinstance(v, pa.ChunkedArray):
                    v = v.combine_chunks()
                self._accumulated_features.setdefault(k, []).append(v)
            return

        table = pa.table(features_array)
        filename = self._format_filename(path_pattern, part)
        if self.s3_filesystem is not None:
            self._write_to_s3(table, filename)
        else:
            self._write_local(table, filename)

    def flush(self) -> None:
        """Write all accumulated features to the output file.

        Called after all selections have been processed in accumulate mode.
        Does nothing if there are no accumulated features or if not
        in accumulate mode.
        """
        if not self._accumulated_features:
            return

        final = {k: pa.concat_arrays(v) for k, v in self._accumulated_features.items()}
        table = pa.table(final)
        filename = self.path_pattern

        if self.s3_filesystem is not None:
            s3_path = self._get_s3_path(filename)
            pq.write_table(table, s3_path, filesystem=self.s3_filesystem)
            logger.info(f"Wrote accumulated output table to S3: {s3_path}")
        else:
            self._write_local(table, filename)

        self._accumulated_features.clear()

    def _get_s3_path(self, file_path: str) -> str:
        """Construct an S3 path by combining the bucket name with a file path.

        Args:
            file_path: The file path within the bucket.

        Returns:
            str: The full S3 path in format "bucket_name/file_path".
        """
        bucket_name = os.getenv("S3_BUCKET_NAME", "")
        return bucket_name + "/" + file_path

    def _write_to_s3(self, table: pa.Table, filename: str) -> None:
        """Write a PyArrow table to S3.

        Args:
            table: The table to write.
            filename: The file path within the bucket.
        """
        s3_path = self._get_s3_path(filename)
        try:
            pq.write_table(table, s3_path, filesystem=self.s3_filesystem)
            logger.info(f"Wrote output table to S3: {s3_path}")
        except Exception:
            logger.exception("Failed to write to S3")
            raise

columns = list(cfg.columns) instance-attribute

exclude = list(cfg.exclude) instance-attribute

name = name instance-attribute

path_pattern = cfg.path_pattern instance-attribute

s3_filesystem = None 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. - storage (bool or dict, optional): Storage configuration. If dict with type "s3", can contain access_key, secret_key, and endpoint_override.

 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 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
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.
            - storage (bool or dict, optional): Storage configuration.
              If dict with type "s3", can contain access_key, secret_key, and endpoint_override.

    Raises:
        ValueError: If required config keys are missing.
    """
    cfg = ParquetOutputConfig.model_validate(config or {})

    self.path_pattern = cfg.path_pattern
    self.columns = list(cfg.columns)
    self.exclude = list(cfg.exclude)
    self.name = name
    self.s3_filesystem = None

    self._accumulate = "{}" not in self.path_pattern
    self._accumulated_features: dict[str, list[pa.Array]] = {}
    storage_cfg = cfg.storage
    if storage_cfg:
        storage_config = StorageConfig.model_validate(storage_cfg)
        if storage_config.type == "s3":
            self.s3_filesystem = get_s3_filesystem(storage_config)

flush() -> None

Write all accumulated features to the output file.

Called after all selections have been processed in accumulate mode. Does nothing if there are no accumulated features or if not in accumulate mode.

Source code in packages/dqm-ml-job/src/dqm_ml_job/outputwriter/parquet.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
def flush(self) -> None:
    """Write all accumulated features to the output file.

    Called after all selections have been processed in accumulate mode.
    Does nothing if there are no accumulated features or if not
    in accumulate mode.
    """
    if not self._accumulated_features:
        return

    final = {k: pa.concat_arrays(v) for k, v in self._accumulated_features.items()}
    table = pa.table(final)
    filename = self.path_pattern

    if self.s3_filesystem is not None:
        s3_path = self._get_s3_path(filename)
        pq.write_table(table, s3_path, filesystem=self.s3_filesystem)
        logger.info(f"Wrote accumulated output table to S3: {s3_path}")
    else:
        self._write_local(table, filename)

    self._accumulated_features.clear()

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

Aggregate and write dataset-level metrics for all selections.

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 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def write_metrics_dict(self, metrics_dict: dict[str, dict[str, Any]]) -> None:
    """Aggregate and write dataset-level metrics for all selections.

    Args:
        metrics_dict: Map of selection names to their computed
            metric dictionaries.
    """
    if len(metrics_dict) <= 0:
        return
    logger.debug(f"Writing metrics for the {len(metrics_dict)} data selections")
    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:
        if metric_name.startswith("__") and metric_name.endswith("__"):
            continue
        col = self._collect_metric(metric_name, metrics_dict, keys)
        if col is not None:
            metrics_table[metric_name] = col
    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.

Handles directory creation if the target path doesn't exist (for local writes).

Parameters:

Name Type Description Default
path_pattern str

Identifier for the data destination.

 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 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
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.

    Handles directory creation if the target path doesn't exist
    (for local writes).

    Args:
        path_pattern: Identifier for the data destination.
        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")

    # Accumulate mode: buffer features for a single flush at the end
    if self._accumulate:
        for k, v in features_array.items():
            if isinstance(v, pa.ChunkedArray):
                v = v.combine_chunks()
            self._accumulated_features.setdefault(k, []).append(v)
        return

    table = pa.table(features_array)
    filename = self._format_filename(path_pattern, part)
    if self.s3_filesystem is not None:
        self._write_to_s3(table, filename)
    else:
        self._write_local(table, filename)