feat: add param read_nums

Author: ChenZiHong-Gavin

graphgen/operators/read/read.py

@@ -53,6 +53,7 @@ def read(
     working_dir: Optional[str] = "cache",
     parallelism: int = 4,
     recursive: bool = True,
+    read_num: Optional[int] = None,
     **reader_kwargs: Any,
 ) -> ray.data.Dataset:
     """
@@ -120,6 +121,9 @@ def read(
             }
         )
 
+        if read_num is not None:
+            combined_ds = combined_ds.limit(read_num)
+
         logger.info("[READ] Successfully read files from %s", input_path)
         return combined_ds