wehub-resource-sync
43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
from typing import TYPE_CHECKING, Iterator, List
|
|
|
|
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
|
|
from ray.data.block import Block
|
|
from ray.data.datasource.file_based_datasource import FileBasedDatasource
|
|
|
|
if TYPE_CHECKING:
|
|
import pyarrow
|
|
|
|
|
|
class TextDatasource(FileBasedDatasource):
|
|
"""Text datasource, for reading and writing text files."""
|
|
|
|
_COLUMN_NAME = "text"
|
|
|
|
def __init__(
|
|
self,
|
|
paths: List[str],
|
|
*,
|
|
drop_empty_lines: bool = False,
|
|
encoding: str = "utf-8",
|
|
**file_based_datasource_kwargs
|
|
):
|
|
super().__init__(paths, **file_based_datasource_kwargs)
|
|
|
|
self.drop_empty_lines = drop_empty_lines
|
|
self.encoding = encoding
|
|
|
|
def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
|
|
data = f.readall()
|
|
|
|
builder = DelegatingBlockBuilder()
|
|
|
|
lines = data.decode(self.encoding).splitlines()
|
|
for line in lines:
|
|
if self.drop_empty_lines and line.strip() == "":
|
|
continue
|
|
item = {self._COLUMN_NAME: line}
|
|
builder.add(item)
|
|
|
|
block = builder.build()
|
|
yield block
|