Fix package containing RestrictionProvider

apache · Aug 1, 2021 · 68ecff9 · 68ecff9
1 parent dce846b
commit 68ecff9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sdks/python/apache_beam/examples/snippets/snippets.py b/sdks/python/apache_beam/examples/snippets/snippets.py
@@ -1647,7 +1647,7 @@ def sdf_basic_example():
   read_next_record = None
 
   # [START SDF_BasicExample]
-  class FileToWordsRestrictionProvider(beam.io.RestrictionProvider):
+  class FileToWordsRestrictionProvider(beam.transforms.core.RestrictionProvider):
     def initial_restriction(self, file_name):
       return OffsetRange(0, os.stat(file_name).st_size)
 
@@ -1680,7 +1680,7 @@ def sdf_basic_example_with_splitting():
   from apache_beam.io.restriction_trackers import OffsetRange
 
   # [START SDF_BasicExampleWithSplitting]
-  class FileToWordsRestrictionProvider(beam.io.RestrictionProvider):
+  class FileToWordsRestrictionProvider(beam.transforms.core.RestrictionProvider):
     def split(self, file_name, restriction):
       # Compute and output 64 MiB size ranges to process in parallel
       split_size = 64 * (1 << 20)