This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 0302bf5529c Implemented MLTransform generate vocab Dataflow benchmark
(#38215)
0302bf5529c is described below
commit 0302bf5529cda834cfabed642e6e8e5373d5a65f
Author: Abdelrahman Ibrahim <[email protected]>
AuthorDate: Thu May 28 15:47:20 2026 +0200
Implemented MLTransform generate vocab Dataflow benchmark (#38215)
* Implemented MLTransform generate vocab Dataflow benchmark
* applied anomaly transform test for reshuffled ZScore order
* remove redundant money nanos fallback
* resolved comments
---
.../beam_Inference_Python_Benchmarks_Dataflow.yml | 12 +
...s_Dataflow_MLTransform_Generate_Vocab_Batch.txt | 42 ++++
.test-infra/tools/refresh_looker_metrics.py | 3 +-
.../apache_beam/examples/ml_transform/README.md | 116 +++++++++
.../ml_transform/mltransform_generate_vocab.py | 265 +++++++++++++++++++++
.../mltransform_generate_vocab_requirements.txt | 26 ++
.../mltransform_generate_vocab_test.py | 191 +++++++++++++++
.../mltransform_generate_vocab_benchmark.py | 56 +++++
.../testing/load_tests/dataflow_cost_benchmark.py | 28 ++-
website/www/site/content/en/performance/_index.md | 1 +
.../en/performance/mltransformvocab/_index.md | 37 +++
website/www/site/data/performance.yaml | 16 ++
12 files changed, 791 insertions(+), 2 deletions(-)
diff --git a/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml
b/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml
index 61de2b54e73..36e962bc2ba 100644
--- a/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml
+++ b/.github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml
@@ -94,6 +94,7 @@ jobs:
${{ github.workspace
}}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_VLLM_Gemma_Batch.txt
${{ github.workspace
}}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Table_Row_Inference_Batch.txt
${{ github.workspace
}}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Table_Row_Inference_Stream.txt
+ ${{ github.workspace
}}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Generate_Vocab_Batch.txt
# The env variables are created and populated in the
test-arguments-action as
"<github.job>_test_arguments_<argument_file_paths_index>"
- name: get current time
run: echo "NOW_UTC=$(date '+%m%d%H%M%S' --utc)" >> $GITHUB_ENV
@@ -214,3 +215,14 @@ jobs:
-PpythonVersion=3.10 \
-PloadTest.requirementsTxtFile=apache_beam/ml/inference/table_row_inference_requirements.txt
\
'-PloadTest.args=${{
env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_10 }}
--autoscaling_algorithm=THROUGHPUT_BASED --max_num_workers=20
--metrics_table=result_table_row_inference_stream
--influx_measurement=result_table_row_inference_stream --mode=streaming
--input_subscription=projects/apache-beam-testing/subscriptions/table_row_inference_benchmark
--window_size_sec=60 --trigger_interval_sec=30 --timeout_ms=900000
--output_table=apache-beam-testing:beam_run_inf [...]
+ - name: run MLTransform Generate Vocab Batch
+ uses: ./.github/actions/gradle-command-self-hosted-action
+ timeout-minutes: 180
+ with:
+ gradle-command: :sdks:python:apache_beam:testing:load_tests:run
+ arguments: |
+
-PloadTest.mainClass=apache_beam.testing.benchmarks.inference.mltransform_generate_vocab_benchmark
\
+ -Prunner=DataflowRunner \
+ -PpythonVersion=3.10 \
+
-PloadTest.requirementsTxtFile=apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt
\
+ '-PloadTest.args=${{
env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_11 }}
--job_name=benchmark-tests-mltransform-generate-vocab-batch-${{env.NOW_UTC}}'
diff --git
a/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Generate_Vocab_Batch.txt
b/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Generate_Vocab_Batch.txt
new file mode 100644
index 00000000000..24723bf8f9c
--- /dev/null
+++
b/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_MLTransform_Generate_Vocab_Batch.txt
@@ -0,0 +1,42 @@
+# 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.
+#
+
+--project=apache-beam-testing
+--region=us-central1
+--runner=DataflowRunner
+--temp_location=gs://temp-storage-for-perf-tests/loadtests
+--staging_location=gs://temp-storage-for-perf-tests/loadtests
+--machine_type=n1-standard-4
+--disk_size_gb=100
+--num_workers=8
+--max_num_workers=16
+--autoscaling_algorithm=THROUGHPUT_BASED
+--worker_zone=us-central1-b
+--sdk_location=container
+--requirements_file=apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt
+--input_options={}
+--publish_to_big_query=true
+--metrics_dataset=beam_run_inference
+--metrics_table=mltransform_generate_vocab_batch
+--influx_measurement=mltransform_generate_vocab_batch
+--input_file=gs://apache-beam-ml/testing/inputs/sentences_50k.txt
+--output_vocab=gs://temp-storage-for-perf-tests/mltransform/vocab_outputs/mltransform_generate_vocab_batch
+--columns=text
+--vocab_size=50000
+--min_frequency=1
+--lowercase=true
+--input_expand_factor=1
+
diff --git a/.test-infra/tools/refresh_looker_metrics.py
b/.test-infra/tools/refresh_looker_metrics.py
index afd8ffa6f86..35122d5afe0 100644
--- a/.test-infra/tools/refresh_looker_metrics.py
+++ b/.test-infra/tools/refresh_looker_metrics.py
@@ -44,7 +44,8 @@ LOOKS_TO_DOWNLOAD = [
("85", ["268", "269", "270", "271", "272"]), # PyTorch Sentiment Batch
DistilBERT base uncased
("86", ["284", "285", "286", "287", "288"]), # VLLM Batch Gemma
("96", ["270", "304", "305", "353", "354"]), # Table Row Inference
Sklearn Batch
- ("106", ["355", "356", "357", "358", "359"]) # Table Row Inference
Sklearn Streaming
+ ("106", ["355", "356", "357", "358", "359"]), # Table Row Inference
Sklearn Streaming
+ ("107", ["360", "361", "362", "363", "364"]), # MLTransform Generate
Vocab Batch
]
def get_look(id: str) -> models.Look:
diff --git a/sdks/python/apache_beam/examples/ml_transform/README.md
b/sdks/python/apache_beam/examples/ml_transform/README.md
new file mode 100644
index 00000000000..a47fba1ff0c
--- /dev/null
+++ b/sdks/python/apache_beam/examples/ml_transform/README.md
@@ -0,0 +1,116 @@
+<!--
+Licensed 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.
+-->
+
+# MLTransform Examples
+
+This directory contains Apache Beam examples for MLTransform pipelines.
+
+## MLTransform - Generate Vocab (Batch only)
+
+`mltransform_generate_vocab.py` builds a vocabulary artifact from batch input
+rows using `MLTransform` + `ComputeAndApplyVocabulary`.
+
+### What it does
+
+1. Reads input rows from JSONL (`--input_file`) or BigQuery (`--input_table`).
+2. Extracts specified columns (`--columns`).
+3. Normalizes and combines text values (`trim`, optional lowercasing).
+4. Runs `ComputeAndApplyVocabulary` with top-k and min-frequency constraints
+ using space-delimited token splitting.
+5. Writes the vocabulary as one token per line.
+
+### Required arguments
+
+- `--output_vocab`
+- `--columns`
+- and one of:
+ - `--input_file`
+ - `--input_table`
+
+### Optional arguments
+
+- `--vocab_size` (default: `50000`)
+- `--min_frequency` (default: `1`)
+- `--lowercase` (default: `true`)
+- `--input_expand_factor` (default: `1`, useful for perf/load testing)
+
+### Local batch example
+
+```sh
+python -m apache_beam.examples.ml_transform.mltransform_generate_vocab \
+ --input_file=/tmp/input.jsonl \
+ --output_vocab=/tmp/vocab.txt \
+ --columns=text,category \
+ --vocab_size=5 \
+ --min_frequency=1 \
+ --lowercase=true \
+ --input_expand_factor=1 \
+ --runner=DirectRunner
+```
+
+### Input format
+
+JSONL input with object rows, for example:
+
+```json
+{"id":"1","text":"Beam beam ML pipeline"}
+{"id":"2","text":"Beam pipeline dataflow"}
+{"id":"3","text":"ML transform beam"}
+{"id":"4","text":"vocab vocab vocab test"}
+{"id":"5","text":"rare_token_once"}
+{"id":"6","text":""}
+{"id":"7","text":null}
+```
+
+The integration tests in `mltransform_generate_vocab_test.py` generate this
+sample data programmatically.
+
+### Output format
+
+One token per line:
+
+1. tokens follow the vocabulary order produced by `ComputeAndApplyVocabulary`.
+
+Example output:
+
+```txt
+beam
+ml
+```
+
+For this sample and config:
+
+```sh
+--columns=text --min_frequency=2 --vocab_size=3
+```
+
+the expected output is:
+
+```txt
+beam
+vocab
+ml
+```
+
+### Additional test datasets
+
+Test data for happy path and null/empty/missing columns is generated inline in
+`mltransform_generate_vocab_test.py`.
+
+### Performance testing pattern
+
+- Small local files: functional correctness and output-stability tests.
+- Large GCS files (or moderate file + `--input_expand_factor`): throughput/cost
+ benchmarking on Dataflow.
+
diff --git
a/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab.py
b/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab.py
new file mode 100644
index 00000000000..f0e4b38aa3f
--- /dev/null
+++
b/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab.py
@@ -0,0 +1,265 @@
+#
+# 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.
+#
+
+"""Batch-only vocabulary generation pipeline using MLTransform.
+
+This pipeline creates a vocabulary artifact from one or more input columns.
+
+Key properties:
+- Batch only (no streaming path).
+- Vocabulary generation via MLTransform ComputeAndApplyVocabulary.
+- Output format: one token per line.
+"""
+
+import argparse
+import json
+import logging
+import tempfile
+from typing import Any
+
+import apache_beam as beam
+from apache_beam.io.filesystems import FileSystems
+from apache_beam.ml.transforms.base import MLTransform
+from apache_beam.ml.transforms.tft import ComputeAndApplyVocabulary
+from apache_beam.options.pipeline_options import PipelineOptions
+
+
+def parse_bool_flag(value: str) -> bool:
+ value_lc = value.strip().lower()
+ if value_lc in ('1', 'true', 't', 'yes', 'y'):
+ return True
+ if value_lc in ('0', 'false', 'f', 'no', 'n'):
+ return False
+ raise ValueError(
+ f'Invalid boolean value {value!r}. Expected true/false style value.')
+
+
+def normalize_text(value: Any, lowercase: bool = True) -> str:
+ if value is None:
+ return ''
+ text = str(value).strip()
+ if lowercase:
+ text = text.lower()
+ return text
+
+
+def _parse_json_line(line: str) -> dict[str, Any]:
+ try:
+ parsed = json.loads(line)
+ if isinstance(parsed, dict):
+ return parsed
+ except json.JSONDecodeError:
+ pass
+ # Treat plain-text or non-object rows as values for the default "text"
column.
+ return {'text': line}
+
+
+def _extract_column_values(row: dict[str, Any],
+ columns: list[str]) -> list[str]:
+ values: list[str] = []
+ for col in columns:
+ if col not in row:
+ continue
+ val = row[col]
+ if val is None:
+ continue
+ if isinstance(val, list):
+ values.extend(str(item) for item in val if item is not None)
+ else:
+ values.append(str(val))
+ return values
+
+
+def _build_vocab_text(
+ row: dict[str, Any], columns: list[str], lowercase: bool) -> str:
+ values = _extract_column_values(row, columns)
+ normalized_values = [
+ normalize_text(value, lowercase=lowercase) for value in values
+ ]
+ non_empty_values = [value for value in normalized_values if value]
+ return ' '.join(non_empty_values)
+
+
+def _resolve_vocab_asset_path(
+ artifact_location: str, vocab_filename: str, column_name: str) -> str:
+ asset_name = f'{vocab_filename}_{column_name}'
+ pattern = (
+ f'{artifact_location.rstrip("/")}'
+ f'/*/transform_fn/assets/{asset_name}')
+ try:
+ match_results = FileSystems.match([pattern])
+ matches = match_results[0].metadata_list if match_results else []
+ except Exception as e:
+ raise ValueError(
+ f'Could not locate vocabulary artifact {asset_name!r} under '
+ f'{artifact_location!r} due to error: {e}') from e
+ if not matches:
+ raise ValueError(
+ f'Could not locate vocabulary artifact {asset_name!r} under '
+ f'{artifact_location!r}.')
+ return matches[0].path
+
+
+def _read_vocab_tokens(vocab_asset_path: str) -> list[str]:
+ tokens = []
+ with FileSystems.open(vocab_asset_path) as f:
+ for raw_line in f:
+ token = raw_line.decode('utf-8').rstrip('\r\n')
+ if token:
+ tokens.append(token)
+ return tokens
+
+
+def _write_vocab_file(output_path: str, tokens: list[str]) -> None:
+ with FileSystems.create(output_path) as f:
+ for token in tokens:
+ f.write((token + '\n').encode('utf-8'))
+
+
+def parse_known_args(argv):
+ parser = argparse.ArgumentParser(
+ description='Generate vocabulary from batch input with MLTransform.')
+ parser.add_argument('--input_file', help='Input JSONL file path.')
+ parser.add_argument(
+ '--input_table',
+ help='Input BigQuery table path in PROJECT:DATASET.TABLE format.')
+ parser.add_argument('--output_vocab', help='Output vocab file prefix/path.')
+ parser.add_argument(
+ '--columns',
+ help='Comma-separated source columns to include in vocabulary.')
+ parser.add_argument(
+ '--vocab_size',
+ type=int,
+ default=50000,
+ help='Maximum vocabulary size (top-K by frequency).')
+ parser.add_argument(
+ '--min_frequency',
+ type=int,
+ default=1,
+ help='Minimum token frequency required to keep token.')
+ parser.add_argument(
+ '--lowercase',
+ default='true',
+ help='Whether to lowercase text before vocabulary generation.')
+ parser.add_argument(
+ '--input_expand_factor',
+ type=int,
+ default=1,
+ help=(
+ 'Batch-only: repeat each input line this many times to scale volume '
+ 'for load/perf testing.'))
+ parser.add_argument(
+ '--artifact_location',
+ default='',
+ help=(
+ 'Artifact directory for MLTransform output. If empty, a temporary '
+ 'local directory is used.'))
+ return parser.parse_known_args(argv)
+
+
+def validate_args(args) -> list[str]:
+ has_input_file = bool(args.input_file)
+ has_input_table = bool(args.input_table)
+ if not has_input_file and not has_input_table:
+ raise ValueError('One of --input_file or --input_table is required.')
+ if has_input_file and has_input_table:
+ raise ValueError('Use exactly one of --input_file or --input_table.')
+ if not args.output_vocab:
+ raise ValueError('--output_vocab is required.')
+ if not args.columns:
+ raise ValueError('--columns is required.')
+ if args.vocab_size is None or args.vocab_size <= 0:
+ raise ValueError('--vocab_size must be > 0.')
+ if args.min_frequency is None or args.min_frequency < 1:
+ raise ValueError('--min_frequency must be >= 1.')
+ if args.input_expand_factor is None or args.input_expand_factor < 1:
+ raise ValueError('--input_expand_factor must be >= 1.')
+ seen = set()
+ deduped_cols = []
+ for col in (c.strip() for c in args.columns.split(',')):
+ if col and col not in seen:
+ seen.add(col)
+ deduped_cols.append(col)
+ return deduped_cols
+
+
+def run(argv=None, test_pipeline=None):
+ known_args, pipeline_args = parse_known_args(argv)
+ columns = validate_args(known_args)
+ lowercase = parse_bool_flag(known_args.lowercase)
+
+ options = PipelineOptions(pipeline_args)
+ artifact_location = known_args.artifact_location
+ if not artifact_location:
+ temp_location = options.get_all_options().get('temp_location')
+ if temp_location and (temp_location.startswith('gs://') or
+ temp_location.startswith('s3://')):
+ artifact_location = (
+ f'{temp_location.rstrip("/")}/mltransform_vocab_artifacts')
+ else:
+ artifact_location = tempfile.mkdtemp(
+ prefix='mltransform_generate_vocab_artifacts_')
+
+ pipeline = test_pipeline or beam.Pipeline(options=options)
+
+ if known_args.input_file:
+ lines = (
+ pipeline
+ | 'ReadInputFile' >> beam.io.ReadFromText(known_args.input_file))
+ if known_args.input_expand_factor > 1:
+ lines = (
+ lines
+ | 'ExpandInputForPerf' >> beam.FlatMap(
+ lambda line, n: [line] * n, known_args.input_expand_factor))
+ rows = lines | 'ParseJSON' >> beam.Map(_parse_json_line)
+ else:
+ rows = pipeline | 'ReadInputTable' >> beam.io.ReadFromBigQuery(
+ table=known_args.input_table)
+
+ vocab_text = (
+ rows
+ | 'BuildVocabText' >>
+ beam.Map(lambda row: _build_vocab_text(row, columns, lowercase))
+ | 'DropEmptyText' >> beam.Filter(bool))
+
+ _ = (
+ vocab_text
+ | 'MLTransformInput' >> beam.Map(lambda text: {'text': text})
+ | 'ApplyMLTransform' >>
+ MLTransform(write_artifact_location=artifact_location).with_transform(
+ ComputeAndApplyVocabulary(
+ columns=['text'],
+ top_k=known_args.vocab_size,
+ frequency_threshold=known_args.min_frequency,
+ split_string_by_delimiter=' ',
+ vocab_filename='vocab')))
+
+ result = pipeline.run()
+ result.wait_until_finish()
+
+ vocab_tokens = _read_vocab_tokens(
+ _resolve_vocab_asset_path(
+ artifact_location=artifact_location,
+ vocab_filename='vocab',
+ column_name='text'))
+ _write_vocab_file(known_args.output_vocab, vocab_tokens)
+ return result
+
+
+if __name__ == '__main__':
+ logging.getLogger().setLevel(logging.INFO)
+ run()
diff --git
a/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt
b/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt
new file mode 100644
index 00000000000..4e00b3b4316
--- /dev/null
+++
b/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab_requirements.txt
@@ -0,0 +1,26 @@
+# 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.
+#
+
+# Dataflow worker requirements for mltransform_generate_vocab load tests.
+# MLTransform TFT operations need a consistent TensorFlow Transform stack;
+# otherwise workers can crash-loop with pandas/numpy ABI mismatches.
+google-cloud-monitoring>=2.27.0
+tensorflow_transform>=1.14.0,<1.15.0
+tensorflow-metadata>=1.14.0,<1.15.0
+tfx-bsl>=1.14.0,<1.15.0
+numpy<2
+pandas<2
+dill
diff --git
a/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab_test.py
b/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab_test.py
new file mode 100644
index 00000000000..3e29bd3aae1
--- /dev/null
+++
b/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab_test.py
@@ -0,0 +1,191 @@
+#
+# 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.
+#
+
+import json
+import os
+import tempfile
+import unittest
+
+try:
+ from apache_beam.examples.ml_transform import mltransform_generate_vocab
+except ImportError:
+ raise unittest.SkipTest('tensorflow_transform is not installed.')
+
+
+class MLTransformGenerateVocabUnitTest(unittest.TestCase):
+ def test_normalize_text(self):
+ text = mltransform_generate_vocab.normalize_text(' Hello Beam ', True)
+ self.assertEqual(text, 'hello beam')
+
+ def test_null_and_empty_handling_helpers(self):
+ normalized_none = mltransform_generate_vocab.normalize_text(None, True)
+ self.assertEqual(normalized_none, '')
+ self.assertEqual(
+ mltransform_generate_vocab._build_vocab_text(
+ {'text': ['Beam', None, ' ', 'Flow']}, ['text'], lowercase=True),
+ 'beam flow')
+
+
+class MLTransformGenerateVocabCliValidationTest(unittest.TestCase):
+ def test_missing_required_args(self):
+ args, _ = mltransform_generate_vocab.parse_known_args([])
+ with self.assertRaisesRegex(ValueError, 'input_file or --input_table'):
+ mltransform_generate_vocab.validate_args(args)
+
+ def test_invalid_numeric_values(self):
+ args, _ = mltransform_generate_vocab.parse_known_args([
+ '--input_file=a.jsonl',
+ '--output_vocab=/tmp/vocab',
+ '--columns=text',
+ '--vocab_size=0',
+ '--min_frequency=0',
+ ])
+ with self.assertRaisesRegex(ValueError, 'vocab_size'):
+ mltransform_generate_vocab.validate_args(args)
+
+ def test_invalid_input_expand_factor(self):
+ args, _ = mltransform_generate_vocab.parse_known_args([
+ '--input_file=a.jsonl',
+ '--output_vocab=/tmp/vocab',
+ '--columns=text',
+ '--input_expand_factor=0',
+ ])
+ with self.assertRaisesRegex(ValueError, 'input_expand_factor'):
+ mltransform_generate_vocab.validate_args(args)
+
+
+class MLTransformGenerateVocabIntegrationTest(unittest.TestCase):
+ def test_batch_pipeline_exact_output_order(self):
+ with tempfile.TemporaryDirectory() as tmpdir:
+ input_path = os.path.join(tmpdir, 'input.jsonl')
+ output_prefix = os.path.join(tmpdir, 'vocab.txt')
+
+ rows = [
+ {
+ 'id': '1', 'text': 'Beam beam ML pipeline'
+ },
+ {
+ 'id': '2', 'text': 'Beam pipeline dataflow'
+ },
+ {
+ 'id': '3', 'text': 'ML transform beam'
+ },
+ {
+ 'id': '4', 'text': 'vocab vocab vocab test'
+ },
+ {
+ 'id': '5', 'text': 'rare_token_once'
+ },
+ {
+ 'id': '6', 'text': ''
+ },
+ {
+ 'id': '7', 'text': None
+ },
+ ]
+ with open(input_path, 'w', encoding='utf-8') as f:
+ for row in rows:
+ f.write(json.dumps(row) + '\n')
+
+ mltransform_generate_vocab.run([
+ f'--input_file={input_path}',
+ f'--output_vocab={output_prefix}',
+ '--columns=text',
+ '--vocab_size=3',
+ '--min_frequency=2',
+ '--lowercase=true',
+ '--runner=DirectRunner',
+ ])
+
+ output_path = output_prefix
+ with open(output_path, 'r', encoding='utf-8') as f:
+ output_tokens = [line.rstrip('\n') for line in f]
+
+ self.assertEqual(set(output_tokens), {'beam', 'vocab', 'ml'})
+ self.assertEqual(len(output_tokens), 3)
+
+ def test_output_is_stable_across_runs(self):
+ with tempfile.TemporaryDirectory() as tmpdir:
+ input_path = os.path.join(tmpdir, 'input.jsonl')
+ output_prefix = os.path.join(tmpdir, 'vocab.txt')
+ rows = [
+ {
+ 'text': 'apple banana'
+ },
+ {
+ 'text': 'banana apple'
+ },
+ {
+ 'text': 'cat dog'
+ },
+ {
+ 'text': 'dog cat'
+ },
+ ]
+ with open(input_path, 'w', encoding='utf-8') as f:
+ for row in rows:
+ f.write(json.dumps(row) + '\n')
+
+ common_args = [
+ f'--input_file={input_path}',
+ '--columns=text',
+ '--vocab_size=4',
+ '--min_frequency=2',
+ '--lowercase=true',
+ '--runner=DirectRunner',
+ ]
+
+ mltransform_generate_vocab.run(
+ common_args + [f'--output_vocab={output_prefix}'])
+
+ output_path = output_prefix
+ with open(output_path, 'r', encoding='utf-8') as f:
+ first_run_tokens = [line.rstrip('\n') for line in f]
+
+ output_prefix_second = os.path.join(tmpdir, 'vocab_second.txt')
+ mltransform_generate_vocab.run(
+ common_args + [f'--output_vocab={output_prefix_second}'])
+
+ with open(output_prefix_second, 'r', encoding='utf-8') as f:
+ second_run_tokens = [line.rstrip('\n') for line in f]
+
+ self.assertEqual(first_run_tokens, second_run_tokens)
+
+ def test_empty_filtered_result_writes_empty_vocab(self):
+ with tempfile.TemporaryDirectory() as tmpdir:
+ input_path = os.path.join(tmpdir, 'input.jsonl')
+ output_prefix = os.path.join(tmpdir, 'vocab.txt')
+ with open(input_path, 'w', encoding='utf-8') as f:
+ f.write(json.dumps({'text': 'beam'}) + '\n')
+
+ mltransform_generate_vocab.run([
+ f'--input_file={input_path}',
+ f'--output_vocab={output_prefix}',
+ '--columns=text',
+ '--vocab_size=10',
+ '--min_frequency=2',
+ '--runner=DirectRunner',
+ ])
+
+ output_path = output_prefix
+ with open(output_path, 'r', encoding='utf-8') as f:
+ output_tokens = [line.rstrip('\n') for line in f]
+ self.assertEqual(output_tokens, [])
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git
a/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_generate_vocab_benchmark.py
b/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_generate_vocab_benchmark.py
new file mode 100644
index 00000000000..98a7b59191d
--- /dev/null
+++
b/sdks/python/apache_beam/testing/benchmarks/inference/mltransform_generate_vocab_benchmark.py
@@ -0,0 +1,56 @@
+#
+# 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.
+#
+
+"""Benchmark test for the batch-only MLTransform Generate Vocab pipeline."""
+
+import logging
+
+from apache_beam.examples.ml_transform import mltransform_generate_vocab
+from apache_beam.testing.load_tests.dataflow_cost_benchmark import
DataflowCostBenchmark
+
+
+class MLTransformGenerateVocabBenchmarkTest(DataflowCostBenchmark):
+ """Runs the MLTransform vocab generation pipeline as a cost benchmark."""
+ def __init__(self):
+ # Namespace used by benchmark dashboards.
+ self.metrics_namespace = 'BeamML_MLTransformVocab'
+ # Monitor throughput on text rows after empty strings are filtered.
+ super().__init__(
+ metrics_namespace=self.metrics_namespace,
+ is_streaming=False,
+ pcollection='DropEmptyText.out0')
+
+ def test(self):
+ """Execute the MLTransform vocab generation pipeline for benchmarking."""
+ extra_opts = {
+ 'input_file': self.pipeline.get_option('input_file'),
+ 'output_vocab': self.pipeline.get_option('output_vocab'),
+ 'artifact_location': self.pipeline.get_option('artifact_location'),
+ 'columns': self.pipeline.get_option('columns'),
+ 'vocab_size': self.pipeline.get_option('vocab_size'),
+ 'min_frequency': self.pipeline.get_option('min_frequency'),
+ 'lowercase': self.pipeline.get_option('lowercase'),
+ 'input_expand_factor': self.pipeline.get_option('input_expand_factor'),
+ }
+
+ self.result = mltransform_generate_vocab.run(
+ self.pipeline.get_full_options_as_args(**extra_opts))
+
+
+if __name__ == '__main__':
+ logging.basicConfig(level=logging.INFO)
+ MLTransformGenerateVocabBenchmarkTest().run()
diff --git
a/sdks/python/apache_beam/testing/load_tests/dataflow_cost_benchmark.py
b/sdks/python/apache_beam/testing/load_tests/dataflow_cost_benchmark.py
index 3df96c03a8a..3300831b8ea 100644
--- a/sdks/python/apache_beam/testing/load_tests/dataflow_cost_benchmark.py
+++ b/sdks/python/apache_beam/testing/load_tests/dataflow_cost_benchmark.py
@@ -193,6 +193,32 @@ class DataflowCostBenchmark(LoadTest):
"""Query Cloud Monitoring for per-PCollection throughput."""
name = (
pcollection_name if pcollection_name is not None else self.pcollection)
+
+ def _point_numeric_value(point) -> float:
+ value = point.value
+ # point.value is proto-plus; use the underlying protobuf oneof if
present.
+ msg = getattr(value, '_pb', value)
+ if msg is not None:
+ try:
+ active_field = msg.WhichOneof('value')
+ except AttributeError:
+ active_field = None
+ if active_field == 'double_value':
+ return float(value.double_value)
+ if active_field == 'int64_value':
+ return float(value.int64_value)
+ if active_field == 'distribution_value':
+ # Use aligned mean for distribution-valued points.
+ distribution = value.distribution_value
+ if distribution.count > 0:
+ return float(distribution.mean)
+ return 0.0
+ if active_field == 'money_value':
+ money = value.money_value
+ nanos = getattr(money, 'nanos', 0)
+ return float(money.units) + (float(nanos) / 1_000_000_000.0)
+ return 0.0
+
interval = monitoring_v3.TimeInterval(
start_time=start_time, end_time=end_time)
aggregation = monitoring_v3.Aggregation(
@@ -221,7 +247,7 @@ class DataflowCostBenchmark(LoadTest):
for key, req in requests.items():
time_series = self.monitoring_client.list_time_series(request=req)
values = [
- point.value.double_value for series in time_series
+ _point_numeric_value(point) for series in time_series
for point in series.points
]
metrics[f"AvgThroughput{key}"] = sum(values) / len(
diff --git a/website/www/site/content/en/performance/_index.md
b/website/www/site/content/en/performance/_index.md
index 1624c58efe2..350524a3c86 100644
--- a/website/www/site/content/en/performance/_index.md
+++ b/website/www/site/content/en/performance/_index.md
@@ -59,3 +59,4 @@ See the following pages for performance measures recorded
when running various B
- [TensorFlow MNIST Image Classification](/performance/tensorflowmnist)
- [VLLM Gemma Batch Completion Tesla T4 GPU](/performance/vllmgemmabatchtesla)
- [Table Row Inference Sklearn Batch](/performance/tablerowinference)
+- [MLTransform Generate Vocab (batch)](/performance/mltransformvocab)
diff --git a/website/www/site/content/en/performance/mltransformvocab/_index.md
b/website/www/site/content/en/performance/mltransformvocab/_index.md
new file mode 100644
index 00000000000..f7844a5a09a
--- /dev/null
+++ b/website/www/site/content/en/performance/mltransformvocab/_index.md
@@ -0,0 +1,37 @@
+---
+title: "MLTransform Generate Vocab Performance"
+---
+
+<!--
+Licensed 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
+-->
+
+# MLTransform Generate Vocab Performance
+
+**Model**: Apache Beam MLTransform — ComputeAndApplyVocabulary (Generate Vocab)
+**Accelerator**: CPU (batch)
+**Host**: MLTransform vocab pipeline on Dataflow (batch)
+
+This batch pipeline computes a vocabulary from input text columns using
+`ComputeAndApplyVocabulary` and writes vocabulary artifacts for downstream use.
+
+See the [glossary](/performance/glossary) for definitions.
+
+Full pipeline implementation is available here:
+https://github.com/apache/beam/blob/master/sdks/python/apache_beam/examples/ml_transform/mltransform_generate_vocab.py
+
+## What is the estimated cost to run the pipeline?
+
+{{< performance_looks io="mltransformvocab" read_or_write="write"
section="cost" >}}
+
+## How has various metrics changed when running the pipeline for different
Beam SDK versions?
+
+{{< performance_looks io="mltransformvocab" read_or_write="write"
section="version" >}}
+
+## How has various metrics changed over time when running the pipeline?
+
+{{< performance_looks io="mltransformvocab" read_or_write="write"
section="date" >}}
+
diff --git a/website/www/site/data/performance.yaml
b/website/www/site/data/performance.yaml
index 8841bbb58ec..9725a4af99a 100644
--- a/website/www/site/data/performance.yaml
+++ b/website/www/site/data/performance.yaml
@@ -283,3 +283,19 @@ looks:
title: AvgThroughputBytesPerSec by Version
- id: P7wKZy6tQFWbbDfm4HzfCJnsQrVgfGsJ
title: AvgThroughputElementsPerSec by Version
+ mltransformvocab:
+ write:
+ folder: 107
+ cost:
+ - id: CccNCHQ23Js2tmxTDYVnjxwfx2zSFSjY
+ title: RunTime and EstimatedCost
+ date:
+ - id: n7Gn7c7KSW52ZgXCQXndHgWtjVPK8W33
+ title: AvgThroughputBytesPerSec by Date
+ - id: 3tcg6cSh6tBg6FmpTxgzyMtvFqWKkQn3
+ title: AvgThroughputElementsPerSec by Date
+ version:
+ - id: tPYwJngPDBsKjK3DNgGHGsCyTpxdfmTB
+ title: AvgThroughputBytesPerSec by Version
+ - id: cC75NnCbQT3mQmKVHtDxzptXpwPb64qz
+ title: AvgThroughputElementsPerSec by Version
