|
| 1 | +# |
| 2 | +# Licensed to the Apache Software Foundation (ASF) under one or more |
| 3 | +# contributor license agreements. See the NOTICE file distributed with |
| 4 | +# this work for additional information regarding copyright ownership. |
| 5 | +# The ASF licenses this file to You under the Apache License, Version 2.0 |
| 6 | +# (the "License"); you may not use this file except in compliance with |
| 7 | +# the License. You may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, software |
| 12 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | +# See the License for the specific language governing permissions and |
| 15 | +# limitations under the License. |
| 16 | +# |
| 17 | + |
| 18 | +"""Categorical encoding pipeline using MLTransform for batch processing. |
| 19 | +
|
| 20 | +This pipeline demonstrates MLTransform's ComputeAndApplyVocabulary transform |
| 21 | +for categorical feature encoding. It can either read input data from a file |
| 22 | +or generate synthetic test data, computes vocabulary on categorical columns, |
| 23 | +and converts categorical values to integer indices. |
| 24 | +
|
| 25 | +Example usage with input file: |
| 26 | + python mltransform_one_hot_encoding.py \ |
| 27 | + --input_file=gs://bucket/input.jsonl \ |
| 28 | + --output_file=gs://bucket/output.jsonl \ |
| 29 | + --artifact_location=gs://bucket/artifacts \ |
| 30 | + --categorical_columns=category \ |
| 31 | + --runner=DataflowRunner \ |
| 32 | + --project=PROJECT \ |
| 33 | + --region=us-central1 \ |
| 34 | + --temp_location=gs://bucket/temp |
| 35 | +
|
| 36 | +Example usage with synthetic data: |
| 37 | + python mltransform_one_hot_encoding.py \ |
| 38 | + --output_file=gs://bucket/output.jsonl \ |
| 39 | + --categorical_columns=category \ |
| 40 | + --num_records=100000 \ |
| 41 | + --runner=DataflowRunner \ |
| 42 | + --project=PROJECT \ |
| 43 | + --region=us-central1 |
| 44 | +""" |
| 45 | + |
| 46 | +import argparse |
| 47 | +import json |
| 48 | +import logging |
| 49 | +import tempfile |
| 50 | +from typing import Any |
| 51 | + |
| 52 | +import apache_beam as beam |
| 53 | +from apache_beam.ml.transforms.base import MLTransform |
| 54 | +from apache_beam.ml.transforms.tft import ComputeAndApplyVocabulary |
| 55 | +from apache_beam.runners.runner import PipelineResult |
| 56 | + |
| 57 | + |
| 58 | +def parse_json_line(line: str) -> dict[str, Any]: |
| 59 | + """Parse a JSON line into a dictionary.""" |
| 60 | + try: |
| 61 | + return json.loads(line) |
| 62 | + except json.JSONDecodeError as e: |
| 63 | + raise ValueError(f"Failed to parse JSON line: {line[:200]}...") from e |
| 64 | + |
| 65 | + |
| 66 | +def parse_text_line(line: str, |
| 67 | + categorical_columns: list[str]) -> dict[str, Any]: |
| 68 | + """Parse plain text line into the first categorical column.""" |
| 69 | + text_value = line.strip() |
| 70 | + if not text_value: |
| 71 | + text_value = 'unknown' |
| 72 | + return {categorical_columns[0]: text_value} |
| 73 | + |
| 74 | + |
| 75 | +def format_json_output(element: Any) -> str: |
| 76 | + """Format output element as JSON string.""" |
| 77 | + def to_json_compatible(value: Any) -> Any: |
| 78 | + """Recursively convert non-JSON types (e.g. numpy arrays/scalars).""" |
| 79 | + if isinstance(value, dict): |
| 80 | + return {k: to_json_compatible(v) for k, v in value.items()} |
| 81 | + if isinstance(value, (list, tuple)): |
| 82 | + return [to_json_compatible(v) for v in value] |
| 83 | + |
| 84 | + # MLTransform outputs may include numpy scalar/ndarray values. |
| 85 | + if hasattr(value, 'tolist'): |
| 86 | + return to_json_compatible(value.tolist()) |
| 87 | + if hasattr(value, 'item'): |
| 88 | + try: |
| 89 | + return to_json_compatible(value.item()) |
| 90 | + except (TypeError, ValueError): |
| 91 | + pass |
| 92 | + return value |
| 93 | + |
| 94 | + if hasattr(element, 'as_dict'): |
| 95 | + return json.dumps(to_json_compatible(element.as_dict())) |
| 96 | + if hasattr(element, '_asdict'): |
| 97 | + return json.dumps(to_json_compatible(element._asdict())) |
| 98 | + return json.dumps(to_json_compatible(dict(element))) |
| 99 | + |
| 100 | + |
| 101 | +def generate_synthetic_record(index: int, |
| 102 | + categorical_columns: list[str]) -> dict[str, str]: |
| 103 | + """Generate a deterministic synthetic record with categorical values.""" |
| 104 | + categories = [ |
| 105 | + 'electronics', |
| 106 | + 'clothing', |
| 107 | + 'food', |
| 108 | + 'books', |
| 109 | + 'sports', |
| 110 | + 'home', |
| 111 | + 'toys', |
| 112 | + 'health', |
| 113 | + 'automotive', |
| 114 | + 'garden' |
| 115 | + ] |
| 116 | + colors = [ |
| 117 | + 'red', |
| 118 | + 'blue', |
| 119 | + 'green', |
| 120 | + 'yellow', |
| 121 | + 'black', |
| 122 | + 'white', |
| 123 | + 'purple', |
| 124 | + 'orange', |
| 125 | + 'pink', |
| 126 | + 'gray' |
| 127 | + ] |
| 128 | + sizes = ['small', 'medium', 'large', 'xlarge', 'tiny', 'huge'] |
| 129 | + |
| 130 | + record = {} |
| 131 | + for col in categorical_columns: |
| 132 | + if col.lower() in ['category', 'type', 'product']: |
| 133 | + record[col] = categories[index % len(categories)] |
| 134 | + elif col.lower() in ['color', 'colour']: |
| 135 | + record[col] = colors[index % len(colors)] |
| 136 | + elif col.lower() in ['size', 'dimension']: |
| 137 | + record[col] = sizes[index % len(sizes)] |
| 138 | + else: |
| 139 | + # Default to categories for unknown columns |
| 140 | + record[col] = categories[index % len(categories)] |
| 141 | + return record |
| 142 | + |
| 143 | + |
| 144 | +def run( |
| 145 | + argv=None, |
| 146 | + save_main_session=True, |
| 147 | + test_pipeline=None) -> PipelineResult | None: |
| 148 | + """Run the categorical encoding pipeline.""" |
| 149 | + known_args, pipeline_args = parse_known_args(argv) |
| 150 | + |
| 151 | + categorical_columns = [ |
| 152 | + col.strip() for col in known_args.categorical_columns.split(',') |
| 153 | + ] |
| 154 | + |
| 155 | + if not categorical_columns or not categorical_columns[0]: |
| 156 | + raise ValueError("At least one categorical column must be specified") |
| 157 | + |
| 158 | + if not known_args.output_file: |
| 159 | + raise ValueError("--output_file is required") |
| 160 | + |
| 161 | + # Create artifact location if not provided |
| 162 | + artifact_location = known_args.artifact_location |
| 163 | + if not artifact_location: |
| 164 | + artifact_location = tempfile.mkdtemp() |
| 165 | + logging.info("Using temporary artifact location: %s", artifact_location) |
| 166 | + |
| 167 | + pipeline_options = beam.options.pipeline_options.PipelineOptions( |
| 168 | + pipeline_args) |
| 169 | + pipeline_options.view_as( |
| 170 | + beam.options.pipeline_options.SetupOptions |
| 171 | + ).save_main_session = save_main_session |
| 172 | + |
| 173 | + pipeline = test_pipeline or beam.Pipeline(options=pipeline_options) |
| 174 | + |
| 175 | + # Use synthetic data or read from file |
| 176 | + if known_args.input_file: |
| 177 | + # Read and parse input data from file |
| 178 | + if known_args.input_format == 'jsonl': |
| 179 | + parse_input_fn = parse_json_line |
| 180 | + else: |
| 181 | + if len(categorical_columns) > 1: |
| 182 | + logging.warning( |
| 183 | + 'Input format is "text" but multiple categorical columns are ' |
| 184 | + 'specified. Only the first column "%s" will be used for parsing.', |
| 185 | + categorical_columns[0]) |
| 186 | + parse_input_fn = lambda line: parse_text_line(line, categorical_columns) |
| 187 | + raw_data = ( |
| 188 | + pipeline |
| 189 | + | 'ReadFromJSONL' >> beam.io.ReadFromText(known_args.input_file) |
| 190 | + | 'ParseInput' >> beam.Map(parse_input_fn)) |
| 191 | + else: |
| 192 | + # Generate synthetic data |
| 193 | + num_records = known_args.num_records or 100000 |
| 194 | + logging.info("Generating %d synthetic records", num_records) |
| 195 | + |
| 196 | + raw_data = ( |
| 197 | + pipeline |
| 198 | + | 'GenerateSyntheticIndexes' >> beam.Create(range(num_records)) |
| 199 | + | 'BuildSyntheticRecord' >> beam.Map( |
| 200 | + lambda idx: generate_synthetic_record(idx, categorical_columns))) |
| 201 | + |
| 202 | + # Build MLTransform with ComputeAndApplyVocabulary |
| 203 | + ml_transform = MLTransform( |
| 204 | + write_artifact_location=artifact_location, |
| 205 | + ).with_transform( |
| 206 | + ComputeAndApplyVocabulary( |
| 207 | + columns=categorical_columns, vocab_filename='vocab_onehot')) |
| 208 | + |
| 209 | + # Apply MLTransform |
| 210 | + transformed_data = ( |
| 211 | + raw_data |
| 212 | + | 'ValidateAndFilterColumns' >> beam.Filter( |
| 213 | + lambda element: all(col in element for col in categorical_columns)) |
| 214 | + | 'MLTransform' >> ml_transform |
| 215 | + | 'FormatOutput' >> beam.Map(format_json_output)) |
| 216 | + |
| 217 | + # Write output |
| 218 | + _ = ( |
| 219 | + transformed_data |
| 220 | + | 'WriteToJSONL' >> beam.io.WriteToText( |
| 221 | + known_args.output_file, file_name_suffix='.jsonl')) |
| 222 | + |
| 223 | + result = pipeline.run() |
| 224 | + return result |
| 225 | + |
| 226 | + |
| 227 | +def parse_known_args(argv): |
| 228 | + """Parse command-line arguments.""" |
| 229 | + parser = argparse.ArgumentParser( |
| 230 | + description='Categorical encoding pipeline using MLTransform') |
| 231 | + |
| 232 | + parser.add_argument( |
| 233 | + '--input_file', |
| 234 | + help='Input JSONL file path (local or GCS). ' |
| 235 | + 'If not provided, synthetic data will be generated.') |
| 236 | + parser.add_argument( |
| 237 | + '--input_format', |
| 238 | + choices=['jsonl', 'text'], |
| 239 | + default='jsonl', |
| 240 | + help='Input file format for --input_file. Use jsonl for JSON lines ' |
| 241 | + 'or text for plain text lines (default: jsonl).') |
| 242 | + parser.add_argument( |
| 243 | + '--output_file', |
| 244 | + required=True, |
| 245 | + help='Output file prefix for encoded results (JSONL format)') |
| 246 | + parser.add_argument( |
| 247 | + '--artifact_location', |
| 248 | + help='GCS or local path to store MLTransform artifacts ' |
| 249 | + '(vocabulary files). If not provided, a temp location is used.') |
| 250 | + parser.add_argument( |
| 251 | + '--categorical_columns', |
| 252 | + required=True, |
| 253 | + help='Comma-separated list of categorical column names to encode') |
| 254 | + parser.add_argument( |
| 255 | + '--num_records', |
| 256 | + type=int, |
| 257 | + default=100000, |
| 258 | + help='Number of synthetic records to generate if --input_file is not ' |
| 259 | + 'provided (default: 100000)') |
| 260 | + |
| 261 | + return parser.parse_known_args(argv) |
| 262 | + |
| 263 | + |
| 264 | +if __name__ == '__main__': |
| 265 | + logging.getLogger().setLevel(logging.INFO) |
| 266 | + run() |
0 commit comments