deltacat 2.0.0b9__py3-none-any.whl → 2.0.0b11__py3-none-any.whl

deltacat/tests/daft/test_model.py

@@ -0,0 +1,97 @@
+import pytest
+import pyarrow as pa
+from daft import DataType, TimeUnit
+from daft.logical.schema import Field as DaftField
+
+from deltacat.storage.model.transform import IdentityTransform
+from deltacat.storage.model.partition import PartitionKey
+from deltacat.utils.daft import DaftFieldMapper, DaftPartitionKeyMapper
+
+from deltacat.storage.model.schema import Field, Schema
+
+
+class TestDaftFieldMapper:
+    def test_field_mapper_basic_types(self):
+        """Test mapping basic data types between Daft and PyArrow fields"""
+        test_cases = [
+            (DataType.int32(), pa.int32()),
+            (DataType.int64(), pa.int64()),
+            (DataType.float32(), pa.float32()),
+            (DataType.float64(), pa.float64()),
+            (DataType.string(), pa.large_string()),
+            (DataType.bool(), pa.bool_()),
+            (DataType.binary(), pa.large_binary()),
+            (DataType.date(), pa.date32()),
+            (DataType.timestamp(TimeUnit.ns()), pa.timestamp("ns")),
+        ]
+
+        for daft_type, pa_type in test_cases:
+            # Create test fields
+            daft_field = DaftField.create(
+                name="test_field",
+                dtype=daft_type,
+            )
+
+            # Daft to PyArrow
+            pa_field = DaftFieldMapper.map(daft_field)
+            assert pa_field is not None
+            assert pa_field.name == "test_field"
+            assert pa_field.type == pa_type  # type: ignore
+            assert pa_field.nullable is True
+
+            # PyArrow to Daft
+            daft_field_back = DaftFieldMapper.unmap(pa_field)
+            assert daft_field_back is not None
+            assert daft_field_back.name == daft_field.name
+            assert daft_field_back.dtype == daft_field.dtype
+
+
+class TestDaftPartitionKeyMapper:
+    def test_unmap(self):
+        """
+        Test unmap method of DaftPartitionKeyMapper when obj is not None, schema is provided,
+        len(obj.key) is 1, and dc_field is found in the schema.
+
+        This test verifies that the method correctly converts a PartitionKey to a DaftPartitionField
+        when all conditions are met and the field exists in the schema.
+        """
+        # Create a mock schema
+        schema = Schema.of(schema=[Field.of(pa.field("test_field", pa.int32()))])
+        # Create a PartitionKey object
+        partition_key = PartitionKey(
+            key=["test_field"], transform=IdentityTransform(), name="partition_field"
+        )
+
+        result = DaftPartitionKeyMapper.unmap(obj=partition_key, schema=schema)
+        assert result is not None
+        assert result.field.name() == "partition_field"
+        assert DataType._from_pydatatype(result.field.dtype()) == DataType.int32()
+
+    def test_unmap_no_field_locator(self):
+        schema = Schema.of(schema=[Field.of(pa.field("test_field", pa.int32()))])
+        partition_key = PartitionKey(key=[], name="partition_field")
+
+        with pytest.raises(ValueError) as excinfo:
+            DaftPartitionKeyMapper.unmap(partition_key, schema)
+
+        assert "At least 1 PartitionKey FieldLocator is expected" in str(excinfo.value)
+
+    def test_unmap_partition_key_not_found(self):
+        schema = Schema.of(schema=[Field.of(pa.field("test_field", pa.int32()))])
+        partition_key = PartitionKey(
+            key=["test_field_2"], transform=IdentityTransform(), name="partition_field"
+        )
+
+        with pytest.raises(KeyError) as excinfo:
+            DaftPartitionKeyMapper.unmap(partition_key, schema)
+
+        assert "Column test_field_2 does not exist in schema" in str(excinfo.value)
+
+    def test_unmap_partition_name_not_defined(self):
+        schema = Schema.of(schema=[Field.of(pa.field("test_field", pa.int32()))])
+        partition_key = PartitionKey(key=[])
+
+        with pytest.raises(ValueError) as excinfo:
+            DaftPartitionKeyMapper.unmap(partition_key, schema)
+
+        assert "Name is required for PartitionKey conversion" in str(excinfo.value)

deltacat/tests/experimental/catalog/iceberg/test_iceberg_catalog.py

@@ -0,0 +1,71 @@
+import tempfile
+import shutil
+import uuid
+import deltacat
+import pytest
+from deltacat import Field, Schema
+from pyiceberg.catalog import CatalogType
+
+import pyarrow as pa
+
+from deltacat.experimental.catalog.iceberg import IcebergCatalogConfig
+
+
+@pytest.fixture
+def schema_a():
+    return Schema.of(
+        [
+            Field.of(
+                field=pa.field("col1", pa.int32(), nullable=False),
+                field_id=1,
+                is_merge_key=True,
+            )
+        ]
+    )
+
+
+class TestIcebergCatalogInitialization:
+    temp_dir = None
+
+    @classmethod
+    def setup_class(cls):
+        cls.temp_dir = tempfile.mkdtemp()
+
+    @classmethod
+    def teardown_class(cls):
+        shutil.rmtree(cls.temp_dir)
+
+    def test_iceberg_catalog_and_table_create(self, schema_a):
+
+        # Register a random catalog name to avoid concurrent test conflicts
+        catalog_name = str(uuid.uuid4())
+
+        config = IcebergCatalogConfig(
+            type=CatalogType.SQL,
+            properties={
+                "warehouse": self.temp_dir,
+                "uri": f"sqlite:////{self.temp_dir}/sql-catalog.db",
+            },
+        )
+
+        # Initialize with the PyIceberg catalog
+        catalog = deltacat.IcebergCatalog.from_config(config)
+        deltacat.init(
+            {catalog_name: catalog},
+            force=True,
+        )
+
+        table_def = deltacat.create_table(
+            "test_table", catalog=catalog_name, schema=schema_a
+        )
+
+        # Fetch table we just created
+        fetched_table_def = deltacat.get_table("test_table", catalog=catalog_name)
+        assert table_def.table_version == fetched_table_def.table_version
+
+        # For now, just check that we created a table version with an equivalent schema
+        assert table_def.table_version.schema.equivalent_to(schema_a)
+
+        # Sanity check that list namespaces works
+        namespaces = deltacat.list_namespaces(catalog=catalog_name).all_items()
+        assert table_def.table.namespace in [n.namespace for n in namespaces]

deltacat/tests/experimental/daft/test_deltacat_daft_integration.py

@@ -0,0 +1,136 @@
+import daft
+from daft import Table, Identifier
+import pytest
+import uuid
+
+from deltacat.catalog import Catalog as DeltaCATCatalog
+from deltacat.catalog import CatalogProperties
+from deltacat.experimental.daft.daft_catalog import DaftCatalog
+import shutil
+import tempfile
+
+from deltacat.experimental.catalog.iceberg import impl as IcebergCatalog
+from deltacat.experimental.catalog.iceberg import IcebergCatalogConfig
+
+from pyiceberg.catalog import CatalogType
+
+
+class TestCatalogIntegration:
+    @classmethod
+    def setup_method(cls):
+        cls.tmpdir = tempfile.mkdtemp()
+
+    @classmethod
+    def teardown_method(cls):
+        shutil.rmtree(cls.tmpdir)
+
+    def test_create_table(self):
+        """Demonstrate DeltaCAT-Daft integration."""
+        # Create a DeltaCAT catalog
+        catalog_props = CatalogProperties(root=self.tmpdir)
+        dc_catalog = DeltaCATCatalog(catalog_props)
+
+        # Use a random catalog name to prevent namespacing conflicts with other tests
+        # Convert the DeltaCAT catalog to a Daft catalog
+        catalog_name = f"deltacat_{uuid.uuid4().hex[:8]}"
+
+        daft_catalog = DaftCatalog(catalog=dc_catalog, name=catalog_name)
+
+        # Register the catalog with Daft's catalog system
+        daft.attach_catalog(daft_catalog, catalog_name)
+
+        # Create a sample DataFrame
+        df = daft.from_pydict({"id": [1, 2, 3], "value": ["a", "b", "c"]})
+        # Create then get table
+        daft_catalog.create_table(Identifier("example_table"), df)
+        table: Table = daft_catalog.get_table(Identifier("example_table"))
+        assert table.name == "example_table"
+
+    def test_get_table(self):
+        """Test getting a table from the DeltaCAT-Daft catalog."""
+        # Create a DeltaCAT catalog using the existing tmpdir
+        catalog_props = CatalogProperties(root=self.tmpdir)
+        dc_catalog = DeltaCATCatalog(catalog_props)
+
+        # Convert to DaftCatalog and attach to Daft
+        catalog_name = f"deltacat_{uuid.uuid4().hex[:8]}"
+        daft_catalog = DaftCatalog(dc_catalog, catalog_name)
+        daft.attach_catalog(daft_catalog, catalog_name)
+
+        # Create a sample DataFrame and table
+        df = daft.from_pydict({"id": [1, 2, 3], "value": ["a", "b", "c"]})
+        table_name = "test_get_table"
+        daft_catalog.create_table(Identifier(table_name), df)
+
+        # Get the table using different forms of identifiers
+        table2 = daft_catalog.get_table(Identifier(table_name))
+        assert table2 is not None
+        assert table2.name == table_name
+
+        # 3. With namespace. DeltaCAT used the default namespace since it was not provided
+        table3 = daft_catalog.get_table(Identifier("default", table_name))
+        assert table3 is not None
+        assert table3.name == table_name
+
+        # Test non-existent table raises an appropriate error
+        with pytest.raises(ValueError, match="Table nonexistent_table not found"):
+            daft_catalog.get_table(Identifier("nonexistent_table"))
+
+
+class TestIcebergCatalogIntegration:
+    @classmethod
+    def setup_method(cls):
+        cls.tmpdir = tempfile.mkdtemp()
+
+    @classmethod
+    def teardown_method(cls):
+        shutil.rmtree(cls.tmpdir)
+
+    def test_iceberg_catalog_integration(self):
+        # Create a unique warehouse path for this test
+        warehouse_path = self.tmpdir
+
+        # Configure an Iceberg catalog with the warehouse path
+        config = IcebergCatalogConfig(
+            type=CatalogType.SQL,
+            properties={
+                "warehouse": warehouse_path,
+                "uri": f"sqlite:////{warehouse_path}/sql-catalog.db",
+            },
+        )
+        dc_catalog = IcebergCatalog.from_config(config)
+
+        # Convert the DeltaCAT catalog to a Daft catalog
+        catalog_name = f"deltacat_iceberg_{uuid.uuid4().hex[:8]}"
+        daft_catalog = DaftCatalog(dc_catalog, catalog_name)
+        daft.attach_catalog(daft_catalog, catalog_name)
+
+        # Create a sample DataFrame
+        df = daft.from_pydict({"id": [1, 2, 3], "value": ["a", "b", "c"]})
+
+        # Create a table with the Daft catalog
+        table_name = "example_table"
+        namespace = "example_namespace"
+        daft_catalog.create_table(Identifier(namespace, table_name), df)
+
+        # Query that Iceberg table exists using PyIceberg
+        iceberg_catalog = dc_catalog.inner
+
+        # Verify the table exists in the Iceberg catalog
+        tables = iceberg_catalog.list_tables(namespace)
+
+        assert any(
+            t[0] == namespace and t[1] == table_name for t in tables
+        ), f"Table {table_name} not found in Iceberg catalog"
+
+        # Load the table from Iceberg catalog and verify its properties
+        iceberg_table = iceberg_catalog.load_table(f"{namespace}.{table_name}")
+
+        # Check that the schema matches our DataFrame
+        schema = iceberg_table.schema()
+        assert (
+            schema.find_field("id") is not None
+        ), "Field 'id' not fcound in table schema"
+        assert (
+            schema.find_field("value") is not None
+        ), "Field 'value' not found in table schema"

deltacat/tests/{storage → experimental/storage}/rivulet/conftest.py

@@ -3,9 +3,9 @@ import io
 import pytest
 from faker import Faker
 
-from deltacat.storage.rivulet.schema.datatype import Datatype
-from deltacat.storage.rivulet.mvp.Table import MvpTable
-from deltacat.storage.rivulet.schema.schema import Schema
+from deltacat.experimental.storage.rivulet.schema.datatype import Datatype
+from deltacat.experimental.storage.rivulet.mvp.Table import MvpTable
+from deltacat.experimental.storage.rivulet.schema.schema import Schema
 import random
 import string
 from PIL import Image

deltacat/tests/{storage → experimental/storage}/rivulet/fs/test_file_location_provider.py

@@ -2,8 +2,9 @@ import pytest
 
 import pyarrow as pa
 import pyarrow.parquet as pq
-from deltacat import Datatype, Dataset
-from deltacat.storage.rivulet import Schema, Field
+from deltacat.experimental.storage.rivulet.schema.datatype import Datatype
+from deltacat.experimental.storage.rivulet.dataset import Dataset
+from deltacat.experimental.storage.rivulet import Schema, Field
 from deltacat.utils.metafile_locator import _find_partition_path
 
 

deltacat/tests/experimental/storage/rivulet/reader/query_expression.py

@@ -0,0 +1,80 @@
+import pytest
+from deltacat.experimental.storage.rivulet.reader.query_expression import (
+    QueryExpression,
+)
+from deltacat.experimental.storage.rivulet.shard.range_shard import RangeShard
+
+
+@pytest.fixture
+def sample_range_shard():
+    return RangeShard(min_key=5, max_key=15)
+
+
+@pytest.fixture
+def sample_string_shard():
+    return RangeShard(min_key="apple", max_key="zebra")
+
+
+def test_with_key():
+    query = QueryExpression[int]()
+    query.with_key(5)
+    assert query.min_key == 5
+    assert query.max_key == 5
+    with pytest.raises(ValueError):
+        query.with_key(10)
+
+
+def test_with_range():
+    query = QueryExpression[int]()
+    query.with_range(10, 5)
+    assert query.min_key == 5
+    assert query.max_key == 10
+    with pytest.raises(ValueError):
+        query.with_range(20, 25)
+
+
+def test_matches_query():
+    query = QueryExpression[int]()
+    assert query.matches_query(5)
+    assert query.matches_query(-999)
+    query.with_range(10, 20)
+    assert query.matches_query(15)
+    assert not query.matches_query(25)
+    assert not query.matches_query(5)
+
+
+def test_below_query_range():
+    query = QueryExpression[int]()
+    assert not query.below_query_range(5)
+    query.with_range(10, 20)
+    assert query.below_query_range(5)
+    assert not query.below_query_range(15)
+    assert not query.below_query_range(25)
+
+
+def test_with_shard_existing_query(sample_range_shard):
+    query = QueryExpression[int]().with_range(10, 20)
+    new_query = QueryExpression.with_shard(query, sample_range_shard)
+    assert new_query.min_key == 5
+    assert new_query.max_key == 20
+
+
+def test_with_shard_none_shard():
+    query = QueryExpression[int]().with_range(10, 20)
+    result = QueryExpression.with_shard(query, None)
+    assert result.min_key == 10
+    assert result.max_key == 20
+
+
+def test_with_shard_existing_query_string(sample_string_shard):
+    query = QueryExpression[str]().with_range("banana", "yellow")
+    new_query = QueryExpression.with_shard(query, sample_string_shard)
+    assert new_query.min_key == "apple"
+    assert new_query.max_key == "zebra"
+
+
+def test_query_expression_string_matches():
+    query = QueryExpression[str]().with_range("apple", "cat")
+    assert query.matches_query("apple")
+    assert query.matches_query("banana")
+    assert not query.matches_query("dog")

deltacat/tests/experimental/storage/rivulet/reader/test_data_scan.py

@@ -0,0 +1,119 @@
+import pytest
+from deltacat.tests.experimental.storage.rivulet.test_utils import verify_pyarrow_scan
+import pyarrow as pa
+from deltacat.experimental.storage.rivulet import Schema, Field, Datatype
+from deltacat.experimental.storage.rivulet.dataset import Dataset
+
+
+@pytest.fixture
+def combined_schema():
+    return Schema(
+        fields=[
+            Field("id", Datatype.int64(), is_merge_key=True),
+            Field("name", Datatype.string()),
+            Field("age", Datatype.int32()),
+            Field("height", Datatype.int64()),
+            Field("gender", Datatype.string()),
+        ]
+    )
+
+
+@pytest.fixture
+def initial_schema():
+    return Schema(
+        fields=[
+            Field("id", Datatype.int32(), is_merge_key=True),
+            Field("name", Datatype.string()),
+            Field("age", Datatype.int32()),
+        ]
+    )
+
+
+@pytest.fixture
+def extended_schema():
+    return Schema(
+        fields=[
+            Field("id", Datatype.int64(), is_merge_key=True),
+            Field("height", Datatype.int64()),
+            Field("gender", Datatype.string()),
+        ]
+    )
+
+
+@pytest.fixture
+def sample_data():
+    return {
+        "id": [1, 2, 3],
+        "name": ["Alice", "Bob", "Charlie"],
+        "age": [25, 30, 35],
+    }
+
+
+@pytest.fixture
+def extended_data():
+    return {
+        "id": [1, 2, 3],
+        "height": [150, 160, 159],
+        "gender": ["male", "female", "male"],
+    }
+
+
+@pytest.fixture
+def combined_data(sample_data, extended_data):
+    data = sample_data.copy()
+    data.update(extended_data)
+    return data
+
+
+@pytest.fixture
+def parquet_data(tmp_path, sample_data):
+    parquet_path = tmp_path / "test.parquet"
+    table = pa.Table.from_pydict(sample_data)
+    pa.parquet.write_table(table, parquet_path)
+    return parquet_path
+
+
+@pytest.fixture
+def sample_dataset(parquet_data, tmp_path):
+    return Dataset.from_parquet(
+        name="test_dataset",
+        file_uri=str(parquet_data),
+        metadata_uri=str(tmp_path),
+        merge_keys="id",
+    )
+
+
+def test_end_to_end_scan_with_multiple_schemas(
+    sample_dataset,
+    initial_schema,
+    extended_schema,
+    combined_schema,
+    sample_data,
+    extended_data,
+    combined_data,
+):
+    # Verify initial scan.
+    verify_pyarrow_scan(sample_dataset.scan().to_arrow(), initial_schema, sample_data)
+
+    # Add a new schema to the dataset
+    sample_dataset.add_schema(schema=extended_schema, schema_name="schema2")
+    new_data = [
+        {"id": 1, "height": 150, "gender": "male"},
+        {"id": 2, "height": 160, "gender": "female"},
+        {"id": 3, "height": 159, "gender": "male"},
+    ]
+    writer = sample_dataset.writer(schema_name="schema2")
+    writer.write(new_data)
+    writer.flush()
+
+    # Verify scan with the extended schema retrieves only extended datfa
+    verify_pyarrow_scan(
+        sample_dataset.scan(schema_name="schema2").to_arrow(),
+        extended_schema,
+        extended_data,
+    )
+
+    # Verify a combined scan retrieves data matching the combined schema
+    verify_pyarrow_scan(
+        sample_dataset.scan().to_arrow(), combined_schema, combined_data
+    )

deltacat/tests/experimental/storage/rivulet/reader/test_dataset_metastore.py

@@ -0,0 +1,71 @@
+import pytest
+import os
+
+from deltacat.experimental.storage.rivulet.metastore.delta import DeltacatManifestIO
+from deltacat.experimental.storage.rivulet.reader.dataset_metastore import (
+    DatasetMetastore,
+)
+from deltacat.experimental.storage.rivulet.schema.datatype import Datatype
+from deltacat.experimental.storage.rivulet.dataset import Dataset
+from deltacat.experimental.storage.rivulet import Schema
+
+
+@pytest.fixture
+def sample_schema():
+    return Schema(
+        {("id", Datatype.int32()), ("name", Datatype.string())},
+        "id",
+    )
+
+
+@pytest.fixture
+def sample_pydict():
+    return {"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"]}
+
+
+def test_dataset_metastore_e2e(sample_schema, tmp_path):
+    # Setup
+    dataset = Dataset(metadata_uri=tmp_path, dataset_name="dataset")
+    file_provider = dataset._file_provider
+    manifest_io = DeltacatManifestIO(file_provider.uri, dataset._locator)
+
+    # Create multiple manifests
+    manifests_data = [
+        {"sst_files": ["sst1.sst", "sst2.sst"], "level": 1},
+        {"sst_files": ["sst3.sst", "sst4.sst"], "level": 2},
+    ]
+
+    # Create SST files and manifests
+    manifest_paths = []
+    for manifest_data in manifests_data:
+        sst_files = manifest_data["sst_files"]
+        for sst in sst_files:
+            with open(os.path.join(file_provider.uri, sst), "w") as f:
+                f.write("test data")
+
+        manifest_path = manifest_io.write(
+            sst_files, sample_schema, manifest_data["level"]
+        )
+        manifest_paths.append(manifest_path)
+
+    # Initialize DatasetMetastore
+    metastore = DatasetMetastore(
+        file_provider.uri,
+        file_provider,
+        file_provider._locator,
+        manifest_io=manifest_io,
+    )
+
+    # Test manifest generation
+    manifest_accessors = list(metastore.generate_manifests())
+    assert len(manifest_accessors) == len(manifests_data)
+
+    # Verify each manifest accessor
+    for accessor in manifest_accessors:
+        assert accessor.context.schema == sample_schema
+        manifests_data_index = 0 if accessor.context.level == 1 else 1
+        assert accessor.context.level == manifests_data[manifests_data_index]["level"]
+        assert (
+            accessor.manifest.sst_files
+            == manifests_data[manifests_data_index]["sst_files"]
+        )

deltacat/tests/{storage → experimental/storage}/rivulet/schema/test_schema.py

@@ -1,6 +1,6 @@
 import pytest
 import pyarrow as pa
-from deltacat.storage.rivulet import Schema, Field, Datatype
+from deltacat.experimental.storage.rivulet import Schema, Field, Datatype
 
 
 def test_field_initialization():