deltacat 1.1.38__py3-none-any.whl → 2.0.0__py3-none-any.whl

deltacat/experimental/daft/daft_catalog.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+from typing import Tuple, Optional
+
+from deltacat.catalog.model.catalog import Catalog as DCCatalog
+from deltacat.catalog.model.table_definition import TableDefinition
+
+from daft.catalog import Catalog, Identifier, Table
+from daft.dataframe import DataFrame
+from daft.logical.schema import Schema
+from deltacat.constants import DEFAULT_NAMESPACE
+
+
+class DaftCatalog(Catalog):
+    """
+    Wrapper class to create a Daft catalog from a DeltaCAT catalog.
+
+    The initialization of DeltaCAT and Daft catalogs is managed in `deltacat.catalog.catalog.py`. The user
+    is just expected to initialize catalogs through the DeltaCAT public interface (init / put_catalog).
+
+    TODO (mccember) in follow up PR we need to consider how to keep the DeltaCAT Catalogs class and Daft session in sync,
+      and the user-facing entrypoint to get a Daft catalog
+
+    This class itself expects a `Catalog` and will invoke the underlying implementation
+    similar to `deltacat.catalog.delegate.py`, like:
+      catalog.impl.create_namespace(namespace, inner=catalog.inner)
+
+    We cannot route calls through the higher level catalog registry / delegate since this wrapper class is at a lower
+     layer and does not manage registering catalogs.
+    """
+
+    def __init__(self, catalog: DCCatalog, name: str):
+        """
+        Initialize given DeltaCAT catalog. This catalog is also registered with DeltaCAT (via deltacat.put_catalog) given the provided Name
+
+        :param catalog: DeltaCAT Catalog object. If None, the catalog will be fetched from `deltacat.Catalogs`
+            given the catalog name.
+
+        :param name: Name of DeltaCAT catalog. If the name is not yet registered with `deltacat.Catalogs`,
+            it will be registered upon creation to ensure that the DeltaCAT and Daft catalogs keep in sync.
+
+        :param kwargs: Additional keyword arguments passed to deltacat.get_catalog or deltacat.put_catalog,
+                       such as 'namespace' for tests.
+        """
+        self.dc_catalog = catalog
+        self._name = name
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    ###
+    # create_*
+    ###
+    def create_namespace(self, identifier: Identifier | str):
+        """Create a new namespace in the catalog."""
+        if isinstance(identifier, Identifier):
+            identifier = str(identifier)
+        self.dc_catalog.impl.create_namespace(identifier, inner=self.dc_catalog.inner)
+
+    def create_table(
+        self, identifier: Identifier | str, source: Schema | DataFrame, **kwargs
+    ) -> Table:
+        """
+        Create a DeltaCAT table via Daft catalog API
+
+        End users calling create_table through the daft table API may provide kwargs which will be plumbed through
+        to deltacat create_table. For full list of keyword arguments accepted by create_table.
+
+        Note: as of 4/22, Daft create_table does not yet support kwargs. Tracked at: https://github.com/Eventual-Inc/Daft/issues/4195
+
+        :param identifier: Daft table identifier. Sequence of strings of the format (namespace) or (namespace, table)
+            or (namespace, table, table version). If this is a string, it is a dot delimited string of the same format.
+            Identifiers can be created either like Identifier("namespace", "table", "version") OR
+                Identifier.from_str("namespace.table.version")
+
+        :param source: a TableSource, either a Daft DataFrame, Daft Schema, or str (filesystem path)
+        """
+        if isinstance(source, DataFrame):
+            return self._create_table_from_df(identifier, source)
+        elif isinstance(source, Schema):
+            return self._create_table_from_schema(identifier, source)
+        else:
+            raise Exception(
+                f"Expected table source to be Schema or DataFrame. Found: {type(source)}"
+            )
+
+    def _create_table_from_df(
+        self, ident: Identifier | str, source: DataFrame, **kwargs
+    ) -> Table:
+        """
+        Create a table from a DataFrame.
+        """
+        t = self._create_table_from_schema(ident, source.schema(), **kwargs)
+        # TODO (mccember) append data upon creation
+        return t
+
+    def _create_table_from_schema(
+        self, ident: Identifier | str, source: Schema, **kwargs
+    ) -> Table:
+        """
+        Create a table from a schema.
+        """
+        namespace, name, version = self._extract_namespace_name_version(ident)
+
+        # Convert the Daft schema to a DeltaCAT schema
+        # This is a simplified version, would need to be enhanced for production
+        deltacat_schema = self._convert_schema_to_deltacat(source)
+
+        # Create the table in DeltaCAT
+        table_def = self.dc_catalog.impl.create_table(
+            name,
+            namespace=namespace,
+            version=version,
+            schema=deltacat_schema,
+            inner=self.dc_catalog.inner,
+            **kwargs,
+        )
+
+        return DaftTable._from_obj(table_def)
+
+    ###
+    # drop_*
+    ###
+
+    def drop_namespace(self, identifier: Identifier | str):
+        raise NotImplementedError()
+
+    def drop_table(self, identifier: Identifier | str):
+        raise NotImplementedError()
+
+    ###
+    # get_*
+    ###
+
+    def get_table(self, identifier: Identifier | str, **kwargs) -> Table:
+        namespace, table, version = self._extract_namespace_name_version(identifier)
+
+        table_def = self.dc_catalog.impl.get_table(
+            table,
+            namespace=namespace,
+            table_version=version,
+            inner=self.dc_catalog.inner,
+            **kwargs,
+        )
+
+        if not table_def:
+            raise ValueError(f"Table {identifier} not found")
+
+        return DaftTable._from_obj(table_def)
+
+    ###
+    # list_*
+    ###
+
+    def list_namespaces(self, pattern: str | None = None) -> list[Identifier]:
+        raise NotImplementedError("Not implemented")
+
+    def list_tables(self, pattern: str | None = None) -> list[str]:
+        raise NotImplementedError("Not implemented")
+
+    def _extract_namespace_name_version(
+        self, ident: Identifier | str
+    ) -> Tuple[str, str, Optional[str]]:
+        """
+        Extract namespace, name,version from identifier
+
+        Returns a 3-tuple. If no namespace is provided, uses DeltaCAT defualt namespace
+        """
+        default_namespace = DEFAULT_NAMESPACE
+
+        if isinstance(ident, str):
+            ident = Identifier.from_str(ident)
+
+        if isinstance(ident, Identifier):
+            if len(ident) == 1:
+                return (default_namespace, ident[0], None)
+            elif len(ident) == 2:
+                return (ident[0], ident[1], None)
+            elif len(ident) == 3:
+                return (ident[0], ident[1], ident[2])
+            else:
+                raise ValueError(
+                    f"Expected table identifier to be in format (table) or (namespace, table)"
+                    f"or (namespace, table, version). Found: {ident}"
+                )
+
+    def _convert_schema_to_deltacat(self, schema: Schema):
+        """Convert Daft schema to DeltaCAT schema.
+        For now, just use PyArrow schema as intermediary
+        TODO look into how enhancements on schema can be propagated between Daft<=>DeltaCAT
+        """
+        from deltacat.storage.model.schema import Schema as DeltaCATSchema
+
+        return DeltaCATSchema.of(schema=schema.to_pyarrow_schema())
+
+
+class DaftTable(Table):
+    """
+    Wrapper class to create a Daft table from a DeltaCAT table
+    """
+
+    _inner: TableDefinition
+
+    _read_options = set()
+    _write_options = set()
+
+    def __init__(self, inner: TableDefinition):
+        self._inner = inner
+
+    @property
+    def name(self) -> str:
+        """Return the table name."""
+        return self._inner.table_version.table_name
+
+    @staticmethod
+    def _from_obj(obj: object) -> DaftTable:
+        """Returns a DeltaCATTable if the given object can be adapted so."""
+        if isinstance(obj, TableDefinition):
+            t = DaftTable.__new__(DaftTable)
+            t._inner = obj
+            return t
+        raise ValueError(f"Unsupported DeltaCAT table type: {type(obj)}")
+
+    def read(self, **options) -> DataFrame:
+        raise NotImplementedError("Not implemented")
+
+    def write(self, df: DataFrame | object, mode: str = "append", **options):
+        raise NotImplementedError("Not implemented")

deltacat/experimental/storage/iceberg/iceberg_scan_planner.py

@@ -0,0 +1,129 @@
+import logging
+from typing import Optional, Any, Set
+
+from pyiceberg.catalog import Catalog
+from pyiceberg.table import Table
+import deltacat.logs as logs
+
+from deltacat.storage.model.scan.push_down import Pushdown, PartitionFilter
+from deltacat.storage.model.scan.scan_plan import ScanPlan
+from deltacat.storage.model.scan.scan_task import FileScanTask, DataFile
+from deltacat.storage.util.scan_planner import ScanPlanner
+from deltacat.experimental.storage.iceberg.impl import _try_load_iceberg_table
+from deltacat.experimental.storage.iceberg.visitor import IcebergExpressionVisitor
+
+# Initialize DeltaCAT logger
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+class IcebergScanPlanner(ScanPlanner):
+    def __init__(self, catalog: Catalog):
+        self.catalog = catalog
+        self.expression_visitor = IcebergExpressionVisitor()
+
+    @classmethod
+    def _collect_filter_fields(cls, expr: Any) -> Set[str]:
+        """
+        Collects all field names referenced in the filter expression.
+
+        Args:
+            expr: The expression to analyze
+
+        Returns:
+            Set of field names referenced in the expression
+        """
+        fields = set()
+        if hasattr(expr, "field"):
+            fields.add(expr.field)
+        if hasattr(expr, "left"):
+            fields.update(cls._collect_filter_fields(expr.left))
+        if hasattr(expr, "right"):
+            fields.update(cls._collect_filter_fields(expr.right))
+        if hasattr(expr, "expr"):
+            fields.update(cls._collect_filter_fields(expr.expr))
+        if hasattr(expr, "values"):
+            for value in expr.values:
+                fields.update(cls._collect_filter_fields(value))
+        return fields
+
+    def create_scan_plan(
+        self,
+        table_name: str,
+        namespace: Optional[str] = None,
+        pushdown: Optional[Pushdown] = None,
+    ) -> ScanPlan:
+        iceberg_table = _try_load_iceberg_table(
+            self.catalog, namespace=namespace, table_name=table_name
+        )
+
+        # TODO: implement row, column predicate pushdown to Iceberg
+
+        # Get the partition spec
+        partition_spec = iceberg_table.spec()
+
+        # Check if the table is partitioned
+        is_partitioned = len(partition_spec.fields) > 0
+
+        scan = iceberg_table.scan()
+        if is_partitioned:
+            if pushdown and pushdown.partition_filter:
+                filter_fields = self._collect_filter_fields(pushdown.partition_filter)
+                logger.info(
+                    f"Pushdown partition filter is enabled, converting to Iceberg. Fields discovered in filter: {', '.join(sorted(filter_fields))}"
+                )
+                # Handle partition filter if present, DeltaCAT only supports partition-level filters right now
+                iceberg_expression = self._convert_partition_filter(
+                    iceberg_table, pushdown.partition_filter
+                )
+                scan = scan.filter(iceberg_expression)
+
+        file_scan_tasks = []
+        for scan_task in scan.plan_files():
+            file_scan_tasks.append(FileScanTask([DataFile(scan_task.file.file_path)]))
+        return ScanPlan(file_scan_tasks)
+
+    @classmethod
+    def _validate_partition_references(
+        cls, expr: Any, partition_cols: Set[str]
+    ) -> None:
+        """
+        Validates that the expression only references partition columns.
+
+        Args:
+            expr: The expression to validate
+            partition_cols: Set of valid partition column names
+
+        Raises:
+            ValueError: If the expression references a non-partition column
+        """
+        if hasattr(expr, "field"):  # Reference type expression
+            if expr.field not in partition_cols:
+                raise ValueError(
+                    f"Filter references non-partition column: {expr.field}. "
+                    f"Partition columns are: {partition_cols}"
+                )
+        # Recursively validate nested expressions
+        if hasattr(expr, "left"):
+            cls._validate_partition_references(expr.left, partition_cols)
+        if hasattr(expr, "right"):
+            cls._validate_partition_references(expr.right, partition_cols)
+        if hasattr(expr, "expr"):
+            cls._validate_partition_references(expr.expr, partition_cols)
+        if hasattr(expr, "values"):
+            for value in expr.values:
+                cls._validate_partition_references(value, partition_cols)
+
+    def _convert_partition_filter(
+        self, table: Table, partition_filter: PartitionFilter
+    ):
+        """
+        Convert DeltaCAT partition filter to PyIceberg expression,
+        validating that only partition columns are referenced.
+        """
+        partition_cols = set(field.name for field in table.spec().fields)
+
+        # Validate before converting
+        self._validate_partition_references(partition_filter, partition_cols)
+
+        # Convert to PyIceberg expression
+        return self.expression_visitor.visit(partition_filter)