Merge pull request #7868 from vkarampudi/feature/filter-component

rwitcher · web-flow · commit 58435d68d561 · 2026-06-25T11:24:54.000-07:00
[Experimental] Add FilterComponent for row-level example filtering
diff --git a/tfx/components/experimental/filter/__init__.py 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
@@ -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
@@ -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
@@ -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
@@ -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()
