diff --git a/tfx/components/experimental/filter/__init__.py b/tfx/components/experimental/filter/__init__.py new file mode 100644 index 0000000000..55de8dfb39 --- /dev/null +++ b/tfx/components/experimental/filter/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2026 Google LLC. All Rights Reserved. +# +# 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. +"""Filter component experimental module.""" diff --git a/tfx/components/experimental/filter/component.py b/tfx/components/experimental/filter/component.py new file mode 100644 index 0000000000..93d7e9df3e --- /dev/null +++ b/tfx/components/experimental/filter/component.py @@ -0,0 +1,83 @@ +# Copyright 2026 Google LLC. All Rights Reserved. +# +# 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. +"""TFX experimental Filter component definition.""" + +from typing import Optional + +from tfx import types +from tfx.components.experimental.filter import executor +from tfx.dsl.components.base import base_component +from tfx.dsl.components.base import executor_spec +from tfx.types import standard_artifacts +from tfx.types.component_spec import ChannelParameter +from tfx.types.component_spec import ComponentSpec +from tfx.types.component_spec import ExecutionParameter + + +class FilterSpec(ComponentSpec): + """Filter component spec.""" + + PARAMETERS = { + 'filter_fn_path': ExecutionParameter(type=str), + } + INPUTS = { + 'examples': ChannelParameter(type=standard_artifacts.Examples), + } + OUTPUTS = { + 'filtered_examples': ChannelParameter(type=standard_artifacts.Examples), + } + + +class FilterComponent(base_component.BaseComponent): + """A TFX component to filter examples based on a user-defined function. + + The FilterComponent reads examples from each split of the input `examples` + artifact, applies a user-defined filter function using an Apache Beam + pipeline, and writes the filtered examples to the `filtered_examples` output + artifact, preserving the split structure. + + Example usage: + ```python + # Filter out examples where age <= 18 + filter_component = FilterComponent( + examples=example_gen.outputs['examples'], + filter_fn_path='my_filters.custom_filter_fn' + ) + ``` + """ + + SPEC_CLASS = FilterSpec + EXECUTOR_SPEC = executor_spec.ExecutorClassSpec(executor.Executor) + + def __init__(self, + examples: types.BaseChannel, + filter_fn_path: str, + filtered_examples: Optional[types.Channel] = None): + """Construct a FilterComponent. + + Args: + examples: A [BaseChannel] of type [standard_artifacts.Examples]. + filter_fn_path: The Python import path to the filter function. + e.g., 'my_module.my_filter_fn'. The function must have the signature: + `def my_filter_fn(serialized_example: bytes) -> bool` + filtered_examples: Optional output channel of type [standard_artifacts.Examples]. + """ + if filtered_examples is None: + filtered_examples = types.Channel(type=standard_artifacts.Examples) + + spec = FilterSpec( + examples=examples, + filter_fn_path=filter_fn_path, + filtered_examples=filtered_examples) + super().__init__(spec=spec) diff --git a/tfx/components/experimental/filter/component_test.py b/tfx/components/experimental/filter/component_test.py new file mode 100644 index 0000000000..d56f8e5633 --- /dev/null +++ b/tfx/components/experimental/filter/component_test.py @@ -0,0 +1,53 @@ +# Copyright 2026 Google LLC. All Rights Reserved. +# +# 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. +"""Tests for tfx.components.experimental.filter.component.""" + +import tensorflow as tf +from tfx.components.experimental.filter import component +from tfx.types import standard_artifacts +from tfx.types import channel + + +class ComponentTest(tf.test.TestCase): + + def testConstruct(self): + examples = channel.Channel(type=standard_artifacts.Examples) + filter_fn_path = 'my_module.my_filter_fn' + + filter_component = component.FilterComponent( + examples=examples, + filter_fn_path=filter_fn_path + ) + + # Verify input channel + self.assertEqual( + filter_component.inputs['examples'].type, + standard_artifacts.Examples + ) + + # Verify output channel + self.assertEqual( + filter_component.outputs['filtered_examples'].type, + standard_artifacts.Examples + ) + + # Verify parameter + self.assertEqual( + filter_component.exec_properties['filter_fn_path'], + filter_fn_path + ) + + +if __name__ == '__main__': + tf.test.main() diff --git a/tfx/components/experimental/filter/executor.py b/tfx/components/experimental/filter/executor.py new file mode 100644 index 0000000000..6450ba78d3 --- /dev/null +++ b/tfx/components/experimental/filter/executor.py @@ -0,0 +1,90 @@ +# Copyright 2026 Google LLC. All Rights Reserved. +# +# 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. +"""TFX experimental Filter component executor.""" + +import os +from typing import Any, Dict, List + +from absl import logging +import apache_beam as beam +import tensorflow as tf +from tfx import types +from tfx.dsl.components.base import base_beam_executor +from tfx.types import artifact_utils +from tfx.utils import import_utils + + +class Executor(base_beam_executor.BaseBeamExecutor): + """TFX experimental Filter component executor.""" + + def Do(self, input_dict: Dict[str, List[types.Artifact]], + output_dict: Dict[str, List[types.Artifact]], + exec_properties: Dict[str, Any]) -> None: + """Runs the filter Apache Beam pipeline. + + Args: + input_dict: Input dict from input key to a list of Artifacts. + - examples: A list of type `standard_artifacts.Examples` containing + the splits to be filtered. + output_dict: Output dict from output key to a list of Artifacts. + - filtered_examples: A list of type `standard_artifacts.Examples` + where the filtered splits will be written. + exec_properties: A dict of execution properties. + - filter_fn_path: The Python import path to the filter function. + """ + self._log_startup(input_dict, output_dict, exec_properties) + + examples = artifact_utils.get_single_instance(input_dict['examples']) + filtered_examples = artifact_utils.get_single_instance( + output_dict['filtered_examples']) + + # Setup output splits. + split_names = artifact_utils.decode_split_names(examples.split_names) + filtered_examples.split_names = artifact_utils.encode_split_names( + split_names) + filtered_examples.span = examples.span + filtered_examples.version = examples.version + + # Import the user-defined filter function. + filter_fn_path = exec_properties['filter_fn_path'] + logging.info('Importing user filter function from: %s', filter_fn_path) + filter_fn = import_utils.import_class_by_path(filter_fn_path) + + with self._make_beam_pipeline() as pipeline: + for split in split_names: + input_split_uri = artifact_utils.get_split_uri([examples], split) + output_split_uri = artifact_utils.get_split_uri([filtered_examples], + split) + + # Ensure output split directory exists. + tf.io.gfile.makedirs(output_split_uri) + + input_pattern = os.path.join(input_split_uri, '*') + output_prefix = os.path.join(output_split_uri, 'data_tfrecord') + + logging.info('Filtering split %s. Reading from %s, writing to prefix %s', + split, input_pattern, output_prefix) + + # Run the Beam pipeline to read, filter, and write the split. + _ = ( + pipeline + | f'ReadFromTFRecord[{split}]' >> beam.io.ReadFromTFRecord( + input_pattern) + | f'FilterExamples[{split}]' >> beam.Filter(filter_fn) + | f'WriteToTFRecord[{split}]' >> beam.io.WriteToTFRecord( + output_prefix, + file_name_suffix='.gz') + ) + + logging.info('FilterComponent execution completed successfully.') diff --git a/tfx/components/experimental/filter/executor_test.py b/tfx/components/experimental/filter/executor_test.py new file mode 100644 index 0000000000..57e0d951f7 --- /dev/null +++ b/tfx/components/experimental/filter/executor_test.py @@ -0,0 +1,126 @@ +# Copyright 2026 Google LLC. All Rights Reserved. +# +# 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. +"""Tests for tfx.components.experimental.filter.executor.""" + +import os +from typing import List +import tensorflow as tf +from tfx.components.experimental.filter import executor + +from tfx.types import standard_artifacts +from tfx.types import artifact_utils + + +def dummy_filter_fn(serialized_example: bytes) -> bool: + """A simple filter function that parses the example and filters by age.""" + example = tf.train.Example() + example.ParseFromString(serialized_example) + features = example.features.feature + if 'age' in features: + return features['age'].int64_list.value[0] > 18 + return False + + +class ExecutorTest(tf.test.TestCase): + + def setUp(self): + super().setUp() + self._output_data_dir = os.path.join( + os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()), + self._testMethodName) + self._input_data_dir = os.path.join(self.get_temp_dir(), 'input') + + def _create_test_examples(self, examples_artifact: standard_artifacts.Examples, split_name: str, values: List[int]): + """Creates a TFRecord file with test examples for the given split.""" + split_dir = artifact_utils.get_split_uri([examples_artifact], split_name) + tf.io.gfile.makedirs(split_dir) + file_path = os.path.join(split_dir, 'data.tfrecord.gz') + + options = tf.io.TFRecordOptions(compression_type='GZIP') + with tf.io.TFRecordWriter(file_path, options=options) as writer: + for val in values: + example = tf.train.Example() + example.features.feature['age'].int64_list.value.append(val) + writer.write(example.SerializeToString()) + + def testExecutor(self): + # 1. Prepare input and output examples artifacts. + examples_artifact = standard_artifacts.Examples() + examples_artifact.uri = self._input_data_dir + examples_artifact.split_names = artifact_utils.encode_split_names( + ['train', 'eval']) + + # 2. Create input data: + # train split has ages [10, 20, 30] -> filtered should have [20, 30] + # eval split has ages [15, 25] -> filtered should have [25] + self._create_test_examples(examples_artifact, 'train', [10, 20, 30]) + self._create_test_examples(examples_artifact, 'eval', [15, 25]) + + filtered_examples_artifact = standard_artifacts.Examples() + filtered_examples_artifact.uri = self._output_data_dir + + input_dict = {'examples': [examples_artifact]} + output_dict = {'filtered_examples': [filtered_examples_artifact]} + + # Full python import path to the dummy_filter_fn + filter_fn_path = ( + 'tfx.components.experimental.filter.executor_test.dummy_filter_fn') + + exec_properties = {'filter_fn_path': filter_fn_path} + + # 3. Run the executor. + filter_executor = executor.Executor() + filter_executor.Do(input_dict, output_dict, exec_properties) + + # 4. Verify output splits. + decoded_splits = artifact_utils.decode_split_names( + filtered_examples_artifact.split_names) + self.assertEqual(decoded_splits, ['train', 'eval']) + + # 5. Verify the content of the filtered train split. + train_output_dir = artifact_utils.get_split_uri( + [filtered_examples_artifact], 'train') + train_output_files = tf.io.gfile.glob(os.path.join(train_output_dir, '*')) + self.assertNotEmpty(train_output_files) + + train_ages = [] + # Read the output TFRecords. Beam writes sharded files, so we read all matching files. + for file_path in train_output_files: + raw_dataset = tf.data.TFRecordDataset(file_path, compression_type='GZIP') + for raw_record in raw_dataset: + example = tf.train.Example() + example.ParseFromString(raw_record.numpy()) + train_ages.append(example.features.feature['age'].int64_list.value[0]) + + self.assertCountEqual(train_ages, [20, 30]) + + # 6. Verify the content of the filtered eval split. + eval_output_dir = artifact_utils.get_split_uri( + [filtered_examples_artifact], 'eval') + eval_output_files = tf.io.gfile.glob(os.path.join(eval_output_dir, '*')) + self.assertNotEmpty(eval_output_files) + + eval_ages = [] + for file_path in eval_output_files: + raw_dataset = tf.data.TFRecordDataset(file_path, compression_type='GZIP') + for raw_record in raw_dataset: + example = tf.train.Example() + example.ParseFromString(raw_record.numpy()) + eval_ages.append(example.features.feature['age'].int64_list.value[0]) + + self.assertCountEqual(eval_ages, [25]) + + +if __name__ == '__main__': + tf.test.main()