Internal

PiperOrigin-RevId: 832363601

Author: Grain Team

grain/_src/python/BUILD

@@ -335,6 +335,7 @@ py_library(
     srcs = ["shared_memory_array.py"],
     srcs_version = "PY3",
     deps = [
+        "//grain/_src/core:tree_lib",
         "@pypi//numpy:pkg",
     ],
 )

grain/_src/python/dataset/BUILD

@@ -56,6 +56,7 @@ py_library(
         "//grain/_src/python:options",
         "//grain/_src/python:shared_memory_array",
         "//grain/proto:execution_summary_py_pb2",
+        "@abseil-py//absl/flags",
         "@abseil-py//absl/logging",
         "@pypi//cloudpickle:pkg",
         "@pypi//etils:pkg",

grain/_src/python/dataset/transformations/BUILD

@@ -14,6 +14,7 @@ filegroup(
         "map.py",
         "mix.py",
         "prefetch.py",
+        "process_prefetch.py",
         "rebatch.py",
         "repeat.py",
         "shuffle.py",
@@ -352,3 +353,21 @@ py_test(
         "@pypi//numpy:pkg",
     ],
 )
+
+py_test(
+    name = "process_prefetch_test",
+    timeout = "long",
+    srcs = ["process_prefetch_test.py"],
+    shard_count = 50,
+    srcs_version = "PY3",
+    deps = [
+        "//grain/_src/core:transforms",
+        "//grain/_src/python:options",
+        "//grain/_src/python/dataset",
+        "//grain/_src/python/dataset:base",
+        "@abseil-py//absl/logging",
+        "@abseil-py//absl/testing:absltest",
+        "@abseil-py//absl/testing:parameterized",
+        "@pypi//numpy:pkg",
+    ],
+)