refactor: refactor csv_reader using ray data

Author: ChenZiHong-Gavin

graphgen/bases/base_reader.py

@@ -2,6 +2,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Union
 
+import pandas as pd
 import requests
 from ray.data import Dataset
 
@@ -43,6 +44,21 @@ def _should_keep_item(self, item: Dict[str, Any]) -> bool:
             return bool(content)
         return True
 
+    def _validate_batch(self, batch: pd.DataFrame) -> pd.DataFrame:
+        """
+        Validate data format.
+        """
+        if "type" not in batch.columns:
+            raise ValueError(f"Missing 'type' column. Found: {list(batch.columns)}")
+
+        if "text" in batch["type"].values:
+            if self.text_column not in batch.columns:
+                raise ValueError(
+                    f"Missing '{self.text_column}' column for text documents"
+                )
+
+        return batch
+
     @staticmethod
     def _image_exists(path_or_url: str, timeout: int = 3) -> bool:
         """

graphgen/models/reader/csv_reader.py

@@ -1,6 +1,7 @@
-from typing import Any, Dict, List
+from typing import List, Union
 
-import pandas as pd
+import ray
+from ray.data import Dataset
 
 from graphgen.bases.base_reader import BaseReader
 
@@ -13,13 +14,23 @@ class CSVReader(BaseReader):
         - if type is "text", "content" column must be present.
     """
 
-    def read(self, file_path: str) -> List[Dict[str, Any]]:
+    def read(
+        self,
+        input_path: Union[str, List[str]],
+        override_num_blocks: int = None,
+    ) -> Dataset:
+        """
+        Read CSV files and return Ray Dataset.
 
-        df = pd.read_csv(file_path)
-        for _, row in df.iterrows():
-            assert "type" in row, f"Missing 'type' column in document: {row.to_dict()}"
-            if row["type"] == "text" and self.text_column not in row:
-                raise ValueError(
-                    f"Missing '{self.text_column}' in document: {row.to_dict()}"
-                )
-        return self.filter(df.to_dict(orient="records"))
+        :param input_path: Path to CSV file or list of CSV files.
+        :param override_num_blocks: Number of blocks for Ray Dataset reading.
+        :return: Ray Dataset containing validated and filtered data.
+        """
+
+        ds = ray.data.read_csv(input_path, override_num_blocks=override_num_blocks)
+
+        ds = ds.map_batches(self._validate_batch, batch_format="pandas")
+
+        ds = ds.filter(self._should_keep_item)
+
+        return ds

requirements.txt

@@ -20,6 +20,8 @@ requests
 fastapi
 trafilatura
 aiohttp
+ray
+diskcache
 
 leidenalg
 igraph