datazone-sdk 1.0.0__tar.gz

datazone_sdk-1.0.0/PKG-INFO

@@ -0,0 +1,56 @@
+Metadata-Version: 2.4
+Name: datazone-sdk
+Version: 1.0.0
+Summary: Database and Delta storage client library for working with Delta Lake tables
+Author: Team Enigma
+Author-email: enigma@energinet.dk
+Requires-Python: >=3.10
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: datamazing (>=5.1.6)
+Requires-Dist: deltalake (==1.2.1)
+Requires-Dist: obstore (>=0.8.2)
+Requires-Dist: pandas (>=2.0.3,<3)
+Requires-Dist: polars (>=1.33.1)
+Requires-Dist: pyarrow (>=19.0.0)
+Description-Content-Type: text/markdown
+
+# Datazone SDK
+
+Database and Delta storage client library for working with Delta Lake tables.
+
+## Overview
+
+This package provides functionality for interacting with Delta Lake tables, including:
+
+- **Database Client**: High-level client for querying Delta Lake tables with support for time intervals, time travel, and filtering.
+- **Delta Storage**: Low-level components for working with Delta tables, schemas, and data types.
+
+## Installation
+
+```bash
+pip install datazone-sdk
+```
+
+## Usage
+
+```python
+import datazone as dz
+
+# Create a client from a storage account name
+client = dz.DatabaseClient.from_resource_name(
+    storage_account="<storage-account>",
+    container_name="<container-name>",
+    sub_path="<sub-path>",
+)
+
+# Query a table
+df = client.query(
+    table_name="my_table",
+    filters={"column": "value"},
+)
+```

datazone_sdk-1.0.0/README.md

@@ -0,0 +1,35 @@
+# Datazone SDK
+
+Database and Delta storage client library for working with Delta Lake tables.
+
+## Overview
+
+This package provides functionality for interacting with Delta Lake tables, including:
+
+- **Database Client**: High-level client for querying Delta Lake tables with support for time intervals, time travel, and filtering.
+- **Delta Storage**: Low-level components for working with Delta tables, schemas, and data types.
+
+## Installation
+
+```bash
+pip install datazone-sdk
+```
+
+## Usage
+
+```python
+import datazone as dz
+
+# Create a client from a storage account name
+client = dz.DatabaseClient.from_resource_name(
+    storage_account="<storage-account>",
+    container_name="<container-name>",
+    sub_path="<sub-path>",
+)
+
+# Query a table
+df = client.query(
+    table_name="my_table",
+    filters={"column": "value"},
+)
+```

datazone_sdk-1.0.0/datazone/__init__.py

@@ -0,0 +1,2 @@
+from datazone.db.client import DatabaseClient
+from datazone.deltastorage import Field, HyperSlice, Schema, Store, Table

datazone_sdk-1.0.0/datazone/db/__init__.py

@@ -0,0 +1 @@
+from . import client

datazone_sdk-1.0.0/datazone/db/client.py

@@ -0,0 +1,233 @@
+from typing import Optional
+
+import datamazing.pandas as pdz
+import pandas as pd
+import pyarrow.compute as pc
+
+from datazone.deltastorage.slicing import HyperSlice
+from datazone.deltastorage.store import Store
+
+
+class DatabaseClient:
+    def __init__(
+        self,
+        path: str,
+        storage_options: dict[str, str] | None = None,
+        table_prefix: str = "",
+    ):
+        self.store = Store(
+            path=path,
+            storage_options=storage_options,
+        )
+        self.table_prefix = table_prefix
+
+    @classmethod
+    def from_resource_name(
+        cls,
+        storage_account: str,
+        container_name: str = "datasets",
+        sub_path: str = "",
+        table_prefix: str = "",
+    ):
+        """Create a DatabaseClient from resource name (storage account).
+        This assumes the path of the delta lake is of the form:
+        abfss://{container_name}@{storage_account}.dfs.core.windows.net/{sub_path}
+
+        Args:
+            storage_account (str): Storage account name.
+            container_name (str, optional): Container name. Defaults to "datasets".
+            sub_path (str, optional): Sub-path within the container. Defaults to "".
+            table_prefix (str, optional): Table prefix to use (e.g. `mz_` for archive).
+            Defaults to "".
+            credential (optional): Azure credential to use.
+            Defaults to DefaultAzureCredential().
+        """
+        path = f"abfss://{container_name}@{storage_account}.dfs.core.windows.net"
+        if sub_path:
+            path += f"/{sub_path}"
+
+        storage_options = Store._get_func_storage_options()
+
+        return cls(
+            path=path, storage_options=storage_options, table_prefix=table_prefix
+        )
+
+    def get_unit_and_multiple(self, timedelta: pd.Timedelta) -> tuple[str | None, int]:
+        """
+        Get unit and multiple of a timedelta. E.g. for a timedelta of "PT5M" then
+        unit = "minute" and multiple = 5.
+        NOTE: Timedelta must have one and only one non-zero component,
+        i.e. "PT0S" doesnt work, and neither does "PT5M10S".
+
+        Args:
+            timedelta (pd.Timedelta): Timedelta
+
+        Returns:
+            tuple[str, int]: Unit and multiple
+        """
+        components = timedelta.components._asdict()
+
+        # remove plural ending from unit, since
+        # this is the standard pyarrow uses
+        components = {k[:-1]: v for k, v in components.items()}
+
+        non_zero_components = {
+            unit: multiple for unit, multiple in components.items() if multiple != 0
+        }
+
+        if len(non_zero_components) == 0:
+            return None, 0
+
+        if len(non_zero_components) != 1:
+            raise ValueError("Timedelta must have one and only one non-zero multiple.")
+
+        return next(iter(non_zero_components.items()))
+
+    def relative_time_travel_version(
+        self, time_column: str, block: pd.Timedelta, horizon: pd.Timedelta
+    ) -> pc.Expression:
+        """
+        Get value to use for filtering a relative time travel
+        (i.e. the interval [valid-from, valid-to] must contain
+        this value)
+        """
+        unit, multiple = self.get_unit_and_multiple(block)
+
+        if multiple == 0:
+            # `pc.floor_temporal` fails with multiple=0,
+            # but in this case we don't need to floor
+            # the time anyway
+            start_of_block = pc.field("time_utc")
+        else:
+            start_of_block = pc.floor_temporal(
+                pc.field(time_column),
+                multiple=multiple,
+                unit=unit,
+            )
+
+        return start_of_block - horizon.to_pytimedelta()
+
+    def time_travel_filter(
+        self,
+        time_travel: pdz.TimeTravel,
+        time_column: str,
+        valid_from_column: str,
+        valid_to_column: str,
+    ) -> list[HyperSlice]:
+        """Filter delta table on a time travel
+
+        Args:
+            time_travel (pdz.TimeTravel): Time travel
+            time_column (str): Time column name
+            valid_from_column (str): Valid-from column name
+            valid_to_column (str): Valid-to column name
+        """
+        match time_travel.tense:
+            case "absolute":
+                # If the time travel is absolute, we filter
+                # to entries where [valid-from, valid-to]
+                # contains `as_of_time`
+                version = time_travel.as_of_time.to_pydatetime()
+            case "relative":
+                version = self.relative_time_travel_version(
+                    time_column, time_travel.block, time_travel.horizon
+                )
+
+        return [
+            HyperSlice((valid_from_column, "<=", version)),
+            HyperSlice((valid_to_column, ">", version)),
+        ]
+
+    def query(
+        self,
+        table_name: str,
+        time_interval: Optional[pdz.TimeInterval] = None,
+        time_travel: Optional[pdz.TimeTravel] = None,
+        filters: Optional[dict[str, object]] = None,
+        columns: Optional[list[str]] = None,
+        include_validity_period_columns: bool = False,
+        include_generated_columns: bool = False,
+    ) -> pd.DataFrame:
+        """Query table.
+        Query defaults are set to match old Table Storage client behavior.
+        Time travel defaults to "as of now"
+        Validity period columns are dropped by default.
+        Generated columns are dropped by default.
+
+        Args:
+            table_name (str): Name of the table
+            time_interval (Optional[pdz.TimeInterval], optional): Time interval for the
+            query. Defaults to None.
+            time_travel (Optional[pdz.TimeTravel], optional): Time travel information.
+            Defaults to None.
+            filters (Optional[dict[str, object]], optional): Filters to apply to the
+            query.
+            Defaults to None.
+            columns (Optional[list[str]], optional): Columns to return.
+            Selecting columns can significantly improve query performance.
+            Defaults to None, meaning all columns will be returned.
+            include_validity_period_columns (bool, optional): Whether to include
+            validity period columns in the result;
+            (`valid_from_time_utc`, `valid_to_time_utc`).
+            Defaults to False.
+            include_generated_columns (bool, optional): Whether to include generated
+            columns in the result; (e.g. `valid_from_time_utc`, `valid_to_time_utc`).
+            Defaults to False.
+
+        Returns:
+            pd.DataFrame: The result of the query.
+        """
+        # Prefix used to differentiate between operation ("{table_name}")
+        # and historical ("mz_{table_name}").
+        table_name = f"{self.table_prefix}{table_name}"
+
+        table = self.store.get_table(table_name)
+        hyper_slice = []
+
+        if filters:
+            for key, value in filters.items():
+                if isinstance(value, (list, tuple, set)):
+                    hyper_slice.append((key, "in", value))
+                else:
+                    hyper_slice.append((key, "=", value))
+        if time_interval:
+            hyper_slice.append(("time_utc", ">=", time_interval.left))
+            hyper_slice.append(("time_utc", "<=", time_interval.right))
+
+        if time_travel is None:
+            time_travel = pdz.TimeTravel(
+                as_of_time=pd.Timestamp.utcnow(),
+            )
+
+        tt_filter = self.time_travel_filter(
+            time_travel,
+            time_column="time_utc",
+            valid_from_column="valid_from_time_utc",
+            valid_to_column="valid_to_time_utc",
+        )
+
+        hyper_slice.extend(tt_filter)
+        pl_df = table.read(hyper_slice=HyperSlice(hyper_slice), columns=columns)
+
+        pd_df = pl_df.to_pandas()
+
+        # We truncate to second, and change to nanosecond
+        # precision because this was used by the old solution (Azure Table Storage)
+        for col in pd_df.select_dtypes(include=["datetime", "datetimetz"]).columns:
+            pd_df[col] = pd_df[col].dt.floor("s").dt.as_unit("ns")
+
+        # Drop generated columns
+        if not include_generated_columns:
+            generated_cols = []
+            for field in table.schema().fields:
+                if field.generated_as is not None:
+                    generated_cols.append(field.column_name)
+            pd_df = pd_df.drop(columns=generated_cols, errors="ignore")
+
+        # Drop valid-from/to columns
+        if not include_validity_period_columns:
+            pd_df = pd_df.drop(
+                columns=["valid_from_time_utc", "valid_to_time_utc"], errors="ignore"
+            )
+
+        return pd_df

datazone_sdk-1.0.0/datazone/deltastorage/__init__.py

@@ -0,0 +1,4 @@
+from .schema import Field, Schema
+from .slicing import HyperSlice
+from .store import Store
+from .table import Table

datazone_sdk-1.0.0/datazone/deltastorage/data_types.py

@@ -0,0 +1,94 @@
+from abc import ABC, abstractmethod
+
+import pyarrow as pa
+
+
+class DataType(ABC):
+    @classmethod
+    def from_arrow(cls, pa_type: pa.DataType) -> "DataType":
+        if pa.types.is_string(pa_type) or pa.types.is_large_string(pa_type):
+            return String()
+        elif pa.types.is_floating(pa_type):
+            return Float()
+        elif pa.types.is_integer(pa_type):
+            return Int()
+        elif pa.types.is_timestamp(pa_type):
+            return Timestamp(tz=pa_type.tz)
+        elif pa.types.is_date(pa_type):
+            return Date()
+        elif pa.types.is_boolean(pa_type):
+            return Boolean()
+        elif pa.types.is_null(pa_type):
+            return Null()
+        else:
+            raise ValueError(f"Unsupported data type: {pa_type}")
+
+    @abstractmethod
+    def __str__(self):
+        ...
+
+    def __eq__(self, other: "DataType") -> bool:
+        return self.to_arrow() == other.to_arrow()
+
+    @abstractmethod
+    def to_arrow(self) -> pa.DataType:
+        ...
+
+
+class String(DataType):
+    def __str__(self):
+        return "string"
+
+    def to_arrow(self) -> pa.DataType:
+        return pa.string()
+
+
+class Float(DataType):
+    def __str__(self):
+        return "float"
+
+    def to_arrow(self) -> pa.DataType:
+        return pa.float64()
+
+
+class Int(DataType):
+    def __str__(self):
+        return "int"
+
+    def to_arrow(self) -> pa.DataType:
+        return pa.int64()
+
+
+class Timestamp(DataType):
+    def __str__(self):
+        return f"timestamp[{self.tz}]"
+
+    def __init__(self, tz: str):
+        self.tz = tz
+
+    def to_arrow(self) -> pa.DataType:
+        return pa.timestamp("us", tz=self.tz)
+
+
+class Date(DataType):
+    def __str__(self):
+        return "date"
+
+    def to_arrow(self) -> pa.DataType:
+        return pa.date32()
+
+
+class Boolean(DataType):
+    def __str__(self):
+        return "boolean"
+
+    def to_arrow(self) -> pa.DataType:
+        return pa.bool_()
+
+
+class Null(DataType):
+    def __str__(self):
+        return "null"
+
+    def to_arrow(self) -> pa.DataType:
+        return pa.null()

datazone_sdk-1.0.0/datazone/deltastorage/generated_columns.py

@@ -0,0 +1,151 @@
+import copy
+import datetime as dt
+from abc import ABC, abstractmethod
+from typing import Any
+from zoneinfo import ZoneInfo
+
+import polars as pl
+
+
+class GeneratedColumn(ABC):
+    """Class representing a generated column."""
+
+    @property
+    @abstractmethod
+    def base_column_names(self) -> list[str]:
+        ...
+
+    @abstractmethod
+    def to_metadata(self) -> dict:
+        ...
+
+    @abstractmethod
+    def expression(self) -> pl.Expr:
+        """Polars expression to compute the generated column."""
+        ...
+
+    @abstractmethod
+    def get_generated_conditions(self, op: str, value) -> list[tuple[str, Any]]:
+        """Get conditions on the generated column based on a condition
+        on the base column."""
+        ...
+
+    @classmethod
+    def from_metadata(cls, metadata: dict):
+        if metadata["type"] == "date_bucket":
+            return DateBucket(
+                base_column_name=metadata["base_column_name"],
+                as_tz=metadata.get("as_tz", None),
+            )
+        elif metadata["type"] == "concat":
+            return Concat(
+                base_column_names=metadata["base_column_names"],
+                delimiter=metadata["delimiter"],
+            )
+        else:
+            raise ValueError(f"Unknown generated column type: {metadata['type']}")
+
+
+class DateBucket(GeneratedColumn):
+    def __init__(self, base_column_name: str, as_tz: str | None = None):
+        """Generated column which bins a timestamp column to a date,
+        optionally converting to a specific timezone first.
+
+        Args:
+            base_column_name (str): Base column name
+            as_tz (str, optional): Cast to date in this timezone.
+                Defaults to None, meaning it uses the original timezone.
+        """
+        self.base_column_name = base_column_name
+        self.as_tz = as_tz
+
+    def __str__(self):
+        if self.as_tz is not None:
+            return f"date_bucket({self.base_column_name}, as_tz={self.as_tz})"
+        else:
+            return f"date_bucket({self.base_column_name})"
+
+    @property
+    def base_column_names(self) -> list[str]:
+        return [self.base_column_name]
+
+    def to_metadata(self) -> dict:
+        metadata = {
+            "type": "date_bucket",
+            "base_column_name": self.base_column_name,
+        }
+        if self.as_tz is not None:
+            metadata["as_tz"] = self.as_tz
+        return metadata
+
+    def expression(self) -> pl.Expr:
+        expr = pl.col(self.base_column_name)
+        if self.as_tz is not None:
+            expr = expr.dt.convert_time_zone(self.as_tz)
+        return expr.dt.date()
+
+    def get_generated_conditions(
+        self, op: str, value: dt.datetime
+    ) -> list[tuple[str, dt.date]]:
+        """Get conditions on the generated column based on a condition
+        on the base column.
+
+        Args:
+            op (str): Operator
+            value (dt.datetime): Value
+
+        Returns:
+            tuple[str, dt.date]: List of conditions on the generated column
+        """
+        timestamp = copy.copy(value)
+        if self.as_tz is not None:
+            timestamp = timestamp.astimezone(ZoneInfo(self.as_tz))
+        date = timestamp.date()
+
+        match op:
+            case "=":
+                return [("=", date)]
+            case ("<" | "<="):
+                return [("<=", date)]
+            case (">" | ">="):
+                return [(">=", date)]
+            case _:
+                # for other operations, we cannot make any
+                # useful filters on the generated column
+                return []
+
+
+class Concat(GeneratedColumn):
+    def __init__(self, base_column_names: list[str], delimiter: str):
+        """Generated column which concats multiple columns into a single string.
+
+        Args:
+            base_column_names (str): Base column names
+            delimiter (str): Delimiter used for concatenation
+        """
+        self._base_column_names = base_column_names
+        self.delimiter = delimiter
+
+    @property
+    def base_column_names(self) -> list[str]:
+        return self._base_column_names
+
+    def __str__(self):
+        base_cols_str = ", ".join(self.base_column_names)
+        return f"concat([{base_cols_str}], delimiter='{self.delimiter}')"
+
+    def to_metadata(self) -> dict:
+        metadata = {
+            "type": "concat",
+            "base_column_names": self.base_column_names,
+            "delimiter": self.delimiter,
+        }
+        return metadata
+
+    def expression(self) -> pl.Expr:
+        return pl.concat_str(self.base_column_names, separator=self.delimiter)
+
+    def get_generated_conditions(self, op: str, value: Any) -> list[tuple[str, Any]]:
+        # for concat generated columns, we cannot make any
+        # useful filters on the generated column
+        return []

datazone_sdk-1.0.0/datazone/deltastorage/schema.py

@@ -0,0 +1,135 @@
+import json
+from typing import Optional
+
+import polars as pl
+import pyarrow as pa
+
+from .data_types import DataType
+from .generated_columns import GeneratedColumn
+from .slicing import HyperSlice
+
+
+class Field:
+    def __init__(
+        self,
+        column_name: str,
+        data_type: DataType,
+        generated_as: Optional[GeneratedColumn] = None,
+    ):
+        """Class representing a table field.
+
+        Args:
+            column_name (str): Column name
+            data_type (DataType): Data type
+            generated_as (GeneratedColumn, optional): Generated column based on a
+                regular column. Defaults to None, meaning the column is not generated.
+        """
+        self.column_name = column_name
+        self.data_type = data_type
+        self.generated_as = generated_as
+
+    def __eq__(self, other: "Field") -> bool:
+        return self.to_arrow().equals(other.to_arrow(), check_metadata=True)
+
+    def __repr__(self):
+        string = f"{self.column_name}: {self.data_type}"
+        if self.generated_as is not None:
+            string = f"{string} [generated as {self.generated_as}]"
+        return string
+
+    def to_arrow(self):
+        """Convert to pyarrow field. Information about generated columns
+        is stored as metadata."""
+        pa_dtype = self.data_type.to_arrow()
+
+        pa_metadata = None
+        if self.generated_as is not None:
+            pa_metadata = {
+                "generated_column": json.dumps(self.generated_as.to_metadata())
+            }
+
+        return pa.field(self.column_name, pa_dtype, metadata=pa_metadata)
+
+    @classmethod
+    def from_arrow(cls, pa_field: pa.Field) -> "Field":
+        """Convert from pyarrow field"""
+        data_type = DataType.from_arrow(pa_field.type)
+
+        generated_as = None
+
+        if pa_field.metadata is not None:
+            pa_metadata = {k.decode(): v.decode() for k, v in pa_field.metadata.items()}
+            if "generated_column" in pa_metadata:
+                gen_col_metadata = json.loads(pa_metadata["generated_column"])
+                generated_as = GeneratedColumn.from_metadata(gen_col_metadata)
+
+        return cls(pa_field.name, data_type, generated_as)
+
+
+class Schema:
+    def __init__(
+        self,
+        fields: list[Field],
+    ):
+        """Class representing a table schema.
+
+        Args:
+            fields (list[Field]): Schema fields
+        """
+        self.fields = fields
+
+    def __eq__(self, other: "Schema") -> bool:
+        return self.to_arrow().equals(other.to_arrow(), check_metadata=True)
+
+    def __repr__(self):
+        field_reprs = [repr(field) for field in self.fields]
+        return "\n".join(field_reprs)
+
+    def to_arrow(self):
+        """Convert to pyarrow schema."""
+        pa_fields = [field.to_arrow() for field in self.fields]
+        return pa.schema(pa_fields)
+
+    @classmethod
+    def from_arrow(cls, pa_schema: pa.Schema) -> "Schema":
+        """Convert from pyarrow schema"""
+        fields = []
+        for pa_field in pa_schema:
+            field = Field.from_arrow(pa_field)
+            fields.append(field)
+        return cls(fields)
+
+    def add_generated_columns(self, base_df: pl.DataFrame) -> pl.DataFrame:
+        """Add additional columns to a dataframe derived
+        from all generated columns in the schema.
+
+        Args:
+            base_df (pl.DataFrame): Input dataframe
+        """
+        generated_exprs = []
+        for field in self.fields:
+            if field.generated_as is not None:
+                expr = field.generated_as.expression().alias(field.column_name)
+                generated_exprs.append(expr)
+
+        return base_df.with_columns(generated_exprs)
+
+    def add_generated_filters(self, base_slice: HyperSlice) -> HyperSlice:
+        """Add additional filters to a hyperslice based from
+        all generated columns in the schema.
+
+        Args:
+            base_slice (HyperSlice): Input hyperslice
+        """
+        generated_slice = []
+        for col, op, val in base_slice:
+            for field in self.fields:
+                if field.generated_as is None:
+                    continue
+                if col not in field.generated_as.base_column_names:
+                    continue
+                for gen_op, gen_val in field.generated_as.get_generated_conditions(
+                    op, val
+                ):
+                    generated_slice.append((field.column_name, gen_op, gen_val))
+        return HyperSlice(list(base_slice) + generated_slice)

datazone_sdk-1.0.0/datazone/deltastorage/slicing.py

@@ -0,0 +1,22 @@
+from typing import Any
+
+
+class HyperSlice(list[tuple[str, str, Any]]):
+    """A list of tuples representing a n-dimensional slice of a table.
+    Each tuple corresponds to a filter applied on the table,
+    in the form of (column, operator, value). For example, the
+    following hyper slice represents all records in a table
+    where the country is 'Denmark' and the date is greater than
+    2000-01-01:
+
+        [
+            ("country", "=", "Denmark"),
+            ("date", ">", "2000-01-01"),
+        ]
+
+    In SQL, this would be equivalent to the WHERE clause:
+
+        country = 'Denmark' AND date > '2000-01-01'
+    """
+
+    ...

datazone_sdk-1.0.0/datazone/deltastorage/store.py

@@ -0,0 +1,138 @@
+import os
+
+import deltalake as dl
+import obstore as obs
+from deltalake.exceptions import TableNotFoundError as DeltaTableNotFoundError
+
+from .schema import Schema
+from .table import Table
+
+
+class Store:
+    def __init__(
+        self,
+        path: str,
+        storage_options: dict[str, str] | None = None,
+    ):
+        """Class representing a store containing datasets
+
+        Args:
+            path (str): Root directory containing Delta tables
+            storage_options (dict[str, str] | None, optional): Storage options used for
+                remote cloud storage. For more information on available options,
+                go to https://delta-io.github.io/delta-rs/integrations/object-storage/.
+                Defaults to None, corresponding to the local file system.
+        """
+        self.path = path
+        if storage_options is None:
+            storage_options = self._get_func_storage_options()
+        self.storage_options = storage_options
+        # We use obstore to interact with remote
+        # cloud storage for operations not directly
+        # supported by delta-rs (e.g. listing directories)
+        # We could use fsspec, but the `storage_options`
+        # used by delta-rs and fsspec are not compatible
+        self._obstore = obs.store.from_url(
+            url=path,
+            config=storage_options,
+        )
+
+    def __repr__(self):
+        return f"Store('{self.path}')"
+
+    def _get_table_uri(self, table_name: str) -> str:
+        return self.path + "/" + table_name
+
+    @staticmethod
+    def _get_func_storage_options() -> dict[str, str]:
+        """Get storage options.
+        This differ depending on whether we are running
+        in cloud or locally.
+        """
+        if "IDENTITY_ENDPOINT" in os.environ:
+            # When running in Azure Function, the environment variable IDENTITY_ENDPOINT
+            # will be set, and we use the managed identity to access the storage account
+            storage_options = {"azure_msi_endpoint": os.environ["IDENTITY_ENDPOINT"]}
+        else:
+            # When running locally, we use the Azure CLI to authenticate.
+            storage_options = {"use_azure_cli": "true"}
+
+        return storage_options
+
+    @classmethod
+    def from_func_environment_variable(cls, container_name: str = "datasets"):
+        """Create Store instance from environment variable.
+        This uses default storage options, created in `__init__`."""
+        storage_account_name = os.environ["OPERATIONAL_DATA_STORAGE_ACCOUNT"]
+        path = f"abfss://{container_name}@{storage_account_name}.dfs.core.windows.net"
+
+        return cls(path=path)
+
+    def list_tables(self) -> list[str]:
+        """List all Delta tables"""
+        return self._obstore.list_with_delimiter()["common_prefixes"]
+
+    def table_exists(self, table_name: str):
+        """Check if Delta table exists
+
+        Args:
+            table_name (str): Table name
+        """
+        # For some reason `deltalake.DeltaTable.is_deltatable()` can be very slow.
+        # deltalake has an issue open about this:
+        # https://github.com/delta-io/delta-rs/issues/3942
+        # For now we catch the exception when trying to load the table
+        try:
+            _ = dl.DeltaTable(
+                table_uri=self._get_table_uri(table_name),
+                storage_options=self.storage_options,
+                without_files=True,
+            )
+        except DeltaTableNotFoundError:
+            return False
+        return True
+
+    def create_table(
+        self,
+        table_name: str,
+        schema: Schema,
+        partition_by: list[str] | None = None,
+    ) -> Table:
+        """Create Delta table
+
+        Args:
+            table_name (str): Table name.
+            schema (pl.Schema): Table schema.
+            partition_by (list[str]): Partition columns.
+        """
+        if self.table_exists(table_name):
+            raise ValueError(f"Table with name '{table_name}' already exists")
+
+        if schema is None:
+            raise ValueError("Schema must be provided when creating a new table")
+
+        pa_schema = schema.to_arrow()
+
+        dl.DeltaTable.create(
+            table_uri=self._get_table_uri(table_name),
+            schema=pa_schema,
+            storage_options=self.storage_options,
+            partition_by=partition_by,
+            configuration={
+                "delta.deletedFileRetentionDuration": "interval 2 hours",
+                "delta.logRetentionDuration": "interval 4 hours",
+            },
+        )
+
+        return Table(self._get_table_uri(table_name), self.storage_options)
+
+    def get_table(self, table_name: str) -> Table:
+        """Get Delta table
+
+        Args:
+            table_name (str): Table name
+        """
+        if not self.table_exists(table_name):
+            raise ValueError(f"Table with name '{table_name}' does not exist")
+
+        return Table(self._get_table_uri(table_name), self.storage_options)

datazone_sdk-1.0.0/datazone/deltastorage/table.py

@@ -0,0 +1,210 @@
+from typing import Any, Optional
+
+import deltalake as dl
+import polars as pl
+import pyarrow as pa
+
+from .schema import Schema
+from .slicing import HyperSlice
+
+
+def _dnf_to_sql(dnf: list[tuple]) -> str:
+    """Convert DNF expression to SQL expression."""
+    if len(dnf) == 0:
+        return "1=1"
+
+    sql_parts = []
+    for col, op, val in dnf:
+        if op == "in":
+            assert isinstance(val, list)
+            lst = ", ".join([f"'{item}'" for item in val])
+            sql_parts.append(f"{col} IN ({lst})")
+        elif op in [">=", "<=", ">", "<", "="]:
+            sql_parts.append(f"{col} {op} '{val}'")
+        else:
+            raise ValueError(f"Unsupported operation: {op}")
+
+    return " AND ".join(sql_parts)
+
+
+class Table:
+    def __init__(
+        self,
+        table_uri: str,
+        storage_options: dict[str, str] | None = None,
+    ):
+        """Class representing a dataset
+
+        Args:
+            delta_table (dl.DeltaTable): Delta table
+        """
+        self.table_uri = table_uri
+        self.storage_options = storage_options
+
+        self.table_name = self.table_uri.split("/")[-1]
+        self._delta_table = None
+
+    def __repr__(self):
+        return f"Table('{self.table_name}')"
+
+    @property
+    def delta_table(self) -> dl.DeltaTable:
+        """Get the Delta table object.
+        As the `Table`-class is lazily initialized,
+        the `delta_table`-property is initialized when needed.
+        We do not cache it, which creates a little overhead, but reduces
+        the risk of false transaction issues when doing concurrent reads/writes.
+        This is important because using the same instance can lead to transaction
+        issues in delta as DeltaTable uses metadata (transaction id) from
+        the first time the object is instantiated.
+
+        The risk disappears when this instance is created within a lock.
+        """
+        return dl.DeltaTable(self.table_uri, storage_options=self.storage_options)
+
+    def partition_cols(self) -> list[str]:
+        """Get the partition columns of the table"""
+        return self.delta_table.metadata().partition_columns
+
+    def schema(self) -> Schema:
+        """Get the schema of the table"""
+        pa_schema = pa.schema(self.delta_table.schema())
+        return Schema.from_arrow(pa_schema)
+
+    def read(
+        self, hyper_slice: Optional[HyperSlice] = None, columns=None
+    ) -> pl.DataFrame:
+        """Read from Delta table
+
+        Args:
+            hyper_slice (HyperSlice): Hyper sliced used to filter data
+        """
+        if hyper_slice is None:
+            hyper_slice = []
+
+        # add generated filters to hyperslice
+        hyper_slice = self.schema().add_generated_filters(hyper_slice)
+
+        delta_table = self.delta_table
+        partition_cols = delta_table.metadata().partition_columns
+
+        if len(hyper_slice) == 0:
+            file_filters = None
+            partition_filters = None
+        else:
+            file_filters = hyper_slice
+            partition_filters = [f for f in hyper_slice if f[0] in partition_cols]
+
+        pyarrow_table_existing_data = delta_table.to_pyarrow_table(
+            columns=columns,
+            partitions=partition_filters,
+            filters=file_filters,
+        )
+
+        return pl.from_arrow(pyarrow_table_existing_data)
+
+    def _to_writable_pyarrow_table(self, df: pl.DataFrame, schema: Schema) -> pa.Table:
+        """Convert Polars dataframe to pyarrow table with casted schema.
+        The conversion will include generated columns.
+
+        Args:
+            df (pl.DataFrame): Dataframe to convert
+            schema (Schema): Schema to cast the dataframe to
+
+        Returns:
+            pa.Table: PyArrow table with the casted schema
+        """
+        df = schema.add_generated_columns(df)
+        pyarrow_table = df.to_arrow()
+
+        # we need to cast the incoming data to the
+        # table schema. In theory, this should automatically
+        # be casted, but it seems that metadata on fields
+        # gets removed otherwise.
+        pa_schema = schema.to_arrow()
+        return pyarrow_table.select(pa_schema.names).cast(pa_schema)
+
+    def _write_deltalake(
+        self, data: pa.Table, mode: str, predicate: Optional[str]
+    ) -> None:
+        """Write data to Delta Lake using deltalake-python.
+
+        Args:
+            data (pa.Table): PyArrow table to write to Delta Lake
+            mode (str): Write mode, either "overwrite" or "append"
+            predicate (Optional[str]): SQL predicate to filter rows for update or
+                delete operations. If None, the operation will apply to all rows.
+        """
+        dl.write_deltalake(
+            table_or_uri=self.delta_table,
+            data=data,
+            mode=mode,
+            predicate=predicate,
+            schema_mode="merge",
+        )
+
+    def update(self, df: pl.DataFrame, hyper_slice: HyperSlice) -> None:
+        """Update rows in Delta Lake based on a HyperSlice. This will overwrite data
+        in the Delta Lake specified by the HyperSlice.
+
+        Args:
+            df (pl.DataFrame): DataFrame containing the rows to update.
+            hyper_slice (HyperSlice): HyperSlice used to define rows to update.
+                If None, all rows will be updated.
+        """
+        schema = self.schema()
+        data = self._to_writable_pyarrow_table(df=df, schema=schema)
+
+        hyper_slice = schema.add_generated_filters(hyper_slice)
+        if len(hyper_slice) == 0:
+            predicate = None
+        else:
+            predicate = _dnf_to_sql(hyper_slice)
+
+        self._write_deltalake(data=data, mode="overwrite", predicate=predicate)
+
+    def append(self, df: pl.DataFrame) -> None:
+        """Append rows to Delta Lake. This will write data to the Delta Lake.
+
+        Args:
+            df (pl.DataFrame): DataFrame containing the rows to append.
+        """
+        schema = self.schema()
+        data = self._to_writable_pyarrow_table(df=df, schema=schema)
+        self._write_deltalake(data=data, mode="append", predicate=None)
+
+    def optimize(self) -> list[str]:
+        """Optimize Delta table by compacting and vacuuming
+
+        Returns:
+            list[str]: List of removed files
+        """
+        delta_table = self.delta_table
+        metrics = delta_table.optimize.compact()
+
+        vacuumed_files = delta_table.vacuum(
+            dry_run=False,
+        )
+
+        metrics["numFilesVacuumed"] = len(vacuumed_files)
+
+        return metrics
+
+    def delete(self, hyper_slice: HyperSlice) -> dict[str, Any]:
+        """Delete data from Delta table
+
+        Args:
+            hyper_slice (HyperSlice): Hyper slice to delete.
+            If None, all data will be deleted.
+
+        Returns:
+            dict[str, any]: Delete metrics.
+
+            https://docs.databricks.com/gcp/en/delta/history#operation-metrics-keys
+        """
+        if hyper_slice is None or hyper_slice == [()]:
+            predicate = None
+        else:
+            predicate = _dnf_to_sql(hyper_slice)
+
+        return self.delta_table.delete(predicate)

datazone_sdk-1.0.0/pyproject.toml

@@ -0,0 +1,42 @@
+[project]
+name = "datazone-sdk"
+version = "1.0.0"
+description = "Database and Delta storage client library for working with Delta Lake tables"
+authors = [{ name = "Team Enigma", email = "enigma@energinet.dk" }]
+requires-python = ">=3.10"
+readme = "README.md"
+
+[tool.poetry]
+packages = [{ include = "datazone" }]
+requires-poetry = ">=2.2"
+
+[tool.poetry.dependencies]
+pandas = ">=2.0.3,<3"
+polars = ">=1.33.1"
+obstore = ">=0.8.2"
+deltalake = "==1.2.1" # pin to avoid breaking changes in 1.3.0. Follow Github issue here https://github.com/delta-io/delta-rs/issues/3939
+pyarrow = ">=19.0.0"
+datamazing = ">=5.1.6"
+
+[tool.poetry.group.dev.dependencies]
+pre-commit = ">=2.20.0"
+
+[tool.poetry.group.test.dependencies]
+pytest = ">=7"
+pytest-cov = ">=3.0.0"
+mypy = ">=1.19.0"
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.isort]
+multi_line_output = 3
+line_length = 88
+include_trailing_comma = true
+
+[tool.black]
+line_length = 88
+
+[tool.mypy]
+ignore_missing_imports = true