[yaml] add support for matchall (#38512)

* add support for matchall

* fix lint

* fix more lint

* fix gemini comments

* move matchAll to standard providers

Author: derrickaw

sdks/python/apache_beam/yaml/standard_io.yaml

@@ -116,6 +116,7 @@
     'ReadFromTFRecord': 'apache_beam.yaml.yaml_io.read_from_tfrecord'
     'WriteToTFRecord': 'apache_beam.yaml.yaml_io.write_to_tfrecord'
 
+
 # General File Formats
 # Declared as a renaming transform to avoid exposing all
 # (implementation-specific) pandas arguments and aligning with possible Java

sdks/python/apache_beam/yaml/standard_providers.yaml

@@ -57,6 +57,7 @@
   transforms:
     MLTransform: 'apache_beam.yaml.yaml_ml.ml_transform'
     RunInference: 'apache_beam.yaml.yaml_ml.run_inference'
+    MatchAll: 'apache_beam.yaml.yaml_io.match_all'
 
 - type: renaming
   transforms:

sdks/python/apache_beam/yaml/tests/match_all.yaml

@@ -0,0 +1,96 @@
+#
+# 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.
+#
+
+fixtures:
+  - name: TEMP_DIR
+    type: "tempfile.TemporaryDirectory"
+
+pipelines:
+  # 1. Write first matching file to temp directory
+  - pipeline:
+      type: chain
+      transforms:
+        - type: Create
+          config:
+            elements:
+              - {value: "a"}
+        - type: WriteToText
+          config:
+            path: "{TEMP_DIR}/match1.txt"
+
+  # 2. Write second matching file to temp directory
+  - pipeline:
+      type: chain
+      transforms:
+        - type: Create
+          config:
+            elements:
+              - {value: "b"}
+        - type: WriteToText
+          config:
+            path: "{TEMP_DIR}/match2.txt"
+
+  # 3. Write an ignore file to temp directory (should not be matched)
+  - pipeline:
+      type: chain
+      transforms:
+        - type: Create
+          config:
+            elements:
+              - {value: "c"}
+        - type: WriteToText
+          config:
+            path: "{TEMP_DIR}/ignore.txt"
+
+  # 4. Match files using MatchAll and verify only match1 and match2 are found
+  - pipeline:
+      type: chain
+      transforms:
+        - type: Create
+          config:
+            elements:
+              - {pattern: "{TEMP_DIR}/match*.txt*"}
+        - type: MatchAll
+          config:
+            file_pattern: pattern
+        - type: MapToFields
+          config:
+            language: python
+            fields:
+              filename: "path.split('/')[-1].split('-')[0]"
+        - type: AssertEqual
+          config:
+            elements:
+              - {filename: "match1.txt"}
+              - {filename: "match2.txt"}
+
+  # 5. Match non-existent files with ALLOW empty_match_treatment (should succeed but return empty PCollection)
+  - pipeline:
+      type: chain
+      transforms:
+        - type: Create
+          config:
+            elements:
+              - {pattern: "{TEMP_DIR}/does_not_exist*.txt*"}
+        - type: MatchAll
+          config:
+            file_pattern: pattern
+            empty_match_treatment: ALLOW
+        - type: AssertEqual
+          config:
+            elements: []
+

sdks/python/apache_beam/yaml/yaml_io.py

@@ -723,3 +723,60 @@ def write_to_tfrecord(
           num_shards=num_shards,
           shard_name_template=shard_name_template,
           compression_type=getattr(CompressionTypes, compression_type))
+
+
+@beam.ptransform_fn
+def match_all(
+    pcoll,
+    *,
+    file_pattern: Optional[str] = None,
+    empty_match_treatment: str = 'ALLOW',
+):
+  """Matches file patterns from the input PCollection.
+
+  This transform returns a PCollection of matching files, each represented as a
+  Row with path, size_in_bytes, and last_updated_in_seconds fields.
+
+  Args:
+    file_pattern (str): The name of the field in the input PCollection that contains
+      the file pattern string. If not specified and the input PCollection has
+      exactly one field, that field will be used.
+    empty_match_treatment (str): How to treat empty matches. Possible values are
+      'ALLOW', 'DISALLOW', and 'ALLOW_IF_WILDCARD'. Defaults to 'ALLOW'.
+  """
+  from apache_beam.typehints import schemas
+
+  try:
+    field_names = [
+        name for name, _ in schemas.named_fields_from_element_type(
+            pcoll.element_type)
+    ]
+  except Exception:
+    field_names = None
+
+  if field_names:
+    if file_pattern is not None:
+      if file_pattern not in field_names:
+        raise ValueError(
+            f"Field '{file_pattern}' not found in input schema fields: {field_names}"
+        )
+      pattern_field = file_pattern
+    elif len(field_names) == 1:
+      pattern_field = field_names[0]
+    else:
+      raise ValueError(
+          f"Input schema has multiple fields {field_names}. "
+          f"Please specify the 'file_pattern' parameter to select which field "
+          f"contains the file pattern.")
+    patterns = pcoll | beam.Map(lambda x: str(getattr(x, pattern_field)))
+  else:
+    patterns = pcoll
+
+  matched = patterns | beam.io.fileio.MatchAll(
+      empty_match_treatment=empty_match_treatment)
+
+  return matched | beam.Map(
+      lambda x: beam.Row(
+          path=str(x.path), size_in_bytes=int(x.size_in_bytes),
+          last_updated_in_seconds=float(x.last_updated_in_seconds)
+          if x.last_updated_in_seconds is not None else None))

sdks/python/apache_beam/yaml/yaml_io_test.py

@@ -18,6 +18,8 @@
 import io
 import json
 import logging
+import os
+import tempfile
 import unittest
 
 import fastavro
@@ -543,6 +545,122 @@ def test_read_proto(self):
         assert_that(result, equal_to(data))
 
 
+class YamlMatchAllTest(unittest.TestCase):
+  def test_match_all_simple(self):
+    with tempfile.TemporaryDirectory() as temp_dir:
+      file1 = os.path.join(temp_dir, 'file1.txt')
+      file2 = os.path.join(temp_dir, 'file2.txt')
+      for f in [file1, file2]:
+        with open(f, 'w') as fout:
+          fout.write('data')
+
+      with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+          pickle_library='cloudpickle')) as p:
+        result = (
+            p
+            | beam.Create(
+                [beam.Row(pattern=os.path.join(temp_dir, 'file*.txt'))])
+            | YamlTransform(
+                '''
+                type: MatchAll
+                config:
+                  file_pattern: pattern
+                '''))
+        paths = result | beam.Map(lambda row: row.path)
+        assert_that(paths, equal_to([file1, file2]))
+
+  def test_match_all_single_field_default(self):
+    with tempfile.TemporaryDirectory() as temp_dir:
+      file1 = os.path.join(temp_dir, 'file1.txt')
+      with open(file1, 'w') as fout:
+        fout.write('data')
+
+      with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+          pickle_library='cloudpickle')) as p:
+        result = (
+            p
+            | beam.Create([beam.Row(my_sole_pattern=file1)])
+            | YamlTransform(
+                '''
+                type: MatchAll
+                '''))
+        paths = result | beam.Map(lambda row: row.path)
+        assert_that(paths, equal_to([file1]))
+
+  def test_match_all_multiple_fields_error(self):
+    with self.assertRaises(Exception):
+      with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+          pickle_library='cloudpickle')) as p:
+        _ = (
+            p
+            | beam.Create([beam.Row(pattern='foo', other_field='bar')])
+            | YamlTransform(
+                '''
+                type: MatchAll
+                '''))
+
+  def test_match_all_empty_match_disallow_error(self):
+    with self.assertRaises(Exception):
+      with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+          pickle_library='cloudpickle')) as p:
+        _ = (
+            p
+            | beam.Create([beam.Row(pattern='does_not_exist*.txt')])
+            | YamlTransform(
+                '''
+                type: MatchAll
+                config:
+                  empty_match_treatment: DISALLOW
+                '''))
+
+  def test_match_all_invalid_field_error(self):
+    with self.assertRaisesRegex(
+        ValueError, "Field 'invalid_field' not found in input schema fields"):
+      with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+          pickle_library='cloudpickle')) as p:
+        _ = (
+            p
+            | beam.Create([beam.Row(pattern='foo')])
+            | YamlTransform(
+                '''
+                type: MatchAll
+                config:
+                  file_pattern: invalid_field
+                '''))
+
+  def test_match_all_none_timestamp(self):
+    from apache_beam.io.filesystem import FileMetadata
+
+    class MockMatchAll(beam.PTransform):
+      def expand(self, pcoll):
+        return pcoll.pipeline | beam.Create([
+            FileMetadata(
+                path='file.txt',
+                size_in_bytes=100,
+                last_updated_in_seconds=None)
+        ])
+
+    with mock.patch('apache_beam.io.fileio.MatchAll',
+                    return_value=MockMatchAll()):
+      with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+          pickle_library='cloudpickle')) as p:
+        result = (
+            p
+            | beam.Create([beam.Row(pattern='file.txt')])
+            | YamlTransform(
+                '''
+                type: MatchAll
+                '''))
+        assert_that(
+            result,
+            equal_to([
+                beam.Row(
+                    path='file.txt',
+                    size_in_bytes=100,
+                    last_updated_in_seconds=None)
+            ]))
+
+
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)
   unittest.main()