wehub-resource-sync
83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union
|
|
|
|
from ray.data.block import Block
|
|
from ray.data.datasource.file_based_datasource import FileBasedDatasource
|
|
|
|
if TYPE_CHECKING:
|
|
import pyarrow
|
|
|
|
|
|
class CSVDatasource(FileBasedDatasource):
|
|
"""CSV datasource, for reading and writing CSV files."""
|
|
|
|
_FILE_EXTENSIONS = [
|
|
"csv",
|
|
"csv.gz", # gzip-compressed files
|
|
"csv.br", # Brotli-compressed files
|
|
"csv.zst", # Zstandard-compressed files
|
|
"csv.lz4", # lz4-compressed files
|
|
]
|
|
|
|
def __init__(
|
|
self,
|
|
paths: Union[str, List[str]],
|
|
arrow_csv_args: Optional[Dict[str, Any]] = None,
|
|
**file_based_datasource_kwargs,
|
|
):
|
|
from pyarrow import csv
|
|
|
|
super().__init__(paths, **file_based_datasource_kwargs)
|
|
|
|
if arrow_csv_args is None:
|
|
arrow_csv_args = {}
|
|
|
|
self.read_options = arrow_csv_args.pop(
|
|
"read_options", csv.ReadOptions(use_threads=False)
|
|
)
|
|
self.parse_options = arrow_csv_args.pop("parse_options", csv.ParseOptions())
|
|
self.arrow_csv_args = arrow_csv_args
|
|
|
|
def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
|
|
import pyarrow as pa
|
|
from pyarrow import csv
|
|
|
|
# Re-init invalid row handler: https://issues.apache.org/jira/browse/ARROW-17641
|
|
if hasattr(self.parse_options, "invalid_row_handler"):
|
|
self.parse_options.invalid_row_handler = (
|
|
self.parse_options.invalid_row_handler
|
|
)
|
|
|
|
filter_expr = (
|
|
self._predicate_expr.to_pyarrow()
|
|
if self._predicate_expr is not None
|
|
else None
|
|
)
|
|
|
|
try:
|
|
reader = csv.open_csv(
|
|
f,
|
|
read_options=self.read_options,
|
|
parse_options=self.parse_options,
|
|
**self.arrow_csv_args,
|
|
)
|
|
schema = None
|
|
while True:
|
|
try:
|
|
batch = reader.read_next_batch()
|
|
table = pa.Table.from_batches([batch], schema=schema)
|
|
if schema is None:
|
|
schema = table.schema
|
|
if filter_expr is not None:
|
|
table = table.filter(filter_expr)
|
|
|
|
yield table
|
|
except StopIteration:
|
|
return
|
|
except pa.lib.ArrowInvalid as e:
|
|
raise ValueError(
|
|
f"Failed to read CSV file: {path}. "
|
|
"Please check the CSV file has correct format, or filter out non-CSV "
|
|
"file with 'partition_filter' field. See read_csv() documentation for "
|
|
"more details."
|
|
) from e
|