from __future__ import annotations
import tempfile
from typing import Any
from fsspec import AbstractFileSystem, filesystem
from fsspec.core import url_to_fs
from fsspec.implementations.chained import ChainedFileSystem
from .interchange import DEFAULT_REGISTRY, DataFormat, InterchangeRequest, SchemaPolicy
from .schema import SchemaInput, SchemaProvenance, normalize_schema, resolve_schema
[docs]
class DataFileSystem(ChainedFileSystem):
"""Read-only format and schema conversion layered over another filesystem."""
protocol = "fsspec-data"
def __init__(
self,
fo: str,
target_protocol: str | None = None,
target_options: dict[str, Any] | None = None,
fs: AbstractFileSystem | None = None,
provided_format: DataFormat | str | None = None,
requested_format: DataFormat | str | None = None,
provided_schema: SchemaInput = None,
requested_schema: SchemaInput = None,
schema_policy: SchemaPolicy | str = SchemaPolicy.EXACT,
batch_size: int = 1024,
row_limit: int | None = None,
byte_limit: int | None = None,
spool_max_size: int = 8 * 1024 * 1024,
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
if fs is not None and target_protocol is not None:
raise ValueError("provide either fs or target_protocol, not both")
if fs is None:
if target_protocol is None:
fs, fo = url_to_fs(fo, **(target_options or {}))
else:
fs = filesystem(target_protocol, **(target_options or {}))
self.fs = fs
self.fo = fs._strip_protocol(fo)
self.provided_format = DataFormat(provided_format) if provided_format is not None else None
self.requested_format = DataFormat(requested_format) if requested_format is not None else None
self.provided_schema = normalize_schema(provided_schema)
self.requested_schema = normalize_schema(requested_schema)
self.provided_schema_provenance: SchemaProvenance | None = None
self.requested_schema_provenance: SchemaProvenance | None = None
self.schema_policy = SchemaPolicy(schema_policy)
self.batch_size = batch_size
self.row_limit = row_limit
self.byte_limit = byte_limit
self.spool_max_size = spool_max_size
self._sizes: dict[str, int] = {}
[docs]
def _open(
self,
path: str,
mode: str = "rb",
block_size: int | None = None,
autocommit: bool = True,
cache_options: dict[str, Any] | None = None,
**kwargs: Any,
):
del block_size, autocommit, cache_options, kwargs
if mode != "rb":
raise ValueError("fsspec-data is a read-only filesystem")
path = self._strip_protocol(path)
file = tempfile.SpooledTemporaryFile(max_size=self.spool_max_size, mode="w+b") # noqa: SIM115
try:
self._convert(path, file)
self._sizes[path] = file.tell()
file.seek(0)
except Exception:
file.close()
raise
return file
[docs]
def info(self, path: str, **kwargs: Any) -> dict[str, Any]:
del kwargs
path = self._strip_protocol(path)
size = self._sizes.get(path)
if size is None:
with self._open(path) as file:
file.seek(0, 2)
size = file.tell()
return {"name": path, "size": size, "type": "file"}
[docs]
def ls(self, path: str, detail: bool = True, **kwargs: Any):
path = self._strip_protocol(path)
if not path:
return []
entry = self.info(path, **kwargs)
return [entry] if detail else [entry["name"]]
[docs]
def _convert(self, path: str, output) -> None:
provided_format = self.provided_format or _format_from_path(self.fo)
requested_format = self.requested_format or _format_from_path(path)
resolved_provided = resolve_schema(self.provided_schema)
resolved_requested = resolve_schema(self.requested_schema)
provided_schema = resolved_provided.schema if resolved_provided is not None else None
requested_schema = resolved_requested.schema if resolved_requested is not None else None
self.provided_schema_provenance = resolved_provided.provenance if resolved_provided is not None else None
self.requested_schema_provenance = resolved_requested.provenance if resolved_requested is not None else None
with self.fs.open(self.fo, "rb") as source:
decoded = DEFAULT_REGISTRY.get(provided_format).iter_batches(
source,
schema=provided_schema,
batch_size=self.batch_size,
row_limit=self.row_limit,
byte_limit=self.byte_limit,
)
if provided_schema is not None and not decoded.schema.equals(provided_schema, check_metadata=False):
raise ValueError("provided schema does not match the source schema")
provided_schema = provided_schema or decoded.schema
requested_schema = requested_schema or provided_schema
plan = InterchangeRequest(
provided_format,
requested_format,
provided_schema,
requested_schema,
self.schema_policy,
).plan()
batches = (plan.apply_batch(batch) for batch in decoded)
DEFAULT_REGISTRY.get(requested_format).encode_batches_to(batches, output, schema=requested_schema)
_SUFFIX_FORMATS = {
".arrow": DataFormat.ARROW,
".csv": DataFormat.CSV,
".ipc": DataFormat.ARROW,
".jsonl": DataFormat.JSONL,
".ndjson": DataFormat.JSONL,
".parquet": DataFormat.PARQUET,
".pq": DataFormat.PARQUET,
}