Implement format handlers for Parquet, JSON, and CSV; remove obsolete fsspec utility functions and metadata helpers; refactor schema conversion logic; enhance Writer class with retry logic for dataset writing; introduce utility functions for disk usage tracking and file name management; update dependencies in pyproject.toml.

Author: legout

pydala/cache.py

@@ -0,0 +1,85 @@
+import posixpath
+import os
+from typing import Any
+
+from fsspec.implementations.cache_mapper import AbstractCacheMapper
+from fsspec.implementations.cached import SimpleCacheFileSystem
+from loguru import logger
+
+from .helpers.security import safe_join, validate_path
+
+
+class FileNameCacheMapper(AbstractCacheMapper):
+    def __init__(self, directory: str):
+        self.directory = validate_path(directory)
+
+    def __call__(self, path: str) -> str:
+        validated_path = validate_path(path)
+        full_path = safe_join(self.directory, validated_path)
+        parent_dir = posixpath.dirname(full_path)
+        os.makedirs(parent_dir, exist_ok=True)
+        return validated_path
+
+
+class MonitoredSimpleCacheFileSystem(SimpleCacheFileSystem):
+    def __init__(self, verbose: bool = False, **kwargs):
+        self._verbose = verbose
+        super().__init__(**kwargs)
+        self._mapper = FileNameCacheMapper(kwargs.get("cache_storage", "~/.tmp"))
+
+    def _check_file(self, path: str):
+        self._check_cache()
+        cache_path = self._mapper(path)
+        for storage in self.storage:
+            fn = posixpath.join(storage, cache_path)
+            if posixpath.exists(fn):
+                return fn
+            if self._verbose:
+                logger.info(f"Downloading {self.protocol[0]}://{path}")
+
+    def size(self, path: str):
+        cached_file = self._check_file(self._strip_protocol(path))
+        if cached_file is None:
+            return self.fs.size(path)
+        return posixpath.getsize(cached_file)
+
+    def __getattribute__(self, item: str) -> Any:
+        if item in self._delegated_methods:
+            return lambda *args, **kwargs: getattr(type(self), item).__get__(self, type(self))(*args, **kwargs)
+        if item in {"__reduce_ex__", "__reduce__"}:
+            raise AttributeError(item)
+        if item == "transaction":
+            return type(self).transaction.__get__(self, type(self))
+        if item in {"_cache", "transaction_type"}:
+            return getattr(type(self), item)
+        if item == "__class__":
+            return type(self)
+        return self._delegate_to_fs(item)
+
+    def _delegate_to_fs(self, item: str) -> Any:
+        d = object.__getattribute__(self, "__dict__")
+        fs = d.get("fs")
+        if item in d:
+            return d[item]
+        if fs is None:
+            return super().__getattribute__(item)
+        if item in fs.__dict__:
+            return fs.__dict__[item]
+        cls = type(fs)
+        m = getattr(cls, item, None)
+        if m is None:
+            raise AttributeError(f"'{item}' not found in underlying fs")
+        if callable(m) and not hasattr(m, "__self__") or m.__self__ is None:
+            return m.__get__(fs, cls)
+        return m
+
+    _delegated_methods = {
+        "size", "glob", "load_cache", "_open", "save_cache", "close_and_update",
+        "__init__", "__getattribute__", "__reduce__", "_make_local_details", "open",
+        "cat", "cat_file", "cat_ranges", "get", "read_block", "tail", "head", "info",
+        "ls", "exists", "isfile", "isdir", "_check_file", "_check_cache", "_mkcache",
+        "clear_cache", "clear_expired_cache", "pop_from_cache", "local_file",
+        "_paths_from_path", "get_mapper", "open_many", "commit_many", "hash_name",
+        "__hash__", "__eq__", "to_json", "to_dict", "cache_size", "pipe_file", "pipe",
+        "start_transaction", "end_transaction"
+    }

pydala/catalog.py

@@ -10,6 +10,75 @@
 from pydala.helpers.polars import pl
 
 from .dataset import CsvDataset, JsonDataset, ParquetDataset, PyarrowDataset
+
+from abc import ABC, abstractmethod
+
+class AbstractLoader(ABC):
+    def _matches_format(self, params) -> bool:
+        raise NotImplementedError
+
+    def _read_data(self, catalog, params, **kwargs) -> pl.DataFrame:
+        raise NotImplementedError
+
+    def _get_dataset_class(self):
+        raise NotImplementedError
+
+    def load(self, catalog, table_name, as_dataset: bool, with_metadata: bool = False, **kwargs):
+        params = catalog._get_table_params(table_name=table_name)
+        if not self._matches_format(params):
+            return None
+        if not as_dataset:
+            df = self._read_data(catalog, params, **kwargs)
+            catalog.ddb_con.register(table_name, df)
+            return df
+        cls = self._get_dataset_class(with_metadata)
+        return cls(params.path, filesystem=catalog.fs[params.filesystem], name=table_name, ddb_con=catalog.ddb_con, **kwargs)
+
+class ParquetLoader(AbstractLoader):
+    def _matches_format(self, params) -> bool:
+        return 'parquet' in params.format.lower()
+
+    def _read_data(self, catalog, params, **kwargs) -> pl.DataFrame:
+        fs = catalog.fs[params.filesystem]
+        if params.path.endswith('.parquet'):
+            return fs.read_parquet(params.path, **kwargs)
+        return fs.read_parquet_dataset(params.path, **kwargs)
+
+    def _get_dataset_class(self, with_metadata: bool = True):
+        return ParquetDataset if with_metadata else PyarrowDataset
+
+class CsvLoader(AbstractLoader):
+    def _matches_format(self, params) -> bool:
+        return 'csv' in params.format.lower()
+
+    def _read_data(self, catalog, params, **kwargs) -> pl.DataFrame:
+        fs = catalog.fs[params.filesystem]
+        if params.path.endswith('.csv'):
+            return fs.read_csv(params.path, **kwargs)
+        return fs.read_csv_dataset(params.path, **kwargs)
+
+    def _get_dataset_class(self, with_metadata: bool = True):
+        return CsvDataset
+
+class JsonLoader(AbstractLoader):
+    def _matches_format(self, params) -> bool:
+        return 'json' in params.format.lower()
+
+    def _read_data(self, catalog, params, **kwargs) -> pl.DataFrame:
+        fs = catalog.fs[params.filesystem]
+        if params.path.endswith('.json'):
+            return fs.read_json(params.path, **kwargs)
+        return fs.read_json_dataset(params.path, **kwargs)
+
+    def _get_dataset_class(self, with_metadata: bool = True):
+        return JsonDataset
+
+# Registry
+LOADERS = {
+    'parquet': ParquetLoader(),
+    'csv': CsvLoader(),
+    'json': JsonLoader(),
+}
 from .filesystem import FileSystem
 from .helpers.misc import delattr_rec, get_nested_keys, getattr_rec, setattr_rec
 from .helpers.sql import get_table_names
@@ -162,87 +231,22 @@ def files(self, table_name: str) -> list[str]:
         )
 
     def load_parquet(
-        self, table_name: str, as_dataset=True, with_metadata: bool = True, **kwargs
+        self, table_name: str, as_dataset: bool = True, with_metadata: bool = True, **kwargs
     ) -> ParquetDataset | PyarrowDataset | pl.DataFrame | None:
-        params = self._get_table_params(table_name=table_name)
-
-        if "parquet" not in params.format.lower():
-            return
-        if not as_dataset:
-            if params.path.endswith(".parquet"):
-                df = self.fs[params.filesystem].read_parquet(params.path, **kwargs)
-                self.ddb_con.register(table_name, df)
-                return df
-
-            df = self.fs[params.filesystem].read_parquet_dataset(params.path, **kwargs)
-            self.ddb_con.register(table_name, df)
-            return df
-
-        if with_metadata:
-            return ParquetDataset(
-                params.path,
-                filesystem=self.fs[params.filesystem],
-                name=table_name,
-                ddb_con=self.ddb_con,
-                **kwargs,
-            )
-
-        return PyarrowDataset(
-            params.path,
-            filesystem=self.fs[params.filesystem],
-            name=table_name,
-            ddb_con=self.ddb_con,
-            **kwargs,
-        )
+        """Load Parquet table as DataFrame or dataset."""
+        return self.load(table_name, as_dataset=as_dataset, with_metadata=with_metadata, **kwargs)
 
     def load_csv(
         self, table_name: str, as_dataset: bool = True, **kwargs
     ) -> CsvDataset | pl.DataFrame | None:
-        params = self._get_table_params(table_name=table_name)
-
-        if "csv" not in params.format.lower():
-            return
-        if not as_dataset:
-            if params.path.endswith(".csv"):
-                df = self.fs[params.filesystem].read_parquet(params.path, **kwargs)
-                self.ddb_con.register(table_name, df)
-                return df
-
-            df = self.fs[params.filesystem].read_parquet_dataset(params.path, **kwargs)
-            self.ddb_con.register(table_name, df)
-            return df
-
-        return CsvDataset(
-            params.path,
-            filesystem=self.fs[params.filesystem],
-            name=table_name,
-            ddb_con=self.ddb_con,
-            **kwargs,
-        )
+        """Load CSV table as DataFrame or dataset."""
+        return self.load(table_name, as_dataset=as_dataset, with_metadata=False, **kwargs)
 
     def load_json(
         self, table_name: str, as_dataset: bool = True, **kwargs
     ) -> JsonDataset | pl.DataFrame | None:
-        params = self._get_table_params(table_name=table_name)
-
-        if "json" not in params.format.lower():
-            return
-        if not as_dataset:
-            if params.path.endswith(".json"):
-                df = self.fs[params.filesystem].read_json(params.path, **kwargs)
-                self.ddb_con.register(table_name, df)
-                return df
-
-            df = self.fs[params.filesystem].read_json_dataset(params.path, **kwargs)
-            self.ddb_con.register(table_name, df)
-            return df
-        return JsonDataset(
-            params.path,
-            filesystem=self.fs[params.filesystem],
-            name=table_name,
-            ddb_con=self.ddb_con,
-            **kwargs,
-        )
+        """Load JSON table as DataFrame or dataset."""
+        return self.load(table_name, as_dataset=as_dataset, with_metadata=False, **kwargs)
 
     def load(
         self,
@@ -253,30 +257,15 @@ def load(
         **kwargs,
     ):
         params = self._get_table_params(table_name=table_name)
-
-        if params.format.lower() == "parquet":
-            if table_name not in self.table and not reload:
-                self.table[table_name] = self.load_parquet(
-                    table_name,
-                    as_dataset=as_dataset,
-                    with_metadata=with_metadata,
-                    **kwargs,
-                )
-            return self.table[table_name]
-
-        elif params.format.lower() == "csv":
-            if table_name not in self.table and not reload:
-                self.table[table_name] = self.load_csv(
-                    table_name, as_dataset=as_dataset, **kwargs
-                )
-            return self.table[table_name]
-
-        elif params.format.lower() == "json":
-            if table_name not in self.table and not reload:
-                self.table[table_name] = self.load_json(table_name, **kwargs)
-            return self.table[table_name]
-
-        # return None
+        format_lower = params.format.lower()
+        loader = LOADERS.get(format_lower)
+        if loader is None:
+            return None
+        if table_name not in self.table and not reload:
+            self.table[table_name] = loader.load(
+                self, table_name, as_dataset, with_metadata, **kwargs
+            )
+        return self.table[table_name]
 
     # def _ddb_table_mapping(self, table_name: str):
     #     params = getattr_rec(self._catalog, self._get_table_from_table_name(table_name=table_name))