planar 0.8.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

planar/data/dataset.py

@@ -0,0 +1,272 @@
+"""PlanarDataset implementation for working with Ducklake tables."""
+
+import asyncio
+from typing import Literal, Self
+
+import ibis
+import polars as pl
+import pyarrow as pa
+from ibis.backends.duckdb import Backend as DuckDBBackend
+from ibis.common.exceptions import TableNotFound
+from pydantic import BaseModel
+
+from planar.config import PlanarConfig
+from planar.files.storage.config import LocalDirectoryConfig, S3Config
+from planar.logging import get_logger
+from planar.session import get_config
+
+from .exceptions import DataError, DatasetAlreadyExistsError, DatasetNotFoundError
+
+logger = get_logger(__name__)
+
+
+class PlanarDataset(BaseModel):
+    """Reference to a Ducklake table.
+
+    This class provides a simple interface for working with datasets in Ducklake,
+    handling creation, reading, writing, and deletion of tabular data.
+    """
+
+    # TODO: Add support for schema name (ie. namespace)
+    name: str  # Table name in Ducklake
+    # TODO: Add snapshot version: no version = latest, otherwise time travel on read operations
+    # TODO: Add partition support? A Dataset representation could be a table with a partition column
+
+    model_config = {"arbitrary_types_allowed": True}
+    # TODO: Add serialization metadata to make clear this is a dataset reference
+    # like EntityField.
+
+    @classmethod
+    async def create(cls, name: str, if_not_exists: bool = True) -> Self:
+        """Create a dataset reference.
+
+        Note: The actual table is created when data is first written to avoid
+        DuckDB's requirement that tables have at least one column.
+
+        Args:
+            name: Name of the dataset
+            if_not_exists: If True, don't raise error if dataset exists. default: True
+            catalog: Catalog name in Ducklake
+
+        Returns:
+            PlanarDataset instance
+
+        Raises:
+            DatasetAlreadyExistsError: If dataset exists and if_not_exists=False
+        """
+        dataset = cls(name=name)
+
+        # Check if dataset already exists
+        if await dataset.exists():
+            if not if_not_exists:
+                raise DatasetAlreadyExistsError(f"Dataset {name} already exists")
+            logger.debug("dataset already exists", dataset_name=name)
+        else:
+            logger.debug("dataset reference created", dataset_name=name)
+
+        return dataset
+
+    async def exists(self) -> bool:
+        """Check if the dataset exists in Ducklake."""
+        con = await self._get_connection()
+        try:
+            # TODO: Query for the table name directly
+            tables = await asyncio.to_thread(con.list_tables)
+            return self.name in tables
+        except Exception as e:
+            logger.warning("failed to check dataset existence", error=str(e))
+            return False
+
+    async def write(
+        self,
+        data: pl.DataFrame | ibis.Table | list | dict,
+        mode: Literal["overwrite", "append"] = "append",
+    ) -> None:
+        """Write data to the dataset.
+
+        Args:
+            data: Data to write (Polars DataFrame, PyArrow Table, or Ibis expression)
+            mode: Write mode - "append" or "overwrite"
+        """
+        con = await self._get_connection()
+        overwrite = mode == "overwrite"
+
+        try:
+            if not await self.exists():
+                await asyncio.to_thread(
+                    con.create_table, self.name, data, overwrite=overwrite
+                )
+            else:
+                # TODO: Explore if workflow context can be used to set metadata
+                # on the snapshot version for lineage
+                if isinstance(data, pl.DataFrame):
+                    await asyncio.to_thread(
+                        con.insert,
+                        self.name,
+                        ibis.memtable(data),
+                        overwrite=overwrite,
+                    )
+                else:
+                    await asyncio.to_thread(
+                        con.insert, self.name, data, overwrite=overwrite
+                    )
+
+            logger.debug(
+                "wrote data to dataset",
+                dataset_name=self.name,
+                mode=mode,
+            )
+        except Exception as e:
+            raise DataError(f"Failed to write data: {e}") from e
+
+    async def read(
+        self,
+        columns: list[str] | None = None,
+        limit: int | None = None,
+    ) -> ibis.Table:
+        """Read data as an Ibis table expression.
+
+        Args:
+            columns: Optional list of columns to select
+            limit: Optional row limit
+
+        Returns:
+            Ibis table expression that can be further filtered using Ibis methods
+        """
+        con = await self._get_connection()
+
+        try:
+            table = await asyncio.to_thread(con.table, self.name)
+
+            if columns:
+                table = table.select(columns)
+
+            if limit:
+                table = table.limit(limit)
+
+            return table
+        except TableNotFound as e:
+            raise DatasetNotFoundError(f"Dataset {self.name} not found") from e
+        except Exception as e:
+            raise DataError(f"Failed to read data: {e}") from e
+
+    async def to_polars(self) -> pl.DataFrame:
+        """Read entire dataset as Polars DataFrame."""
+        table = await self.read()
+        return await asyncio.to_thread(table.to_polars)
+
+    async def to_pyarrow(self) -> pa.Table:
+        """Read entire dataset as PyArrow Table."""
+        table = await self.read()
+        return await asyncio.to_thread(table.to_pyarrow)
+
+    async def delete(self) -> None:
+        """Delete the dataset."""
+        con = await self._get_connection()
+        try:
+            await asyncio.to_thread(con.drop_table, self.name, force=True)
+            logger.info("deleted dataset", dataset_name=self.name)
+        except Exception as e:
+            raise DataError(f"Failed to delete dataset: {e}") from e
+
+    async def _get_connection(self) -> DuckDBBackend:
+        """Get Ibis connection to Ducklake."""
+        config = get_config()
+
+        if not config.data:
+            raise DataError(
+                "Data configuration not found. Please configure 'data' in your planar.yaml"
+            )
+
+        # TODO: Add cached connection pooling or memoize the connection
+        return await self._create_connection(config)
+
+    async def _create_connection(self, config: PlanarConfig) -> DuckDBBackend:
+        """Create Ibis DuckDB connection with Ducklake."""
+        data_config = config.data
+        if not data_config:
+            raise DataError("Data configuration not found")
+
+        # Connect to DuckDB with Ducklake extension
+        con = await asyncio.to_thread(ibis.duckdb.connect, extensions=["ducklake"])
+
+        # Build Ducklake connection string based on catalog type
+        catalog_config = data_config.catalog
+
+        if catalog_config.type == "duckdb":
+            metadata_path = catalog_config.path
+        elif catalog_config.type == "postgres":
+            # Use connection components to build postgres connection string
+            pg = catalog_config
+            metadata_path = f"postgres:dbname={pg.db}"
+            if pg.host:
+                metadata_path += f" host={pg.host}"
+            if pg.port:
+                metadata_path += f" port={pg.port}"
+            if pg.user:
+                metadata_path += f" user={pg.user}"
+            if pg.password:
+                metadata_path += f" password={pg.password}"
+        elif catalog_config.type == "sqlite":
+            metadata_path = f"sqlite:{catalog_config.path}"
+        else:
+            raise ValueError(f"Unsupported catalog type: {catalog_config.type}")
+
+        # Build ATTACH statement
+        attach_sql = f"ATTACH 'ducklake:{metadata_path}'"
+
+        # Add data path from storage config
+        storage = data_config.storage
+        if isinstance(storage, LocalDirectoryConfig):
+            data_path = storage.directory
+        elif isinstance(storage, S3Config):
+            data_path = f"s3://{storage.bucket_name}/"
+        else:
+            # Generic fallback
+            data_path = getattr(storage, "path", None) or getattr(
+                storage, "directory", "."
+            )
+
+        attach_sql += f" (DATA_PATH '{data_path}')"
+
+        # Attach to Ducklake
+        try:
+            await asyncio.to_thread(con.raw_sql, attach_sql)
+        except Exception as e:
+            raise DataError(f"Failed to attach to Ducklake: {e}") from e
+
+        # Check available catalogs (what ibis calls schemas) and use the correct one
+        try:
+            catalogs = await asyncio.to_thread(con.list_catalogs)
+            logger.debug("available catalogs", catalogs=catalogs)
+
+            # Find the ducklake catalog (it might have a different name pattern)
+            ducklake_catalog = None
+            for cat in catalogs:
+                if "ducklake" in cat.lower() or cat == data_config.catalog_name:
+                    ducklake_catalog = cat
+                    break
+
+            if ducklake_catalog:
+                if not ducklake_catalog.replace("_", "").replace("-", "").isalnum():
+                    raise DataError(f"Invalid catalog name format: {ducklake_catalog}")
+                await asyncio.to_thread(con.raw_sql, f"USE {ducklake_catalog}")
+                logger.debug("using catalog", catalog=ducklake_catalog)
+            else:
+                catalog_name = data_config.catalog_name
+                if not catalog_name.replace("_", "").replace("-", "").isalnum():
+                    raise DataError(f"Invalid catalog name format: {catalog_name}")
+                await asyncio.to_thread(
+                    con.raw_sql,
+                    f"CREATE SCHEMA IF NOT EXISTS {catalog_name}",
+                )
+                await asyncio.to_thread(con.raw_sql, f"USE {catalog_name}")
+                logger.debug(
+                    "created and using catalog", catalog=data_config.catalog_name
+                )
+
+        except Exception as e:
+            logger.warning("failed to set catalog, using default", error=str(e))
+            # Continue without setting catalog - will use qualified names
+
+        return con

planar/data/exceptions.py

@@ -0,0 +1,19 @@
+"""Exceptions for the Planar data module."""
+
+
+class DataError(Exception):
+    """Base exception for data-related errors."""
+
+    pass
+
+
+class DatasetNotFoundError(DataError):
+    """Raised when a dataset is not found."""
+
+    pass
+
+
+class DatasetAlreadyExistsError(DataError):
+    """Raised when trying to create a dataset that already exists."""
+
+    pass

planar/data/test_dataset.py

@@ -0,0 +1,354 @@
+"""Tests for PlanarDataset."""
+
+import polars as pl
+import pyarrow as pa
+import pytest
+from ibis import literal
+
+from planar import PlanarApp
+from planar.data import PlanarDataset
+from planar.data.config import DataConfig, SQLiteCatalogConfig
+from planar.data.exceptions import (
+    DataError,
+    DatasetAlreadyExistsError,
+    DatasetNotFoundError,
+)
+from planar.files.storage.config import LocalDirectoryConfig
+from planar.workflows import step
+
+
+@pytest.fixture
+def data_config(tmp_path):
+    """Create a test data configuration."""
+    data_dir = tmp_path / "data"
+    data_dir.mkdir(exist_ok=True)
+
+    catalog_path = data_dir / "test.sqlite"
+    storage_path = data_dir / "ducklake_files"
+    storage_path.mkdir(exist_ok=True)
+
+    return DataConfig(
+        catalog=SQLiteCatalogConfig(type="sqlite", path=str(catalog_path)),
+        storage=LocalDirectoryConfig(backend="localdir", directory=str(storage_path)),
+    )
+
+
+@pytest.fixture(name="app")
+def app_fixture(data_config):
+    """Create a PlanarApp with data configuration."""
+    app = PlanarApp()
+    # Add data config to the app's config
+    app.config.data = data_config
+    return app
+
+
+@pytest.mark.asyncio
+async def test_dataset_create(client):
+    """Test creating a dataset reference."""
+    dataset = await PlanarDataset.create("test_table")
+    assert dataset.name == "test_table"
+
+    # Dataset reference exists but table isn't created until first write
+    assert not await dataset.exists()
+
+    # Write some data to actually create the table
+    df = pl.DataFrame({"id": [1], "name": ["test"]})
+    await dataset.write(df, mode="overwrite")
+
+    # Now it should exist
+    assert await dataset.exists()
+
+    # Cleanup
+    await dataset.delete()
+
+
+@pytest.mark.asyncio
+async def test_dataset_create_if_not_exists(client):
+    """Test creating a dataset with if_not_exists behavior."""
+    # Create dataset and write data to make it exist
+    dataset1 = await PlanarDataset.create("test_table")
+    df = pl.DataFrame({"id": [1], "name": ["test"]})
+    await dataset1.write(df, mode="overwrite")
+
+    # Create again with if_not_exists=True (default) - should not raise
+    dataset2 = await PlanarDataset.create("test_table", if_not_exists=True)
+    assert dataset2.name == dataset1.name
+
+    # Create again with if_not_exists=False - should raise
+    with pytest.raises(DatasetAlreadyExistsError):
+        await PlanarDataset.create("test_table", if_not_exists=False)
+
+    # Cleanup
+    await dataset1.delete()
+
+
+@pytest.mark.asyncio
+async def test_dataset_write_and_read_polars(client):
+    """Test writing and reading data with Polars."""
+    dataset = await PlanarDataset.create("test_polars")
+
+    # Create test data
+    df = pl.DataFrame(
+        {
+            "id": [1, 2, 3],
+            "name": ["Alice", "Bob", "Charlie"],
+            "amount": [100.5, 200.0, 150.75],
+        }
+    )
+
+    # Write data
+    await dataset.write(df, mode="overwrite")
+
+    # Read data back
+    result = await dataset.to_polars()
+
+    # Verify
+    assert result.shape == df.shape
+    assert set(result.columns) == set(df.columns)
+    assert result["id"].to_list() == [1, 2, 3]
+    assert result["name"].to_list() == ["Alice", "Bob", "Charlie"]
+
+    # Cleanup
+    await dataset.delete()
+
+
+@pytest.mark.asyncio
+async def test_dataset_write_and_read_pyarrow(client):
+    """Test writing and reading data with PyArrow."""
+    dataset = await PlanarDataset.create("test_pyarrow")
+
+    # Create test data
+    table = pa.table(
+        {
+            "id": [1, 2, 3],
+            "name": ["Alice", "Bob", "Charlie"],
+            "amount": [100.5, 200.0, 150.75],
+        }
+    )
+
+    # Write data
+    await dataset.write(table, mode="overwrite")
+
+    # Read data back
+    result = await dataset.to_pyarrow()
+
+    # Verify
+    assert result.num_rows == table.num_rows
+    assert result.column_names == table.column_names
+
+    # Cleanup
+    await dataset.delete()
+
+
+@pytest.mark.asyncio
+async def test_dataset_append_mode(client):
+    """Test appending data to a dataset."""
+    dataset = await PlanarDataset.create("test_append")
+
+    # Write initial data
+    df1 = pl.DataFrame({"id": [1, 2], "value": ["a", "b"]})
+    await dataset.write(df1, mode="overwrite")
+
+    # Append more data
+    df2 = pl.DataFrame({"id": [3, 4], "value": ["c", "d"]})
+    await dataset.write(df2, mode="append")
+
+    result = await dataset.to_polars()
+
+    # Verify
+    assert len(result) == 4
+    assert set(result["id"].to_list()) == {1, 2, 3, 4}
+    assert set(result["value"].to_list()) == {"a", "b", "c", "d"}
+
+    # Cleanup
+    await dataset.delete()
+
+
+@pytest.mark.asyncio
+async def test_dataset_overwrite_replaces_existing(client):
+    """Overwrite should replace existing rows completely."""
+    dataset = await PlanarDataset.create("test_overwrite")
+
+    df1 = pl.DataFrame({"id": [1, 2], "value": ["a", "b"]})
+    await dataset.write(df1, mode="overwrite")
+    result1 = await dataset.to_polars()
+    assert result1.shape == (2, 2)
+
+    df2 = pl.DataFrame({"id": [3], "value": ["c"]})
+    await dataset.write(df2, mode="overwrite")
+    result2 = await dataset.to_polars()
+    assert result2.shape == (1, 2)
+    assert result2["id"].to_list() == [3]
+    assert result2["value"].to_list() == ["c"]
+
+    await dataset.delete()
+
+
+@pytest.mark.asyncio
+async def test_dataset_read_with_filter(client):
+    """Test reading data with Ibis filtering."""
+    dataset = await PlanarDataset.create("test_filter")
+
+    # Write test data
+    df = pl.DataFrame({"id": range(1, 11), "value": range(10, 101, 10)})
+    await dataset.write(df, mode="overwrite")
+
+    table = await dataset.read()
+    filtered_table = table.filter(table.value > literal(50))
+    filtered_df = filtered_table.to_polars()
+
+    assert len(filtered_df) == 5
+    assert all(v > 50 for v in filtered_df["value"].to_list())
+
+    # Cleanup
+    await dataset.delete()
+
+
+@pytest.mark.asyncio
+async def test_dataset_read_with_columns_and_limit(client):
+    """Test reading specific columns with limit."""
+    dataset = await PlanarDataset.create("test_select")
+
+    # Write test data
+    df = pl.DataFrame(
+        {
+            "id": range(1, 11),
+            "name": [f"user_{i}" for i in range(1, 11)],
+            "value": range(10, 101, 10),
+        }
+    )
+    await dataset.write(df, mode="overwrite")
+
+    # Read specific columns with limit
+    table = await dataset.read(columns=["id", "name"], limit=5)
+    result_df = table.to_polars()
+
+    # Verify
+    assert len(result_df) == 5
+    assert set(result_df.columns) == {"id", "name"}
+    assert "value" not in result_df.columns
+
+    # Cleanup
+    await dataset.delete()
+
+
+@pytest.mark.asyncio
+async def test_dataset_not_found(client):
+    """Test reading from non-existent dataset."""
+    dataset = PlanarDataset(name="nonexistent")
+
+    # Check exists returns False
+    assert not await dataset.exists()
+
+    # Try to read - should raise
+    with pytest.raises(DatasetNotFoundError):
+        await dataset.read()
+
+
+@pytest.mark.asyncio
+async def test_dataset_delete(client):
+    """Test deleting a dataset."""
+    dataset = await PlanarDataset.create("test_delete")
+
+    # Write some data
+    df = pl.DataFrame({"id": [1, 2, 3]})
+    await dataset.write(df)
+
+    # Verify it exists
+    assert await dataset.exists()
+
+    # Delete it
+    await dataset.delete()
+
+    # Verify it's gone
+    assert not await dataset.exists()
+
+
+@pytest.mark.asyncio
+async def test_dataset_write_list_of_dicts(client):
+    """Write list-of-dicts input and read back with Polars."""
+    dataset = await PlanarDataset.create("test_list_of_dicts")
+
+    rows = [{"id": 1, "name": "a"}, {"id": 2, "name": "b"}]
+    await dataset.write(rows, mode="overwrite")
+
+    result = await dataset.to_polars()
+    assert set(result.columns) == {"id", "name"}
+    assert sorted(result["id"].to_list()) == [1, 2]
+
+    await dataset.delete()
+
+
+@pytest.mark.asyncio
+async def test_dataset_write_dict_of_lists(client):
+    """Write dict-of-lists input and read back with Polars."""
+    dataset = await PlanarDataset.create("test_dict_of_lists")
+
+    data = {"id": [1, 2], "name": ["a", "b"]}
+    await dataset.write(data, mode="overwrite")
+
+    result = await dataset.to_polars()
+    assert result.shape == (2, 2)
+    assert set(result["name"].to_list()) == {"a", "b"}
+
+    await dataset.delete()
+
+
+@pytest.mark.asyncio
+async def test_dataset_workflow_serialization(client):
+    """Test that PlanarDataset can be used as workflow input/output."""
+
+    @step()
+    async def create_data() -> PlanarDataset:
+        """Create a dataset with sample data."""
+        dataset = await PlanarDataset.create("workflow_data")
+
+        df = pl.DataFrame(
+            {"product": ["A", "B", "C", "D"], "sales": [100, 200, 150, 300]}
+        )
+        await dataset.write(df, mode="overwrite")
+
+        return dataset
+
+    @step()
+    async def analyze_data(dataset: PlanarDataset) -> float:
+        """Analyze the dataset and return total sales."""
+        df = await dataset.to_polars()
+        return float(df["sales"].sum())
+
+    # Test basic workflow functionality without API
+    dataset = await create_data()
+    total = await analyze_data(dataset)
+
+    # Verify results
+    assert total == 750.0  # Sum of [100, 200, 150, 300]
+
+    # Cleanup
+    await dataset.delete()
+
+
+@pytest.mark.asyncio
+async def test_no_data_config_error(client):
+    """Test error when data config is not set."""
+    # Remove data config
+    client.app.config.data = None
+
+    dataset = PlanarDataset(name="test")
+
+    with pytest.raises(DataError, match="Data configuration not found"):
+        await dataset._get_connection()
+
+
+@pytest.mark.asyncio
+async def test_write_with_invalid_input_raises(client):
+    """Unknown input types to write() should raise a DataError."""
+
+    class Foo:
+        pass
+
+    dataset = await PlanarDataset.create("test_invalid_input")
+
+    with pytest.raises(DataError):
+        await dataset.write(Foo(), mode="overwrite")  # type: ignore
+
+    await dataset.delete()

planar/dependencies.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import sys
+from importlib import import_module
+from typing import Mapping, Tuple
+
+
+# mapping: public_name -> (relative_submodule, attribute_in_submodule)
+# This is a PEP 562 compliant way to lazily import modules
+# which is a way to avoid circular dependencies in __init__.py.
+def lazy_exports(module_name: str, mapping: Mapping[str, Tuple[str, str]]) -> None:
+    mod = sys.modules[module_name]
+    mod.__all__ = list(mapping.keys())  # type: ignore
+
+    def __getattr__(name: str):
+        try:
+            submod, attr = mapping[name]
+        except KeyError:
+            raise AttributeError(
+                f"module {module_name!r} has no attribute {name!r}"
+            ) from None
+        obj = getattr(import_module(submod, module_name), attr)
+        setattr(mod, name, obj)  # cache
+        return obj
+
+    def __dir__():
+        return sorted(set(mod.__dict__.keys()) | set(mod.__all__))
+
+    mod.__getattr__ = __getattr__  # PEP 562
+    mod.__dir__ = __dir__