oasis-data-manager 0.2.3__py3-none-any.whl

oasis_data_manager/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.2.3'

oasis_data_manager/complex/complex.py

@@ -0,0 +1,173 @@
+import logging
+import pathlib
+from io import BytesIO
+from typing import List, Type
+
+import httpcore
+import httpx
+import pandas as pd
+
+from oasis_data_manager.df_reader.config import clean_config, get_df_reader
+from oasis_data_manager.df_reader.reader import OasisReader
+from oasis_data_manager.filestore.backends.local import LocalStorage
+
+logger = logging.getLogger(__name__)
+
+
+class Adjustment:
+    """
+    Adjustments are any Pandas adjustments made after the data is fetched and filtered by SQL if applicable.
+    """
+
+    @classmethod
+    def apply(cls, df):
+        return df
+
+
+class ComplexData:
+    adjustments: List[Type[Adjustment]] = []
+    filename: str = ""
+    url: str = ""
+    fetch_required: bool = True
+    sql: str = ""
+
+    def __init__(self, storage=None):
+        if not storage:
+            storage = LocalStorage()
+        self.storage = storage
+
+    def fetch(self):
+        raise NotImplementedError
+
+    def get_sql(self):
+        return self.sql
+
+    def adjust(self, reader) -> OasisReader:
+        """
+        Hook to apply any adjustments.
+
+        TODO adjustments are filters? Functions fun on the readers df i.e
+        apply, should be change filter to apply then or is that confusing with pandas?
+        """
+        if self.adjustments:
+            return reader.filter([a.apply for a in self.adjustments])
+        return reader
+
+    def to_dataframe(self, result) -> pd.DataFrame:
+        """
+        Hook to allow conversion of the fetch() to a dataframe, in the case of the
+        RestComplexData class for example, the result is json, so can be fed
+        directly to a dataframe and wrapped into our df_reader.
+        """
+        return pd.DataFrame(result)
+
+    def get_df_reader(self, filepath, **kwargs) -> OasisReader:
+        df_reader_config = clean_config(
+            {"filepath": filepath, "engine": "OasisDaskReader"}
+        )
+        df_reader_config["engine"]["options"]["storage"] = self.storage
+
+        return get_df_reader(df_reader_config, **kwargs)
+
+    def to_reader(self, fetch_result) -> OasisReader:
+        if fetch_result:
+            # When the result has been fetched, apply to_dataframe and pass directly into the df_reader
+            df = self.to_dataframe(fetch_result)
+            return self.get_df_reader(
+                None,  # TODO - None value instead of filename_or_buffer, improve this.
+                dataframe=df,
+                has_read=True,
+            )
+        else:
+            # Not fetched, let df_reader fetch as normal.
+            return self.get_df_reader(self.filename if self.filename else self.url)
+
+    def run(self):
+        # CSV and Parquet files are read directly by the df_reader, so fetch() is not needed.
+        # Only formats the df_reader can't handle (e.g. custom binary formats) require fetch().
+        if self.fetch_required and self.filename:
+            extension = pathlib.Path(self.filename).suffix
+            self.fetch_required = extension not in [".parquet", ".pq", ".csv"]
+
+        fetch_result = None
+        if self.fetch_required:
+            fetch_result = self.fetch()
+
+        reader = self.to_reader(fetch_result)
+
+        sql = self.get_sql()
+        if hasattr(reader, "sql") and sql:
+            reader = reader.sql(sql)
+
+        reader = self.adjust(reader)
+
+        # TODO store file? return for now
+        return reader
+
+
+class FileStoreComplexData(ComplexData):
+    def to_dataframe(self, result) -> pd.DataFrame:
+        """
+        As this is only called on filetypes not handled by the df_reader, this will always be custom.
+        """
+        raise NotImplementedError
+
+    def fetch(self):
+        with self.storage.open(self.filename, "rb") as f:
+            result = BytesIO(f.read())
+        return result
+
+
+class RestComplexData(ComplexData):
+    exceptions = (
+        httpx.RequestError,
+        httpcore.ReadTimeout,
+        httpcore.ConnectTimeout,
+        httpcore.ConnectError,
+    )
+
+    url: str
+    timeout: int = 10
+
+    def handle_error(self, exception, **kwargs):
+        logger.warning(
+            "Exception in complex data call",
+            extra={"exception": exception, "uri": self.url},
+        )
+        return None
+
+    def handle_response(self, response) -> dict:
+        return response.json()
+
+    def get_uri(self) -> str:
+        return self.url
+
+    def get_headers(self) -> dict:
+        return {}
+
+    def get_post_json(self) -> dict:
+        return {}
+
+    def fetch(self):
+        uri = self.get_uri()
+        headers = self.get_headers()
+        post_json = self.get_post_json()
+        timeout = self.timeout
+
+        try:
+            if post_json:
+                response = httpx.post(
+                    uri, json=post_json, headers=headers, timeout=timeout
+                )
+            else:
+                response = httpx.get(uri, headers=headers, timeout=timeout)
+
+            if response.status_code == 200:
+                return self.handle_response(response)
+            else:
+                return self.handle_error(
+                    exception="Unexpected status in complex data call",
+                    post_json=post_json,
+                )
+        except self.exceptions as e:
+            return self.handle_error(exception=str(e), post_json=post_json)

oasis_data_manager/complex/examples.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+
+import h5py
+import pandas as pd
+
+from oasis_data_manager.complex.complex import Adjustment, FileStoreComplexData
+from oasis_data_manager.df_reader.reader import OasisDaskReader
+from oasis_data_manager.filestore.backends.local import LocalStorage
+
+
+class AddColAdjustment(Adjustment):
+    @classmethod
+    def apply(cls, df):
+        df["else"] = "test"
+        return df
+
+
+class FloodDataExample(FileStoreComplexData):
+    # Note this file needs to be sourced. TODO - is this public? If not is there a public/file we can include?
+    filename = "tropical_cyclone_10synth_tracks_150arcsec_rcp26_KNA_2080.hdf5"
+    adjustments = [AddColAdjustment]
+    sql = "SELECT * FROM table WHERE event_id > 3000"
+
+    def to_dataframe(self, result) -> pd.DataFrame:
+        result = h5py.File(result)
+
+        df = pd.DataFrame(list(result["event_id"]), columns=["event_id"])
+        for m in ["event_name", "date", "frequency", "orig"]:
+            df[m] = list(result[m])
+        df = df.reset_index()
+
+        return df
+
+
+if __name__ == "__main__":
+    reader = OasisDaskReader
+    test_storage = LocalStorage("/tmp")
+
+    result = FloodDataExample(storage=test_storage).run()
+    print(result.as_pandas())

oasis_data_manager/config.py

@@ -0,0 +1,36 @@
+import importlib
+
+_ALIASES = {
+    # Storage backends
+    "LocalStorage": "oasis_data_manager.filestore.backends.local.LocalStorage",
+    "AwsS3Storage": "oasis_data_manager.filestore.backends.aws.AwsS3Storage",
+    "AzureABFSStorage": "oasis_data_manager.filestore.backends.azure.AzureABFSStorage",
+    # Reader backends
+    "OasisPandasReader": "oasis_data_manager.df_reader.backends.pandas.OasisPandasReader",
+    "OasisPandasReaderCSV": "oasis_data_manager.df_reader.backends.pandas.OasisPandasReaderCSV",
+    "OasisPandasReaderParquet": "oasis_data_manager.df_reader.backends.pandas.OasisPandasReaderParquet",
+    "OasisDaskReader": "oasis_data_manager.df_reader.backends.dask.OasisDaskReader",
+    "OasisDaskReaderCSV": "oasis_data_manager.df_reader.backends.dask.OasisDaskReaderCSV",
+    "OasisDaskReaderParquet": "oasis_data_manager.df_reader.backends.dask.OasisDaskReaderParquet",
+    "OasisPyarrowReader": "oasis_data_manager.df_reader.backends.pyarrow.OasisPyarrowReader",
+}
+
+
+class ConfigError(Exception):
+    pass
+
+
+def load_class(path, base=None):
+    path = _ALIASES.get(path, path)
+    path_split = path.rsplit(".", 1)
+    if len(path_split) != 2:
+        raise ConfigError(f"'{path}' is not a valid class path (expected 'module.ClassName' or a known alias)")
+
+    module_path, cls_name = path_split
+    module = importlib.import_module(module_path)
+    cls = getattr(module, cls_name)
+
+    if base and not issubclass(cls, base):
+        raise ConfigError(f"'{cls.__name__}' does not extend '{base.__name__}'")
+
+    return cls

oasis_data_manager/df_reader/backends/base.py

@@ -0,0 +1,119 @@
+import pathlib
+from typing import Iterable
+
+from ...filestore.backends.base import BaseStorage
+
+
+class OasisReader:
+    """
+    Base reader.
+
+    as_pandas(), sql() & filter() can all be chained with self.has_read controlling whether the base
+    read (read_csv/read_parquet) needs to be triggered. This is because in the case of spark
+    we need to read differently depending on if the intention is to do sql or filter.
+    """
+
+    def __init__(
+        self,
+        filename_or_buffer,
+        storage: BaseStorage,
+        *args,
+        dataframe=None,
+        has_read=False,
+        **kwargs,
+    ):
+        self.filename_or_buffer = filename_or_buffer
+        self.storage = storage
+        self._df = dataframe
+        self.has_read = has_read
+        self.reader_args = args
+        self.reader_kwargs = kwargs
+
+        if not filename_or_buffer:
+            if dataframe is None and not has_read:
+                raise RuntimeError(
+                    "Reader must be initialised with either a "
+                    "filename_or_buffer or by passing a dataframe "
+                    "and has_read=True"
+                )
+            else:
+                self.read_from_dataframe()
+
+        if (
+            filename_or_buffer
+            and isinstance(self.filename_or_buffer, str)
+            and self.filename_or_buffer.lower().endswith(".zip")
+        ):
+            self.reader_kwargs["compression"] = "zip"
+
+    @property
+    def df(self):
+        self._read()
+        return self._df
+
+    @df.setter
+    def df(self, other):
+        self._df = other
+
+    def read_csv(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def read_parquet(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _read(self):
+        if not self.has_read:
+            if hasattr(self.filename_or_buffer, "name"):
+                parts = pathlib.Path(self.filename_or_buffer.name).parts
+            else:
+                parts = pathlib.Path(self.filename_or_buffer).parts
+
+            is_parquet = any(
+                part.endswith((".parquet", ".pq")) for part in parts
+            )
+
+            # Set has_read before calling read to prevent re-entrant calls (e.g. Dask
+            # readers access self.df internally during read). Reset on failure so the
+            # read can be retried.
+            self.has_read = True
+            try:
+                if is_parquet:
+                    self.read_parquet(*self.reader_args, **self.reader_kwargs)
+                else:
+                    # assume the file is csv if not parquet
+                    self.read_csv(*self.reader_args, **self.reader_kwargs)
+            except Exception:
+                self.has_read = False
+                raise
+
+        return self
+
+    def copy_with_df(self, df):
+        return type(self)(
+            self.filename_or_buffer, self.storage, dataframe=df, has_read=self.has_read
+        )
+
+    def filter(self, filters):
+        self._read()
+
+        df = self._df
+        for df_filter in filters if isinstance(filters, Iterable) else [filters]:
+            df = df_filter(df)
+
+        return self.copy_with_df(df)
+
+    def sql(self, sql):
+        if sql:
+            self._read()
+            return self.apply_sql(sql)
+        return self
+
+    def query(self, fn):
+        return fn(self.df)
+
+    def as_pandas(self):
+        self._read()
+        return self._df
+
+    def read_from_dataframe(self):
+        pass

oasis_data_manager/df_reader/backends/dask.py

@@ -0,0 +1,191 @@
+import io
+import pathlib
+import logging
+
+import dask
+import dask_geopandas as dgpd
+from dask import dataframe as dd
+from dask_sql import Context
+from dask_sql.utils import ParsingException
+from distributed import Client
+
+from ..exceptions import InvalidSQLException
+from .base import OasisReader
+
+dask.config.set(
+    {"dataframe.convert-string": False}
+)  # allows dask sql to support pyarrow
+logger = logging.getLogger("oasis_data_manager.df_reader.reader")
+
+
+class OasisDaskReader(OasisReader):
+    sql_table_name = "table"
+
+    def __init__(self, *args, client_address=None, **kwargs):
+        if client_address:
+            self.client = Client(client_address, set_as_default=False)
+        else:
+            self.client = None
+
+        self.sql_context = Context()
+        self.table_names = [self.sql_table_name]
+        self.pre_sql_columns = []
+
+        super().__init__(*args, **kwargs)
+
+    def copy_with_df(self, df):
+        if not isinstance(df, dd.DataFrame):
+            df = dd.from_pandas(df, npartitions=1)
+        res = super().copy_with_df(df)
+        res.client = self.client
+        return res
+
+    def apply_geo(self, shape_filename_path, *args, drop_geo=True, **kwargs):
+        """
+        Read in a shape file and return the _read file with geo data joined.
+        """
+        # TODO: fix this so that it can work with non local files
+        # with self.storage.open(self.shape_filename_path) as f:
+        #     shape_df = dgpd.read_file(f, npartitions=1)
+
+        shape_df = dgpd.read_file(shape_filename_path, npartitions=1)
+
+        # for situations where the columns in the source data are different.
+        lon_col = kwargs.get("geo_lon_col", "longitude")
+        lat_col = kwargs.get("geo_lat_col", "latitude")
+
+        df_columns = self.df.columns.tolist()
+        if lat_col not in df_columns or lon_col not in df_columns:
+            logger.warning("Invalid shape file provided")
+            # temp until we decide on handling, i.e don't return full data if it fails.
+            return self.copy_with_df(dd.DataFrame.from_dict({}, npartitions=1))
+
+        df = self.df.copy()
+
+        # convert read df to geo
+        df["geometry"] = dgpd.points_from_xy(df, lon_col, lat_col)
+        df = dgpd.from_dask_dataframe(df)
+
+        # Make sure they're using the same projection reference
+        df.crs = shape_df.crs
+
+        # join the datasets, matching `geometry` to points within the shape df
+        df = df.sjoin(shape_df, how="inner")
+
+        if drop_geo:
+            df = df.drop(shape_df.columns.tolist() + ["index_right"], axis=1)
+
+        return self.copy_with_df(df)
+
+    def apply_sql(self, sql):
+        try:
+            # Initially this was the filename, but some filenames are invalid for the table,
+            # is it ok to call it the same name all the time? Mapped to DaskDataTable in case
+            # we need to change this.
+            self.sql_context.create_table("DaskDataTable", self.df)
+            formatted_sql = sql.replace(self.sql_table_name, "DaskDataTable")
+
+            # Combine columns from join() tables with current df columns for case restoration
+            col_map = {}
+            for col in list(self.pre_sql_columns) + list(self.df.columns):
+                col_map.setdefault(col.lower(), col)
+
+            # dask expects the columns to be lower case, which won't match some data
+            df = self.sql_context.sql(
+                formatted_sql,
+                config_options={"sql.identifier.case_sensitive": False},
+            )
+            # which means we then need to map the columns back to the original
+            # and allow for any aggregations to be retained
+            df.columns = [col_map.get(v.lower(), v) for v in df.columns]
+
+            return self.copy_with_df(df)
+        except ParsingException:
+            raise InvalidSQLException
+
+    def join(self, df, table_name):
+        """
+        Creates a secondary table as a sql table in order to allow joins when apply_sql is called.
+        """
+        if table_name in self.table_names:
+            raise RuntimeError(
+                f"Table name already in use: [{','.join(self.table_names)}]"
+            )
+        self.pre_sql_columns.extend(df.columns)
+        self.sql_context.create_table(table_name, df)
+        self.table_names.append(table_name)
+        return self
+
+    def read_from_dataframe(self):
+        if not isinstance(self.df, dd.DataFrame):
+            self.df = dd.from_pandas(self.df, npartitions=1)
+
+    def as_pandas(self):
+        super().as_pandas()
+        if self.client:
+            return self.client.compute(self.df).result()
+        else:
+            return self.df.compute()
+
+    def read_dict(self, data):
+        self.df = dd.DataFrame.from_dict(data)
+
+    def read_csv(self, *args, **kwargs):
+        # remove standard pandas kwargs which will case an issue in dask.
+        dask_safe_kwargs = kwargs.copy()
+        dask_safe_kwargs.pop("memory_map", None)
+        dask_safe_kwargs.pop("low_memory", None)
+
+        filename_or_buffer = self.filename_or_buffer
+        if isinstance(filename_or_buffer, pathlib.PosixPath):
+            filename_or_buffer = str(self.filename_or_buffer)
+
+        if isinstance(filename_or_buffer, io.TextIOWrapper) or isinstance(
+            filename_or_buffer, io.BufferedReader
+        ):
+            filename_or_buffer = filename_or_buffer.name
+
+        # django files
+        if hasattr(filename_or_buffer, "path"):
+            filename_or_buffer = filename_or_buffer.path
+
+        _, uri = self.storage.get_storage_url(filename_or_buffer, encode_params=False)
+        self.df = dd.read_csv(
+            uri,
+            *args,
+            **dask_safe_kwargs,
+            storage_options=self.storage.get_fsspec_storage_options(),
+        )
+
+    def read_parquet(self, *args, **kwargs):
+        if isinstance(self.filename_or_buffer, str):
+            _, uri = self.storage.get_storage_url(
+                self.filename_or_buffer, encode_params=False
+            )
+            filename = uri
+            kwargs["storage_options"] = self.storage.get_fsspec_storage_options()
+        else:
+            filename = self.filename_or_buffer
+
+        self.df = dd.read_parquet(
+            filename,
+            *args,
+            **kwargs,
+        )
+
+        # dask-sql doesn't handle categorical columns, but we need to be careful
+        # how we convert them, if an assign is used we will end up stopping
+        # the `Predicate pushdown optimization` within dask-sql from applying the
+        # sql to the read_parquet filters.
+        categories_to_convert = {}
+        for col in self.df.select_dtypes(include="category").columns:
+            categories_to_convert[col] = self.df[col].dtype.categories.dtype
+        self.df = self.df.astype(categories_to_convert)
+
+
+class OasisDaskReaderCSV(OasisDaskReader):
+    pass
+
+
+class OasisDaskReaderParquet(OasisDaskReader):
+    pass

oasis_data_manager/df_reader/backends/pandas.py

@@ -0,0 +1,85 @@
+import logging
+
+import pandas as pd
+
+
+try:
+    import geopandas as gpd
+except ModuleNotFoundError:
+    gpd = None
+
+from .base import OasisReader
+from ..exceptions import MissingOptionalDependency
+logger = logging.getLogger("oasis_data_manager.df_reader.reader")
+
+
+class OasisPandasReader(OasisReader):
+    def _read_with(self, read_fn, *args, **kwargs):
+        if isinstance(self.filename_or_buffer, str) and not self.filename_or_buffer.startswith(("http://", "https://")):
+            _, uri = self.storage.get_storage_url(
+                self.filename_or_buffer, encode_params=False
+            )
+            self.df = read_fn(
+                uri,
+                *args,
+                **kwargs,
+                storage_options=self.storage.get_fsspec_storage_options(),
+            )
+        else:
+            self.df = read_fn(self.filename_or_buffer, *args, **kwargs)
+
+    def read_csv(self, *args, **kwargs):
+        if pd.__version__ >= "3":  # remove unsupported options issue https://github.com/OasisLMF/OasisLMF/issues/1896
+            kwargs.pop('low_memory', None)
+        self._read_with(pd.read_csv, *args, **kwargs)
+
+    def read_parquet(self, *args, **kwargs):
+        self._read_with(pd.read_parquet, *args, **kwargs)
+
+    def apply_geo(self, shape_filename_path, *args, drop_geo=True, **kwargs):
+        """
+        Read in a shape file and return the _read file with geo data joined.
+        """
+        # TODO: fix this so that it can work with non local files
+        # with self.storage.open(self.shape_filename_path) as f:
+        #     shape_df = gpd.read_file(f)
+
+        if gpd is None:
+            raise MissingOptionalDependency(
+                "Missing optional dependency 'geopandas' for 'apply_geo' method, install package using `pip install oasis-data-manager[extra]`")
+
+        shape_df = gpd.read_file(shape_filename_path)
+
+        # for situations where the columns in the source data are different.
+        lon_col = kwargs.get("geo_lon_col", "longitude")
+        lat_col = kwargs.get("geo_lat_col", "latitude")
+
+        df_columns = self.df.columns.tolist()
+        if lat_col not in df_columns or lon_col not in df_columns:
+            logger.warning("Invalid shape file provided")
+            # temp until we decide on handling, i.e don't return full data if it fails.
+            return self.copy_with_df(pd.DataFrame.from_dict({}))
+
+        # convert read df to geo
+        df = gpd.GeoDataFrame(
+            self.df, geometry=gpd.points_from_xy(self.df[lon_col], self.df[lat_col])
+        )
+
+        # Make sure they're using the same projection reference
+        df.crs = shape_df.crs
+
+        # join the datasets, matching `geometry` to points within the shape df
+        df = df.sjoin(shape_df, how="inner")
+
+        if drop_geo:
+            df = df.drop(shape_df.columns.tolist() + ["index_right"], axis=1)
+
+        return self.copy_with_df(df)
+
+
+class OasisPandasReaderCSV(OasisPandasReader):
+    pass
+
+
+class OasisPandasReaderParquet(OasisPandasReader):
+    pass