nuthatch 0.1.0__py3-none-any.whl

nuthatch/__init__.py

@@ -0,0 +1,14 @@
+"""
+Nuthatch is a library for caching data based on the function call and its arguments.
+
+It caches data in a variety of backends optimized for different data types.
+"""
+from .config import config_parameter
+from .nuthatch import cache
+
+# Trigger backend registration
+import nuthatch.backends #noqa
+
+__version__ = "0.1.0"
+
+__all__ = ["config_parameter", "cache", "__version__", "cli"]

nuthatch/backend.py

@@ -0,0 +1,301 @@
+"""The backend module is used to register and get backends.
+
+It also contains the abstract base class for all backends.
+"""
+from abc import ABC, abstractmethod
+from os.path import join
+import fsspec
+import sqlalchemy
+
+registered_backends = {}
+default_backends = {}
+
+def register_backend(backendClass):
+    """Register a backend class with the nuthatch system.
+
+    This function registers a backend class so it can be used by the nuthatch
+    caching system. The backend class must have a `backend_name` attribute.
+    Optionally, if the backend class has a `default_for_type` attribute, it
+    will be set as the default backend for that data type.
+
+    Args:
+        backendClass: The backend class to register. Must have a `backend_name`
+                     attribute and optionally a `default_for_type` attribute.
+
+    Returns:
+        The registered backend class (allows use as a decorator).
+
+    Example:
+        @register_backend
+        class MyBackend:
+            backend_name = "my_backend"
+            default_for_type = "my_data_type"
+    """
+    registered_backends[backendClass.backend_name] = backendClass
+
+    if 'default_for_type' in backendClass.__dict__:
+        default_backends[backendClass.default_for_type] = backendClass.backend_name
+
+    return backendClass
+
+def get_backend_by_name(backend_name):
+    """Retrieve a registered backend class by its name.
+
+    Args:
+        backend_name (str): The name of the backend to retrieve.
+
+    Returns:
+        The backend class associated with the given name.
+
+    Raises:
+        KeyError: If no backend is registered with the given name.
+    """
+    return registered_backends[backend_name]
+
+def get_default_backend(data_type):
+    """Get the default backend name for a specific data type.
+
+    Args:
+        data_type (str): The data type to get the default backend for.
+
+    Returns:
+        str: The name of the default backend for the data type, or 'basic'
+             if no default is set.
+    """
+    if data_type in default_backends:
+        return default_backends[data_type]
+    else:
+        return 'basic'
+
+class NuthatchBackend(ABC):
+    config_parameters = []
+
+    def __init__(self, cacheable_config, cache_key, namespace, args, backend_kwargs):
+        """The abstract base class for all cacheable backends.
+
+        This is the base class for all cacheable backends in the nuthatch system.
+        Each backend instance manages a specific cache entry with its own
+        configuration and storage mechanism.
+
+        Args:
+            cacheable_config (dict): Configuration dictionary containing static or
+                dynamic parameters. Required parameters should be listed as strings
+                in the backend's `config_parameters` class attribute.
+            cache_key (str): Unique identifier for the cache entry. This key is
+                used to distinguish between different cached items.
+            namespace (str, optional): Optional namespace to organize cache entries.
+                If provided, cache entries will be stored under this namespace.
+            args (dict): Key-value pairs of arguments and values that were passed
+                to the function being cached. These are used to generate cache
+                keys and may influence backend behavior.
+            backend_kwargs (dict): User-configurable keyword arguments specific to
+                the backend implementation. For example, the zarr backend uses
+                these for per-argument-value chunking configuration.
+
+        Note:
+            This is an abstract base class. Subclasses must implement the abstract
+            methods: `write()`, `read()`, `delete()`, `exists()`, `get_file_path()`,
+            and `sync()`.
+        """
+
+        # Store base
+        self.config = cacheable_config
+        self.cache_key = cache_key
+        self.namespace = namespace
+        self.backend_kwargs = backend_kwargs
+        self.args = args
+
+
+    @abstractmethod
+    def write(self, data, upsert=False, primary_keys=None):
+        """Write data to the backend.
+
+        This method is responsible for writing data to the backend. It should
+        be implemented by subclasses to handle the specific storage mechanism
+        of the backend.
+
+        Args:
+            data (any): The data to write to the backend.
+            upsert (bool, optional): Whether to perform an upsert operation.
+                If True, the data will be inserted or updated based on the
+                primary keys provided.
+            primary_keys (list, optional): List of primary key columns to use
+                for upsert operations.
+        """
+        pass
+
+    @abstractmethod
+    def read(self, engine):
+        """Read data from the backend.
+
+        This method is responsible for reading data from the backend. It should
+        be implemented by subclasses to handle the specific storage mechanism
+        of the backend.
+
+        Args:
+            engine (str or type): The data processing engine to use for
+                reading data from the backend.
+
+        Returns:
+            Any: The data read from the backend.
+        """
+        pass
+
+    @abstractmethod
+    def delete(self):
+        """Delete data from the backend.
+
+        This method is responsible for deleting data from the backend. It should
+        be implemented by subclasses to handle the specific storage mechanism
+        of the backend.
+        """
+        pass
+
+    @abstractmethod
+    def exists(self):
+        """Check if the data exists in the backend.
+
+        This method is responsible for checking if the data exists in the backend.
+        It should be implemented by subclasses to handle the specific storage
+        mechanism of the backend.
+
+        Returns:
+            bool: True if the data exists in the backend, False otherwise.
+        """
+        pass
+
+    @abstractmethod
+    def get_file_path(self):
+        """Get the file path of the data in the backend.
+
+        This method is responsible for returning the file path of the data in the
+        backend. It should be implemented by subclasses to handle the specific
+        storage mechanism of the backend.
+
+        Returns:
+            str: The file path of the data in the backend.
+        """
+        pass
+
+    @abstractmethod
+    def sync(self, from_backend):
+        """Sync data from one backend to another.
+
+        This method is responsible for syncing data from one backend to another.
+        It should be implemented by subclasses to handle the specific storage
+        mechanism of the backend.
+
+        Args:
+            from_backend (NuthatchBackend): The backend to sync data from.
+        """
+        pass
+
+
+class FileBackend(NuthatchBackend):
+    """Base class for all backends that rely on a filesystem."""
+
+    config_parameters = ["filesystem", "filesystem_options"]
+
+    def __init__(self, cacheable_config, cache_key, namespace, args, backend_kwargs, extension):
+        super().__init__(cacheable_config, cache_key, namespace, args, backend_kwargs)
+
+        self.base_path = self.config['filesystem']
+
+        if namespace:
+            self.raw_cache_path = join(self.base_path, namespace, cache_key)
+        else:
+            self.raw_cache_path = join(self.base_path, cache_key)
+
+        self.temp_cache_path = join(self.base_path, 'temp', cache_key)
+        self.extension = extension
+        self.path = self.raw_cache_path + '.' + extension
+        if 'filesystem_options' not in self.config:
+            self.config['filesystem_options'] = {}
+
+        if fsspec.utils.get_protocol(self.path) == 'file':
+            # If the protocol is a local filesystem, we need to create the directory if it doesn't exist
+            self.fs = fsspec.core.url_to_fs(self.path, auto_mkdir=True, **self.config['filesystem_options'])[0]
+        else:
+            self.fs = fsspec.core.url_to_fs(self.path, **self.config['filesystem_options'])[0]
+
+
+    def exists(self):
+        return (self.fs.exists(self.path))
+
+    def delete(self):
+        self.fs.rm(self.path, recursive=True)
+
+    def get_file_path(self):
+        return self.path
+
+    def get_cache_key(self, path):
+        if path.endswith('.' + self.extension):
+            path = path[:-len('.' + self.extension)]
+
+        if path.startswith(self.base_path):
+            path = path[len(self.base_path):]
+
+        stripped = fsspec.core.strip_protocol(self.base_path)
+        if path.startswith(stripped):
+            path = path[len(stripped):]
+
+        if path.startswith('/'):
+            path = path[1:]
+
+        if self.namespace:
+            if path.startswith(self.namespace):
+                path = path[len(self.namespace):]
+
+        if path.startswith('/'):
+            path = path[1:]
+
+        return path
+
+
+    def sync(self, from_backend):
+        from_backend.fs.get(from_backend.path, self.path)
+
+@register_backend
+class NullBackend(FileBackend):
+
+    backend_name = 'null'
+
+    def __init__(self, cacheable_config, cache_key, namespace, args, backend_kwargs):
+        super().__init__(cacheable_config, cache_key, namespace, args, backend_kwargs, 'null')
+
+    def read(self, engine=None):
+        raise RuntimeError("Null backend used only for import/export path manipulation. Cannot read from null backend.")
+
+    def write(self, engine=None):
+        raise RuntimeError("Null backend used only for import/export path manipulation. Cannot write to null backend.")
+
+
+class DatabaseBackend(NuthatchBackend):
+    """Base class for all backends that rely on a database."""
+
+    config_parameters = ["driver", "host", "port", "database", "username", "password", "write_username", "write_password"]
+
+    def __init__(self, cacheable_config, cache_key, namespace, args, backend_kwargs):
+        super().__init__(cacheable_config, cache_key, namespace, args, backend_kwargs)
+
+        database_url = sqlalchemy.URL.create(self.config['driver'],
+                                username = self.config['username'],
+                                password = self.config['password'],
+                                host = self.config['host'],
+                                port = self.config['port'],
+                                database = self.config['database'])
+        self.engine = sqlalchemy.create_engine(database_url)
+
+        if 'write_username' in self.config and 'write_password' in self.config:
+            write_database_url = sqlalchemy.URL.create(self.config['driver'],
+                                username = self.config['write_username'],
+                                password = self.config['write_password'],
+                                host = self.config['host'],
+                                port = self.config['port'],
+                                database = self.config['database'])
+            self.write_engine = sqlalchemy.create_engine(write_database_url)
+        else:
+            self.write_engine = self.engine
+
+    def sync(self, local_backend):
+        raise NotImplementedError("Backend syncing not implemented for database-like backends.")

nuthatch/backends/__init__.py

@@ -0,0 +1,8 @@
+from .sql import SQLBackend
+from .delta import DeltaBackend
+from .basic import BasicBackend
+from .zarr import ZarrBackend
+from .terracotta import TerracottaBackend
+from .parquet import ParquetBackend
+
+__all__ = ["SQLBackend", "DeltaBackend", "BasicBackend", "ZarrBackend", "TerracottaBackend", "ParquetBackend"]

nuthatch/backends/basic.py

@@ -0,0 +1,28 @@
+from nuthatch.backend import FileBackend, register_backend
+import pickle
+
+@register_backend
+class BasicBackend(FileBackend):
+    """
+    Basic backend for caching data in a pickle file.
+    """
+
+    backend_name = "basic"
+
+    def __init__(self, cacheable_config, cache_key, namespace, args, backend_kwargs):
+        super().__init__(cacheable_config, cache_key, namespace, args, backend_kwargs, 'pkl')
+
+
+    def write(self, data, upsert=False, primary_keys=None):
+        if upsert:
+            raise ValueError("Basic/pickle backend does not support upsert.")
+
+        with self.fs.open(self.path, 'wb') as f:
+            pickle.dump(data, f)
+
+
+    def read(self, engine=None):
+        # Check to make sure the verify exists
+        if self.fs.exists(self.path):
+            with self.fs.open(self.path, 'rb') as f:
+                return pickle.load(f)

nuthatch/backends/delta.py

@@ -0,0 +1,46 @@
+from nuthatch.backend import FileBackend, register_backend
+from deltalake import DeltaTable, write_deltalake
+import dask.dataframe as dd
+import pandas as pd
+import dask_deltatable as ddt
+
+@register_backend
+class DeltaBackend(FileBackend):
+    """
+    Delta backend for caching tabular data in a delta table.
+
+    This backend supports dask and pandas dataframes.
+    """
+
+    backend_name = "delta"
+    default_for_type = pd.DataFrame
+
+    def __init__(self, cacheable_config, cache_key, namespace, args, backend_kwargs):
+        super().__init__(cacheable_config, cache_key, namespace, args, backend_kwargs, 'delta')
+
+
+    def write(self, data, upsert=False, primary_keys=None):
+        """Write a pandas dataframe to a delta table."""
+        if upsert:
+            raise ValueError("Delta backend does not support upsert.")
+
+        if isinstance(data, dd.DataFrame):
+            print("""Warning: Dask datafame passed to delta backend. Will run `compute()`
+                      on the dataframe prior to storage. This will fail if the dataframe
+                      does not fit in memory. Use `backend=parquet` to handle parallel writing of dask dataframes.""")
+            write_data = data.compute()
+        elif isinstance(data, pd.DataFrame):
+            write_data = data
+        else:
+            raise RuntimeError("Delta backend only supports dask and pandas engines.")
+
+        write_deltalake(self.path, write_data, mode='overwrite', schema_mode='overwrite')
+
+
+    def read(self, engine=None):
+        if engine == 'pandas' or engine == pd.DataFrame or engine is None:
+            return DeltaTable(self.path).to_pandas()
+        elif engine == 'dask' or engine == dd.DataFrame:
+            return ddt.read_deltalake(self.path)
+        else:
+            raise RuntimeError("Delta backend only supports dask and pandas engines.")

nuthatch/backends/parquet.py

@@ -0,0 +1,130 @@
+from nuthatch.backend import FileBackend, register_backend
+from pandas.api.types import is_datetime64_any_dtype as is_datetime
+import dask.dataframe as dd
+import pandas as pd
+
+
+def write_parquet_helper(df, path, partition_on=None):
+    """Helper to write parquets."""
+    print(path)
+    df.to_parquet(
+        path,
+        overwrite=True,
+        partition_on=partition_on,
+        engine="pyarrow",
+        write_metadata_file=True,
+        write_index=False,
+    )
+
+def read_from_parquet(cache_path):
+    """Read from a deltatable into a pandas dataframe."""
+    return dd.read_parquet(cache_path, engine='pyarrow', ignore_metadata_file=True)
+
+
+
+@register_backend
+class ParquetBackend(FileBackend):
+    """
+    Parquet backend for caching tabular data in a parquet file.
+
+    This backend supports dask and pandas dataframes.
+    """
+
+    backend_name = 'parquet'
+    default_for_type = dd.DataFrame
+
+    def __init__(self, cacheable_config, cache_key, namespace, args, backend_kwargs):
+        super().__init__(cacheable_config, cache_key, namespace, args, backend_kwargs, 'parquet')
+
+    def write(self, data, upsert=False, primary_keys=None):
+        if isinstance(data, dd.DataFrame):
+            self.write_to_parquet(data, self.path, self.temp_cache_path, upsert=upsert, primary_keys=primary_keys)
+        elif isinstance(data, pd.DataFrame):
+            if upsert:
+                raise RuntimeError("Parquet backend does not support upsert for pandas engine.")
+
+            part = None
+            if hasattr(data, 'cache_partition'):
+                part = data.cache_partition
+
+            data.to_parquet(self.path, partition_cols=part, engine='pyarrow')
+        else:
+            raise RuntimeError("Delta backend only supports dask and pandas engines.")
+
+    def read(self, engine):
+        if engine == 'pandas' or engine == pd.DataFrame or engine is None:
+            return pd.read_parquet(self.path)
+        elif engine == 'dask' or engine == dd.DataFrame:
+            return dd.read_parquet(self.path, engine='pyarrow', ignore_metadata_file=True)
+        else:
+            raise RuntimeError("Delta backend only supports dask and pandas engines.")
+
+    def write_to_parquet(self, df, cache_path, temp_cache_path, upsert=False, primary_keys=None):
+        """Write a pandas or dask dataframe to a parquet."""
+        part = None
+        if hasattr(df, 'cache_partition'):
+            part = df.cache_partition
+
+        if upsert and self.fs.exists(cache_path):
+            print("Found existing cache for upsert.")
+            if primary_keys is None:
+                raise ValueError("Upsert may only be performed with primary keys specified")
+
+            if isinstance(df, pd.DataFrame):
+                print("Auto converting pandas to dask dataframe.")
+                df = dd.from_pandas(df)
+
+            if not isinstance(df, dd.DataFrame):
+                raise RuntimeError("Upsert is only supported by dask dataframes for parquet")
+
+            existing_df = read_from_parquet(cache_path)
+
+            # Record starting partitions
+            start_parts = df.npartitions
+            existing_parts = existing_df.npartitions
+
+            # Coearce dtypes before joining
+            for key in primary_keys:
+                if is_datetime(existing_df[key].dtype):
+                    # The only way I could get this to work was by removing timezones
+                    # many attempts to coerc df to existing df with the correct tz
+                    df[key] = df[key].dt.tz_localize(None)
+                    df[key] = dd.to_datetime(df[key], utc=False)
+                    existing_df[key] = existing_df[key].dt.tz_localize(None)
+                    existing_df[key] = dd.to_datetime(existing_df[key], utc=False)
+                elif df[key].dtype != existing_df[key].dtype:
+                    df[key] = df[key].astype(existing_df[key].dtype)
+
+
+            outer_join = existing_df.merge(df, how = 'outer', on=primary_keys, indicator = True, suffixes=('_drop',''))
+            new_rows = outer_join[(outer_join._merge == 'right_only')].drop('_merge', axis = 1)
+            cols_to_drop = [x for x in new_rows.columns if x.endswith('_drop')]
+            new_rows = new_rows.drop(columns=cols_to_drop)
+
+            # Now concat with existing df
+            new_rows = new_rows.astype(existing_df.dtypes)
+            new_rows = new_rows[list(existing_df.columns)]
+            final_df = dd.concat([existing_df, new_rows])
+
+            if len(new_rows.index) > 0:
+                final_df = final_df.repartition(npartitions=start_parts + existing_parts)
+
+                # Coearce dtypes and make the columns the same order
+
+                print("Copying cache for ``consistent'' upsert.")
+                if self.fs.exists(temp_cache_path):
+                    self.fs.rm(temp_cache_path, recursive=True)
+
+                write_parquet_helper(final_df, temp_cache_path, part)
+                print("Successfully appended rows to temp parquet. Overwriting existing cache.")
+
+                if self.fs.exists(cache_path):
+                    self.fs.rm(cache_path, recursive=True)
+
+                self.fs.cp(temp_cache_path, cache_path, recursive=True)
+
+            else:
+                print("No rows to upsert.")
+        else:
+            write_parquet_helper(df, cache_path, part)
+

nuthatch/backends/sql.py

@@ -0,0 +1,147 @@
+from nuthatch.backend import DatabaseBackend, register_backend
+from os.path import join
+import hashlib
+import sqlalchemy
+import uuid
+import pandas as pd
+import dask.dataframe as dd
+
+def hashed_table_name(table_name):
+    """Return a qualified postgres table name."""
+    return hashlib.md5(table_name.encode()).hexdigest()
+
+@register_backend
+class SQLBackend(DatabaseBackend):
+    """
+    SQL backend for caching tabular data in a SQL database.
+
+    This backend supports dask and pandas dataframes.
+    """
+
+    backend_name = "sql"
+
+    def __init__(self, cacheable_config, cache_key, namespace, args, backend_kwargs):
+        super().__init__(cacheable_config, cache_key, namespace, args, backend_kwargs)
+
+        if backend_kwargs and 'hash_table_name' in backend_kwargs:
+            self.table_name = hashed_table_name(cache_key)
+        else:
+            self.table_name = cache_key
+
+        if namespace:
+            self.table_name = namespace + '.' + self.table_name
+
+            if not self.write_engine.dialect.has_schema(self.write_engine, namespace):
+                self.write_engine.execute(sqlalchemy.schema.CreateSchema(namespace))
+
+
+
+    def write(self, data, upsert=False, primary_keys=None):
+        if upsert and self.exists():
+            if primary_keys is None or not isinstance(primary_keys, list):
+                raise ValueError("Upsert may only be performed with primary keys specified as a list.")
+
+            if not isinstance(data, dd.DataFrame):
+                raise RuntimeError("Upsert is only supported by dask dataframes for parquet")
+
+            with self.engine.begin() as conn:
+                print("SQL cache exists for upsert.")
+                # If it already exists...
+
+                # Extract the primary key columns for SQL constraint
+                for key in primary_keys:
+                    if key not in data.columns:
+                        raise ValueError("Dataframe MUST contain all primary keys as columns")
+
+                # Write a temporary table
+                temp_table_name = f"temp_{uuid.uuid4().hex[:6]}"
+
+                if isinstance(data, pd.DataFrame):
+                    data.to_sql(temp_table_name, self.engine, index=False)
+                elif isinstance(data, dd.DataFrame):
+                    data.to_sql(temp_table_name, uri=self.engine.url.render_as_string(hide_password=False), index=False, parallel=True, chunksize=10000)
+                else:
+                    raise RuntimeError("Did not return dataframe type.")
+
+                index_sql_txt = ", ".join([f'"{i}"' for i in primary_keys])
+                columns = list(data.columns)
+                headers = primary_keys + list(set(columns) - set(primary_keys))
+                headers_sql_txt = ", ".join(
+                    [f'"{i}"' for i in headers]
+                )  # index1, index2, ..., column 1, col2, ...
+
+                # col1 = exluded.col1, col2=excluded.col2
+                # Excluded statement updates values of rows where primary keys conflict
+                update_column_stmt = ", ".join([f'"{col}" = EXCLUDED."{col}"' for col in columns])
+
+                # For the ON CONFLICT clause, postgres requires that the columns have unique constraint
+                # To add if not exists must drop if exists and then add. In a transaction this is consistent
+
+                # Constraint IDs need to be globally unique to not conflict in the database
+                constraint_id =  hashlib.md5(self.cache_key.encode()).hexdigest()[:10]
+                query_pk = f"""
+                ALTER TABLE "{self.table_name}" DROP CONSTRAINT IF EXISTS unique_constraint_for_{constraint_id} CASCADE;
+                """
+
+                print("Adding a unique to contraint to table if it doesn't exist.")
+                conn.exec_driver_sql(query_pk)
+
+                query_pk = f"""
+                ALTER TABLE "{self.table_name}" ADD CONSTRAINT unique_constraint_for_{constraint_id}
+                UNIQUE ({index_sql_txt});
+                """
+                conn.exec_driver_sql(query_pk)
+
+                # Compose and execute upsert query
+                query_upsert = f"""INSERT INTO "{self.table_name}" ({headers_sql_txt})
+                                   SELECT {headers_sql_txt} FROM "{temp_table_name}"
+                                   ON CONFLICT ({index_sql_txt}) DO UPDATE
+                                   SET {update_column_stmt};
+                                   """
+                print("Upserting.")
+                conn.exec_driver_sql(query_upsert)
+                conn.exec_driver_sql(f"DROP TABLE {temp_table_name}")
+        else:
+            try:
+                if isinstance(data, pd.DataFrame):
+                    data.to_sql(self.table_name, self.write_engine, if_exists='replace', index=False)
+                    return data
+                elif isinstance(data, dd.DataFrame):
+                    data.to_sql(self.table_name, self.engine.url.render_as_string(hide_password=False), if_exists='replace', index=False, parallel=True, chunksize=10000)
+                else:
+                    raise RuntimeError("Did not return dataframe type.")
+
+                # Also log the table name in the tables table
+                pd_name = {'table_name': [self.cache_key], 'table_key': [self.table_name], 'created_at': [pd.Timestamp.now()]}
+                pd_name = pd.DataFrame(pd_name)
+                pd_name.to_sql('cache_tables', self.write_engine, if_exists='append')
+            except sqlalchemy.exc.InterfaceError:
+                raise RuntimeError("Error connecting to database.")
+
+    def read(self, engine=None):
+        if engine == 'pandas' or engine == pd.DataFrame or engine =='dask' or engine == dd.DataFrame or engine is None:
+            try:
+                data = pd.read_sql_query(f'select * from "{self.table_name}"', con=self.engine)
+                if engine == 'dask' or engine == dd.DataFrame:
+                    return dd.from_pandas(data)
+                else:
+                    return data
+            except sqlalchemy.exc.InterfaceError:
+                raise RuntimeError("Error connecting to database.")
+        else:
+            raise RuntimeError("SQL backend only supports pandas engine.")
+
+    def exists(self):
+        try:
+            insp = sqlalchemy.inspect(self.engine)
+            return insp.has_table(self.table_name)
+        except sqlalchemy.exc.InterfaceError:
+            raise RuntimeError("Error connecting to database.")
+
+    def get_file_path(self):
+        return join(self.engine.url.render_as_string(), self.table_name)
+
+    def delete(self):
+        metadata = sqlalchemy.MetaData()
+        table = sqlalchemy.Table(self.table_name, metadata)
+        table.drop(self.engine, checkfirst=True)