Coverage for packages/dqm-ml-job/src/dqm_ml_job/outputwriter/parquet.py: 83%

1"""Parquet output writer for persisting pipeline results. 

2 

3This module contains the ParquetOutputWriter class that writes 

4metrics and features to Parquet files. 

5""" 

6 

7import logging 

8import os 

9from pathlib import Path 

10from typing import Any 

11 

12from dqm_ml_core.models.global_ import StorageConfig 

13from dqm_ml_core.models.outputs import ParquetOutputConfig 

14import pyarrow as pa 

15import pyarrow.parquet as pq 

16 

17from dqm_ml_job.utils.s3 import get_s3_filesystem 

18 

19logger = logging.getLogger(__name__) 

20 

21 

22class ParquetOutputWriter: 

23 """ 

24 Output writer that saves processed features to a Parquet file. 

25 """ 

26 

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

28 """ 

29 Initialize a ParquetOutputWriter. 

30 

31 Args: 

32 name: Unique name for this output writer. 

33 config: Configuration dictionary with keys: 

34 - path_pattern (str): Output file path format string. 

35 - columns (List[str]): Columns to save. 

36 - storage (bool or dict, optional): Storage configuration. 

37 If dict with type "s3", can contain access_key, secret_key, and endpoint_override. 

38 

39 Raises: 

40 ValueError: If required config keys are missing. 

41 """ 

42 cfg = ParquetOutputConfig.model_validate(config or {}) 

43 

44 self.path_pattern = cfg.path_pattern 

45 self.columns = list(cfg.columns) 

46 self.exclude = list(cfg.exclude) 

47 self.name = name 

48 self.s3_filesystem = None 

49 

50 self._accumulate = "{}" not in self.path_pattern 

51 self._accumulated_features: dict[str, list[pa.Array]] = {} 

52 storage_cfg = cfg.storage 

53 if storage_cfg: 53 ↛ 54line 53 didn't jump to line 54 because the condition on line 53 was never true

54 storage_config = StorageConfig.model_validate(storage_cfg) 

55 if storage_config.type == "s3": 

56 self.s3_filesystem = get_s3_filesystem(storage_config) 

57 

58 @staticmethod 

59 def _collect_metric(metric_name: str, metrics_dict: dict[str, dict[str, Any]], keys: list[str]) -> pa.Array | None: 

60 values = [] 

61 for key in keys: 

62 val = metrics_dict[key][metric_name] 

63 if isinstance(val, pa.FixedSizeListArray): 

64 return None 

65 if isinstance(val, pa.Array): 

66 values.extend(val.to_pylist()) 

67 else: 

68 values.append(val) 

69 return pa.array(values) 

70 

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

72 """Aggregate and write dataset-level metrics for all selections. 

73 

74 Args: 

75 metrics_dict: Map of selection names to their computed 

76 metric dictionaries. 

77 """ 

78 if len(metrics_dict) <= 0: 78 ↛ 79line 78 didn't jump to line 79 because the condition on line 78 was never true

79 return 

80 logger.debug(f"Writing metrics for the {len(metrics_dict)} data selections") 

81 keys = list(metrics_dict.keys()) 

82 metric_names = list(metrics_dict[keys[0]].keys()) 

83 metrics_table = {"selection": pa.array(keys)} 

84 for metric_name in metric_names: 

85 if metric_name.startswith("__") and metric_name.endswith("__"): 

86 continue 

87 col = self._collect_metric(metric_name, metrics_dict, keys) 

88 if col is not None: 

89 metrics_table[metric_name] = col 

90 self.write_table("", metrics_table) 

91 

92 def _format_filename(self, path_pattern: str, part: int | None = None) -> str: 

93 """Format the output filename from the path pattern and optional part.""" 

94 if part is None: 

95 return self.path_pattern.format(path_pattern, "") 

96 return self.path_pattern.format(path_pattern, part) 

97 

98 def _write_local(self, table: pa.Table, filename: str) -> None: 

99 """Write a table to a local file, creating the parent directory if needed.""" 

100 output_dir = Path(filename).parent 

101 if not Path.exists(output_dir): 

102 logger.info(f"Creating output directory: {output_dir}") 

103 Path.mkdir(output_dir, parents=True, exist_ok=True) 

104 pq.write_table(table, filename) 

105 logger.info(f"Wrote output table to {filename}") 

106 

107 def write_table( 

108 self, 

109 path_pattern: str, 

110 features_array: dict[str, Any], 

111 part: int | None = None, 

112 ) -> None: 

113 """Write a table of features or metrics to a Parquet file. 

114 

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

116 (for local writes). 

117 

118 Args: 

119 path_pattern: Identifier for the data destination. 

120 features_array: Map of column names to pyarrow Arrays. 

121 part: Optional partition index for chunked output. 

122 """ 

123 

124 for key in self.columns: 

125 if key not in features_array: 

126 logger.error(f"Missing {key} in features for output") 

127 

128 # Accumulate mode: buffer features for a single flush at the end 

129 if self._accumulate: 

130 for k, v in features_array.items(): 

131 if isinstance(v, pa.ChunkedArray): 

132 v = v.combine_chunks() 

133 self._accumulated_features.setdefault(k, []).append(v) 

134 return 

135 

136 table = pa.table(features_array) 

137 filename = self._format_filename(path_pattern, part) 

138 if self.s3_filesystem is not None: 138 ↛ 139line 138 didn't jump to line 139 because the condition on line 138 was never true

139 self._write_to_s3(table, filename) 

140 else: 

141 self._write_local(table, filename) 

142 

143 def flush(self) -> None: 

144 """Write all accumulated features to the output file. 

145 

146 Called after all selections have been processed in accumulate mode. 

147 Does nothing if there are no accumulated features or if not 

148 in accumulate mode. 

149 """ 

150 if not self._accumulated_features: 

151 return 

152 

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

154 table = pa.table(final) 

155 filename = self.path_pattern 

156 

157 if self.s3_filesystem is not None: 157 ↛ 158line 157 didn't jump to line 158 because the condition on line 157 was never true

158 s3_path = self._get_s3_path(filename) 

159 pq.write_table(table, s3_path, filesystem=self.s3_filesystem) 

160 logger.info(f"Wrote accumulated output table to S3: {s3_path}") 

161 else: 

162 self._write_local(table, filename) 

163 

164 self._accumulated_features.clear() 

165 

166 def _get_s3_path(self, file_path: str) -> str: 

167 """Construct an S3 path by combining the bucket name with a file path. 

168 

169 Args: 

170 file_path: The file path within the bucket. 

171 

172 Returns: 

173 str: The full S3 path in format "bucket_name/file_path". 

174 """ 

175 bucket_name = os.getenv("S3_BUCKET_NAME", "") 

176 return bucket_name + "/" + file_path 

177 

178 def _write_to_s3(self, table: pa.Table, filename: str) -> None: 

179 """Write a PyArrow table to S3. 

180 

181 Args: 

182 table: The table to write. 

183 filename: The file path within the bucket. 

184 """ 

185 s3_path = self._get_s3_path(filename) 

186 try: 

187 pq.write_table(table, s3_path, filesystem=self.s3_filesystem) 

188 logger.info(f"Wrote output table to S3: {s3_path}") 

189 except Exception: 

190 logger.exception("Failed to write to S3") 

191 raise