gemini-code-assist[bot] commented on code in PR #38404: URL: https://github.com/apache/beam/pull/38404#discussion_r3202149858
########## sdks/python/apache_beam/examples/ml_transform/mltransform_one_hot_encoding.py: ########## @@ -0,0 +1,242 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Categorical encoding pipeline using MLTransform for batch processing. + +This pipeline demonstrates MLTransform's ComputeAndApplyVocabulary transform +for categorical feature encoding. It can either read input data from a file +or generate synthetic test data, computes vocabulary on categorical columns, +and converts categorical values to integer indices. + +Example usage with input file: + python mltransform_one_hot_encoding.py \ + --input_file=gs://bucket/input.jsonl \ + --output_file=gs://bucket/output.jsonl \ + --artifact_location=gs://bucket/artifacts \ + --categorical_columns=category \ + --runner=DataflowRunner \ + --project=PROJECT \ + --region=us-central1 \ + --temp_location=gs://bucket/temp + +Example usage with synthetic data: + python mltransform_one_hot_encoding.py \ + --output_file=gs://bucket/output.jsonl \ + --categorical_columns=category \ + --num_records=100000 \ + --runner=DataflowRunner \ + --project=PROJECT \ + --region=us-central1 +""" + +import argparse +import json +import logging +import tempfile +from typing import Any + +import apache_beam as beam +from apache_beam.ml.transforms.base import MLTransform +from apache_beam.ml.transforms.tft import ComputeAndApplyVocabulary +from apache_beam.runners.runner import PipelineResult + + +def parse_json_line(line: str) -> dict[str, Any]: + """Parse a JSON line into a dictionary.""" + try: + return json.loads(line) + except json.JSONDecodeError as e: + raise ValueError(f"Failed to parse JSON line: {line[:200]}...") from e + + +def parse_text_line(line: str, categorical_columns: list[str]) -> dict[str, Any]: + """Parse plain text line into the first categorical column.""" + text_value = line.strip() + if not text_value: + text_value = 'unknown' + return {categorical_columns[0]: text_value} + + +def format_json_output(element: Any) -> str: + """Format output element as JSON string.""" + def to_json_compatible(value: Any) -> Any: + """Recursively convert non-JSON types (e.g. numpy arrays/scalars).""" + if isinstance(value, dict): + return {k: to_json_compatible(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [to_json_compatible(v) for v in value] + + # MLTransform outputs may include numpy scalar/ndarray values. + if hasattr(value, 'tolist'): + return to_json_compatible(value.tolist()) + if hasattr(value, 'item'): + try: + return to_json_compatible(value.item()) + except (TypeError, ValueError): + pass + return value + + if hasattr(element, 'as_dict'): + return json.dumps(to_json_compatible(element.as_dict())) + if hasattr(element, '_asdict'): + return json.dumps(to_json_compatible(element._asdict())) + return json.dumps(to_json_compatible(dict(element))) + + +def generate_synthetic_record(index: int, categorical_columns: list[str]) -> dict[str, + str]: + """Generate a deterministic synthetic record with categorical values.""" + categories = ['electronics', 'clothing', 'food', 'books', 'sports', + 'home', 'toys', 'health', 'automotive', 'garden'] + colors = ['red', 'blue', 'green', 'yellow', 'black', 'white', + 'purple', 'orange', 'pink', 'gray'] + sizes = ['small', 'medium', 'large', 'xlarge', 'tiny', 'huge'] + + record = {} + for col in categorical_columns: + if col.lower() in ['category', 'type', 'product']: + record[col] = categories[index % len(categories)] + elif col.lower() in ['color', 'colour']: + record[col] = colors[index % len(colors)] + elif col.lower() in ['size', 'dimension']: + record[col] = sizes[index % len(sizes)] + else: + # Default to categories for unknown columns + record[col] = categories[index % len(categories)] + return record + + +def run( + argv=None, + save_main_session=True, + test_pipeline=None) -> PipelineResult | None: + """Run the categorical encoding pipeline.""" + known_args, pipeline_args = parse_known_args(argv) + + categorical_columns = [ + col.strip() for col in known_args.categorical_columns.split(',') + ] + + if not categorical_columns or not categorical_columns[0]: + raise ValueError("At least one categorical column must be specified") + + if not known_args.output_file: + raise ValueError("--output_file is required") + + # Create artifact location if not provided + artifact_location = known_args.artifact_location + if not artifact_location: + artifact_location = tempfile.mkdtemp() + logging.info("Using temporary artifact location: %s", artifact_location) + + pipeline_options = beam.options.pipeline_options.PipelineOptions(pipeline_args) + pipeline_options.view_as( + beam.options.pipeline_options.SetupOptions).save_main_session = save_main_session + + pipeline = test_pipeline or beam.Pipeline(options=pipeline_options) + + # Use synthetic data or read from file + if known_args.input_file: + # Read and parse input data from file + if known_args.input_format == 'jsonl': + parse_input_fn = parse_json_line + else: + parse_input_fn = lambda line: parse_text_line(line, categorical_columns) Review Comment:  When `input_format` is `text`, the pipeline implicitly uses only the first categorical column, which might be unexpected for users who provide multiple columns. To improve clarity and prevent potential confusion, it's good practice to log a warning in this scenario. ```suggestion else: if len(categorical_columns) > 1: logging.warning( 'Input format is "text" but multiple categorical columns are ' 'specified. Only the first column "%s" will be used for ' 'parsing.', categorical_columns[0]) parse_input_fn = lambda line: parse_text_line(line, categorical_columns) ``` ########## sdks/python/apache_beam/examples/ml_transform/mltransform_one_hot_encoding.py: ########## @@ -0,0 +1,242 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Categorical encoding pipeline using MLTransform for batch processing. + +This pipeline demonstrates MLTransform's ComputeAndApplyVocabulary transform +for categorical feature encoding. It can either read input data from a file +or generate synthetic test data, computes vocabulary on categorical columns, +and converts categorical values to integer indices. + +Example usage with input file: + python mltransform_one_hot_encoding.py \ + --input_file=gs://bucket/input.jsonl \ + --output_file=gs://bucket/output.jsonl \ + --artifact_location=gs://bucket/artifacts \ + --categorical_columns=category \ + --runner=DataflowRunner \ + --project=PROJECT \ + --region=us-central1 \ + --temp_location=gs://bucket/temp + +Example usage with synthetic data: + python mltransform_one_hot_encoding.py \ + --output_file=gs://bucket/output.jsonl \ + --categorical_columns=category \ + --num_records=100000 \ + --runner=DataflowRunner \ + --project=PROJECT \ + --region=us-central1 +""" + +import argparse +import json +import logging +import tempfile +from typing import Any + +import apache_beam as beam +from apache_beam.ml.transforms.base import MLTransform +from apache_beam.ml.transforms.tft import ComputeAndApplyVocabulary +from apache_beam.runners.runner import PipelineResult + + +def parse_json_line(line: str) -> dict[str, Any]: + """Parse a JSON line into a dictionary.""" + try: + return json.loads(line) + except json.JSONDecodeError as e: + raise ValueError(f"Failed to parse JSON line: {line[:200]}...") from e + + +def parse_text_line(line: str, categorical_columns: list[str]) -> dict[str, Any]: + """Parse plain text line into the first categorical column.""" + text_value = line.strip() + if not text_value: + text_value = 'unknown' + return {categorical_columns[0]: text_value} + + +def format_json_output(element: Any) -> str: + """Format output element as JSON string.""" + def to_json_compatible(value: Any) -> Any: + """Recursively convert non-JSON types (e.g. numpy arrays/scalars).""" + if isinstance(value, dict): + return {k: to_json_compatible(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [to_json_compatible(v) for v in value] + + # MLTransform outputs may include numpy scalar/ndarray values. + if hasattr(value, 'tolist'): + return to_json_compatible(value.tolist()) + if hasattr(value, 'item'): + try: + return to_json_compatible(value.item()) + except (TypeError, ValueError): + pass + return value + + if hasattr(element, 'as_dict'): + return json.dumps(to_json_compatible(element.as_dict())) + if hasattr(element, '_asdict'): + return json.dumps(to_json_compatible(element._asdict())) + return json.dumps(to_json_compatible(dict(element))) + + +def generate_synthetic_record(index: int, categorical_columns: list[str]) -> dict[str, + str]: + """Generate a deterministic synthetic record with categorical values.""" + categories = ['electronics', 'clothing', 'food', 'books', 'sports', + 'home', 'toys', 'health', 'automotive', 'garden'] + colors = ['red', 'blue', 'green', 'yellow', 'black', 'white', + 'purple', 'orange', 'pink', 'gray'] + sizes = ['small', 'medium', 'large', 'xlarge', 'tiny', 'huge'] + + record = {} + for col in categorical_columns: + if col.lower() in ['category', 'type', 'product']: + record[col] = categories[index % len(categories)] + elif col.lower() in ['color', 'colour']: + record[col] = colors[index % len(colors)] + elif col.lower() in ['size', 'dimension']: + record[col] = sizes[index % len(sizes)] + else: + # Default to categories for unknown columns + record[col] = categories[index % len(categories)] + return record + + +def run( + argv=None, + save_main_session=True, + test_pipeline=None) -> PipelineResult | None: + """Run the categorical encoding pipeline.""" + known_args, pipeline_args = parse_known_args(argv) + + categorical_columns = [ + col.strip() for col in known_args.categorical_columns.split(',') + ] + + if not categorical_columns or not categorical_columns[0]: + raise ValueError("At least one categorical column must be specified") + + if not known_args.output_file: + raise ValueError("--output_file is required") + + # Create artifact location if not provided + artifact_location = known_args.artifact_location + if not artifact_location: + artifact_location = tempfile.mkdtemp() + logging.info("Using temporary artifact location: %s", artifact_location) + + pipeline_options = beam.options.pipeline_options.PipelineOptions(pipeline_args) + pipeline_options.view_as( + beam.options.pipeline_options.SetupOptions).save_main_session = save_main_session + + pipeline = test_pipeline or beam.Pipeline(options=pipeline_options) + + # Use synthetic data or read from file + if known_args.input_file: + # Read and parse input data from file + if known_args.input_format == 'jsonl': + parse_input_fn = parse_json_line + else: + parse_input_fn = lambda line: parse_text_line(line, categorical_columns) + raw_data = ( + pipeline + | 'ReadFromJSONL' >> beam.io.ReadFromText(known_args.input_file) + | 'ParseInput' >> beam.Map(parse_input_fn)) + else: + # Generate synthetic data + num_records = known_args.num_records or 100000 + logging.info("Generating %d synthetic records", num_records) + + raw_data = ( + pipeline + | 'GenerateSyntheticData' >> beam.Create([None]) + | 'ExpandSyntheticIndexes' >> beam.FlatMap(lambda _: range(num_records)) + | 'BuildSyntheticRecord' >> beam.Map( + lambda idx: json.dumps(generate_synthetic_record( + idx, categorical_columns))) + | 'ParseSyntheticJSON' >> beam.Map(parse_json_line)) + + # Build MLTransform with ComputeAndApplyVocabulary + ml_transform = MLTransform( + write_artifact_location=artifact_location, + ).with_transform( + ComputeAndApplyVocabulary( + columns=categorical_columns, + vocab_filename='vocab_onehot')) + + # Apply MLTransform + transformed_data = ( + raw_data + | 'MLTransform' >> ml_transform + | 'FormatOutput' >> beam.Map(format_json_output)) Review Comment:  The pipeline will fail if input records are missing any of the specified `categorical_columns`, as `MLTransform` requires all columns to be present. The tests suggest that invalid records should be filtered out. To make the pipeline more robust, a `beam.Filter` transform should be added to remove records with missing columns before they are processed by `MLTransform`. ```suggestion transformed_data = ( raw_data | 'ValidateAndFilterColumns' >> beam.Filter( lambda element: all(col in element for col in categorical_columns)) | 'MLTransform' >> ml_transform | 'FormatOutput' >> beam.Map(format_json_output)) ``` ########## sdks/python/apache_beam/examples/ml_transform/mltransform_one_hot_encoding.py: ########## @@ -0,0 +1,242 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Categorical encoding pipeline using MLTransform for batch processing. + +This pipeline demonstrates MLTransform's ComputeAndApplyVocabulary transform +for categorical feature encoding. It can either read input data from a file +or generate synthetic test data, computes vocabulary on categorical columns, +and converts categorical values to integer indices. + +Example usage with input file: + python mltransform_one_hot_encoding.py \ + --input_file=gs://bucket/input.jsonl \ + --output_file=gs://bucket/output.jsonl \ + --artifact_location=gs://bucket/artifacts \ + --categorical_columns=category \ + --runner=DataflowRunner \ + --project=PROJECT \ + --region=us-central1 \ + --temp_location=gs://bucket/temp + +Example usage with synthetic data: + python mltransform_one_hot_encoding.py \ + --output_file=gs://bucket/output.jsonl \ + --categorical_columns=category \ + --num_records=100000 \ + --runner=DataflowRunner \ + --project=PROJECT \ + --region=us-central1 +""" + +import argparse +import json +import logging +import tempfile +from typing import Any + +import apache_beam as beam +from apache_beam.ml.transforms.base import MLTransform +from apache_beam.ml.transforms.tft import ComputeAndApplyVocabulary +from apache_beam.runners.runner import PipelineResult + + +def parse_json_line(line: str) -> dict[str, Any]: + """Parse a JSON line into a dictionary.""" + try: + return json.loads(line) + except json.JSONDecodeError as e: + raise ValueError(f"Failed to parse JSON line: {line[:200]}...") from e + + +def parse_text_line(line: str, categorical_columns: list[str]) -> dict[str, Any]: + """Parse plain text line into the first categorical column.""" + text_value = line.strip() + if not text_value: + text_value = 'unknown' + return {categorical_columns[0]: text_value} + + +def format_json_output(element: Any) -> str: + """Format output element as JSON string.""" + def to_json_compatible(value: Any) -> Any: + """Recursively convert non-JSON types (e.g. numpy arrays/scalars).""" + if isinstance(value, dict): + return {k: to_json_compatible(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [to_json_compatible(v) for v in value] + + # MLTransform outputs may include numpy scalar/ndarray values. + if hasattr(value, 'tolist'): + return to_json_compatible(value.tolist()) + if hasattr(value, 'item'): + try: + return to_json_compatible(value.item()) + except (TypeError, ValueError): + pass + return value + + if hasattr(element, 'as_dict'): + return json.dumps(to_json_compatible(element.as_dict())) + if hasattr(element, '_asdict'): + return json.dumps(to_json_compatible(element._asdict())) + return json.dumps(to_json_compatible(dict(element))) + + +def generate_synthetic_record(index: int, categorical_columns: list[str]) -> dict[str, + str]: + """Generate a deterministic synthetic record with categorical values.""" + categories = ['electronics', 'clothing', 'food', 'books', 'sports', + 'home', 'toys', 'health', 'automotive', 'garden'] + colors = ['red', 'blue', 'green', 'yellow', 'black', 'white', + 'purple', 'orange', 'pink', 'gray'] + sizes = ['small', 'medium', 'large', 'xlarge', 'tiny', 'huge'] + + record = {} + for col in categorical_columns: + if col.lower() in ['category', 'type', 'product']: + record[col] = categories[index % len(categories)] + elif col.lower() in ['color', 'colour']: + record[col] = colors[index % len(colors)] + elif col.lower() in ['size', 'dimension']: + record[col] = sizes[index % len(sizes)] + else: + # Default to categories for unknown columns + record[col] = categories[index % len(categories)] + return record + + +def run( + argv=None, + save_main_session=True, + test_pipeline=None) -> PipelineResult | None: + """Run the categorical encoding pipeline.""" + known_args, pipeline_args = parse_known_args(argv) + + categorical_columns = [ + col.strip() for col in known_args.categorical_columns.split(',') + ] + + if not categorical_columns or not categorical_columns[0]: + raise ValueError("At least one categorical column must be specified") + + if not known_args.output_file: + raise ValueError("--output_file is required") + + # Create artifact location if not provided + artifact_location = known_args.artifact_location + if not artifact_location: + artifact_location = tempfile.mkdtemp() + logging.info("Using temporary artifact location: %s", artifact_location) + + pipeline_options = beam.options.pipeline_options.PipelineOptions(pipeline_args) + pipeline_options.view_as( + beam.options.pipeline_options.SetupOptions).save_main_session = save_main_session + + pipeline = test_pipeline or beam.Pipeline(options=pipeline_options) + + # Use synthetic data or read from file + if known_args.input_file: + # Read and parse input data from file + if known_args.input_format == 'jsonl': + parse_input_fn = parse_json_line + else: + parse_input_fn = lambda line: parse_text_line(line, categorical_columns) + raw_data = ( + pipeline + | 'ReadFromJSONL' >> beam.io.ReadFromText(known_args.input_file) + | 'ParseInput' >> beam.Map(parse_input_fn)) + else: + # Generate synthetic data + num_records = known_args.num_records or 100000 + logging.info("Generating %d synthetic records", num_records) + + raw_data = ( + pipeline + | 'GenerateSyntheticData' >> beam.Create([None]) + | 'ExpandSyntheticIndexes' >> beam.FlatMap(lambda _: range(num_records)) + | 'BuildSyntheticRecord' >> beam.Map( + lambda idx: json.dumps(generate_synthetic_record( + idx, categorical_columns))) + | 'ParseSyntheticJSON' >> beam.Map(parse_json_line)) Review Comment:  The current method for generating synthetic data is inefficient. It creates a single element, expands it to the desired number of records, converts each record to a JSON string, and then parses it back into a dictionary. This can be simplified and made more performant by directly generating the dictionaries. ```python raw_data = ( pipeline | 'GenerateIndexes' >> beam.Create(range(num_records)) | 'BuildSyntheticRecord' >> beam.Map( lambda idx: generate_synthetic_record(idx, categorical_columns))) ``` ########## sdks/python/apache_beam/examples/ml_transform/mltransform_one_hot_encoding_test.py: ########## @@ -0,0 +1,257 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Tests for mltransform_one_hot_encoding pipeline.""" + +import json +import logging +import os +import tempfile +import unittest +from typing import Any + +import pytest + +try: + from apache_beam.examples.ml_transform import mltransform_one_hot_encoding + from apache_beam.testing.test_pipeline import TestPipeline + from apache_beam.testing.util import BeamAssertException + from apache_beam.testing.util import assert_that + from apache_beam.testing.util import matches_all +except ImportError: # pylint: disable=bare-except + raise unittest.SkipTest('tensorflow_transform is not installed.') + + +def create_test_input_data() -> list[dict[str, Any]]: + """Create sample test data for one-hot encoding.""" + return [ + {'category': 'electronics', 'color': 'red', 'size': 'small'}, + {'category': 'clothing', 'color': 'blue', 'size': 'medium'}, + {'category': 'electronics', 'color': 'green', 'size': 'large'}, + {'category': 'food', 'color': 'red', 'size': 'small'}, + {'category': 'clothing', 'color': 'blue', 'size': 'medium'}, + ] + + +class OneHotEncodingPipelineTest(unittest.TestCase): + """Unit and integration tests for one-hot encoding pipeline.""" + + def setUp(self): + """Set up test fixtures.""" + self.test_dir = tempfile.mkdtemp() + self.input_file = os.path.join(self.test_dir, 'input.jsonl') + self.output_prefix = os.path.join(self.test_dir, 'output') + self.artifact_location = os.path.join(self.test_dir, 'artifacts') + + # Create test input file + test_data = create_test_input_data() + with open(self.input_file, 'w', encoding='utf-8') as f: + for record in test_data: + f.write(json.dumps(record) + '\n') + + def tearDown(self): + """Clean up test fixtures.""" + import shutil + shutil.rmtree(self.test_dir, ignore_errors=True) + + def test_parse_json_line_valid(self): + """Test parsing valid JSON lines.""" + line = '{"category": "electronics", "color": "red"}' + result = mltransform_one_hot_encoding.parse_json_line(line) + self.assertEqual(result['category'], 'electronics') + self.assertEqual(result['color'], 'red') + + def test_parse_json_line_invalid(self): + """Test parsing invalid JSON lines raises ValueError.""" + with self.assertRaises(ValueError): + mltransform_one_hot_encoding.parse_json_line('not valid json') + + def test_validate_columns_present_all_present(self): + """Test column validation when all columns are present.""" + element = {'category': 'a', 'color': 'b', 'size': 'c'} + result = mltransform_one_hot_encoding.validate_columns_present( + element, ['category', 'color']) + self.assertTrue(result) + + def test_validate_columns_present_missing(self): + """Test column validation when some columns are missing.""" + element = {'category': 'a'} + result = mltransform_one_hot_encoding.validate_columns_present( + element, ['category', 'color']) + self.assertFalse(result) Review Comment:  These tests for `validate_columns_present` are failing because the function does not exist in the `mltransform_one_hot_encoding` module. This validation logic is an implementation detail of the pipeline and its behavior is better tested via an end-to-end test like `test_pipeline_with_missing_columns`. These unit tests should be removed. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
