[data] fix the test, test_chain

Signed-off-by: Xingyu Long <xingyulong97@gmail.com>

Author: xingyu-long

python/ray/data/preprocessors/chain.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from ray.air.util.data_batch_conversion import BatchFormat
 from ray.data import Dataset
@@ -79,9 +79,22 @@ def fit_transform(self, ds: Dataset) -> Dataset:
             ds = preprocessor.fit_transform(ds)
         return ds
 
-    def _transform(self, ds: Dataset) -> Dataset:
+    def _transform(
+        self,
+        ds: Dataset,
+        batch_size: Optional[int],
+        num_cpus: Optional[float] = None,
+        memory: Optional[float] = None,
+        concurrency: Optional[int] = None,
+    ) -> Dataset:
         for preprocessor in self.preprocessors:
-            ds = preprocessor.transform(ds)
+            ds = preprocessor.transform(
+                ds,
+                batch_size=batch_size,
+                num_cpus=num_cpus,
+                memory=memory,
+                concurrency=concurrency,
+            )
         return ds
 
     def _transform_batch(self, df: "DataBatchType") -> "DataBatchType":