File size: 799 Bytes
31086ae
d2a63cc
31086ae
 
d2a63cc
 
 
 
e4316f1
52419fe
 
 
 
 
 
 
31086ae
 
 
d2a63cc
31086ae
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
from typing import List, Union

import ray
from ray.data import Dataset

from graphgen.bases.base_reader import BaseReader


class CSVReader(BaseReader):
    """
    Reader for CSV files.
    Columns:
        - type: The type of the document (e.g., "text", "image", etc.)
        - if type is "text", "content" column must be present.
    """

    def read(self, input_path: Union[str, List[str]]) -> Dataset:
        """
        Read CSV files and return Ray Dataset.

        :param input_path: Path to CSV file or list of CSV files.
        :return: Ray Dataset containing validated and filtered data.
        """

        ds = ray.data.read_csv(input_path)
        ds = ds.map_batches(self._validate_batch, batch_format="pandas")
        ds = ds.filter(self._should_keep_item)
        return ds