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

deltacat/tests/experimental/storage/rivulet/schema/test_schema.py

@@ -0,0 +1,241 @@
+import pytest
+import pyarrow as pa
+from deltacat.experimental.storage.rivulet import Schema, Field, Datatype
+
+
+def test_field_initialization():
+    field = Field(name="test_field", datatype=Datatype.string(), is_merge_key=True)
+    assert field.name == "test_field"
+    assert field.datatype == Datatype.string()
+    assert field.is_merge_key
+
+
+def test_schema_initialization():
+    fields = [("id", Datatype.int64()), ("name", Datatype.string())]
+    schema = Schema(fields, merge_keys=["id"])
+    assert len(schema) == 2
+    assert "id" in schema.keys()
+    assert schema["id"].datatype == Datatype.int64()
+    assert "name" in schema.keys()
+    assert schema["name"].datatype == Datatype.string()
+
+
+def test_merge_key_conflict_on_init():
+    fields = [
+        Field("id", Datatype.int64(), is_merge_key=False),  # Merge key off here
+        ("name", Datatype.string()),
+    ]
+    with pytest.raises(TypeError):
+        Schema(fields, merge_keys=["id"])  # Merge key on here
+
+
+def test_simultaneous_duplicate_field():
+    with pytest.raises(ValueError):
+        Schema(
+            [
+                ("id", Datatype.int32()),
+                ("name", Datatype.string()),
+                ("age", Datatype.int32()),
+                ("age", Datatype.string()),
+            ],
+            merge_keys=["id"],
+        )
+
+
+def test_add_field():
+    schema = Schema()
+    field = Field("new_field", Datatype.float(), True)
+    schema.add_field(field)
+    assert len(schema) == 1
+    assert "new_field" in schema.keys()
+    assert schema["new_field"].datatype == Datatype.float()
+
+    field2 = Field("another_field", Datatype.string(), True)
+    schema.add_field(field2)
+    assert len(schema) == 2
+    assert "another_field" in schema.keys()
+    assert schema["another_field"].datatype == Datatype.string()
+
+    with pytest.raises(ValueError):
+        schema.add_field(field2)
+
+
+def test_setitem_field():
+    schema = Schema()
+    field = Field("test_field", Datatype.int64(), is_merge_key=True)
+    schema["test_field"] = field
+    assert schema["test_field"] == field
+
+
+def test_setitem_datatype():
+    schema = Schema()
+    schema["id"] = (Datatype.int64(), True)
+    schema["test_field"] = Datatype.int64()
+    assert schema["test_field"].name == "test_field"
+    assert schema["test_field"].datatype == Datatype.int64()
+    assert not schema["test_field"].is_merge_key
+
+
+def test_setitem_tuple_with_merge_key():
+    schema = Schema()
+    schema["test_field"] = (Datatype.int64(), True)
+    assert schema["test_field"].name == "test_field"
+    assert schema["test_field"].datatype == Datatype.int64()
+    assert schema["test_field"].is_merge_key
+
+
+def test_setitem_invalid_type():
+    schema = Schema()
+    with pytest.raises(TypeError):
+        schema["test_field"] = "invalid"
+
+
+def test_non_empty_merge_key():
+    with pytest.raises(TypeError):
+        _ = Schema([], merge_keys=["id"])
+
+
+def test_merge_schemas():
+    schema1 = Schema([("id", Datatype.int64())], merge_keys=["id"])
+    schema2 = Schema(
+        [("other_id", Datatype.string()), ("name", Datatype.string())],
+        merge_keys="other_id",
+    )
+    schema1.merge(schema2)
+    assert len(schema1) == 3
+    assert "id" in schema1.keys()
+    assert "name" in schema1.keys()
+    assert "other_id" in schema1.keys()
+
+
+def test_merge_schemas_same_merge_key():
+    schema1 = Schema(
+        [("id", Datatype.int64()), ("name", Datatype.string())], merge_keys=["id"]
+    )
+    schema2 = Schema(
+        [("id", Datatype.int64()), ("other_name", Datatype.string())],
+        merge_keys="id",
+    )
+    schema1.merge(schema2)
+    assert len(schema1) == 3
+    assert "id" in schema1.keys()
+    assert "name" in schema1.keys()
+    assert "other_name" in schema1.keys()
+
+
+def test_merge_schema_conflict():
+    schema1 = Schema([("id", Datatype.int64())], merge_keys=["id"])
+    schema1_dup = Schema([("id", Datatype.int64())], merge_keys=["id"])
+    schema2 = Schema([("id", Datatype.string())], merge_keys=["id"])
+
+    with pytest.raises(ValueError):
+        schema1.merge(schema2)
+
+    schema1.merge(
+        schema1_dup
+    )  # Merging the same field is allowed (unlike using add_field)
+    assert schema1["id"].datatype == Datatype.int64()
+    assert len(schema1) == 1
+
+
+def test_to_pyarrow_schema():
+    fields = [("id", Datatype.int64()), ("name", Datatype.string())]
+    schema = Schema(fields, merge_keys=["id"])
+    pa_schema = schema.to_pyarrow()
+    assert isinstance(pa_schema, pa.Schema)
+    assert len(pa_schema) == 2
+    assert pa_schema.field("id").type == pa.int64()
+    assert pa_schema.field("name").type == pa.string()
+
+
+def test_from_pyarrow_schema():
+    pa_schema = pa.schema([("id", pa.int64()), ("name", pa.string())])
+    schema = Schema.from_pyarrow(pa_schema, merge_keys=["id"])
+    assert len(schema) == 2
+    assert schema["id"].is_merge_key
+
+
+def test_from_pyarrow_schema_invalid_merge_keys():
+    pa_schema = pa.schema([("id", pa.int64()), ("name", pa.string())])
+    with pytest.raises(ValueError):
+        Schema.from_pyarrow(pa_schema, merge_keys=["bad_key"])
+
+
+def test_get_field():
+    schema = Schema([("id", Datatype.int64())], merge_keys=["id"])
+    field = schema["id"]
+    assert field.name == "id"
+    assert field.datatype == Datatype.int64()
+
+
+def test_set_field():
+    schema = Schema([("id", Datatype.int64())], merge_keys=["id"])
+    schema["name"] = Field("name", Datatype.string())
+    assert len(schema) == 2
+    assert "name" in schema.keys()
+    assert schema["name"].datatype == Datatype.string()
+
+
+def test_delete_field():
+    schema = Schema(
+        [("name", Datatype.string()), ("zip", Datatype.int32())], merge_keys=["name"]
+    )
+    del schema["zip"]
+    assert "zip" not in schema.keys()
+    assert "name" in schema.keys()
+
+
+def test_delete_merge_key_field():
+    schema = Schema([("id", Datatype.int64())], merge_keys=["id"])
+    with pytest.raises(ValueError):
+        del schema["id"]
+
+
+def test_schema_iter():
+    fields = [
+        Field("id", Datatype.int32(), is_merge_key=True),
+        Field("name", Datatype.string()),
+    ]
+    schema = Schema(fields)
+    iter_result = list(iter(schema))
+    assert len(iter_result) == 2
+    assert all(isinstance(item, str) for item in iter_result)
+
+
+def test_merge_all():
+    schema1 = Schema(
+        [
+            Field("id", Datatype.int64(), is_merge_key=True),
+            Field("name", Datatype.string()),
+        ]
+    )
+    schema2 = Schema(
+        [
+            Field("age", Datatype.int32()),
+            Field("email", Datatype.string(), is_merge_key=True),
+        ]
+    )
+    merged_schema = Schema.merge_all([schema1, schema2])
+    assert len(merged_schema) == 4
+
+
+def test_schema_values():
+    fields = [
+        Field("id", Datatype.int64(), is_merge_key=True),
+        Field("name", Datatype.string()),
+    ]
+    schema = Schema(fields)
+    values = list(schema.values())
+    assert len(values) == 2
+    assert all(isinstance(v, Field) for v in values)
+
+
+def test_schema_items():
+    fields = [
+        Field("id", Datatype.int64(), is_merge_key=True),
+        Field("name", Datatype.string()),
+    ]
+    schema = Schema(fields)
+    items = list(schema.items())
+    assert len(items) == 2
+    assert all(isinstance(k, str) and isinstance(v, Field) for k, v in items)

deltacat/tests/experimental/storage/rivulet/shard/test_range_shard.py

@@ -0,0 +1,162 @@
+import pytest
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from deltacat.experimental.storage.rivulet.dataset import Dataset
+from deltacat.experimental.storage.rivulet.shard.range_shard import (
+    RangeShard,
+    RangeShardingStrategy,
+)
+
+
+@pytest.fixture
+def sample_numeric_dataset(tmp_path):
+    """
+    Creates a small Parquet file with integer-based min/max keys and
+    initializes a Dataset from it. Merge key is 'id' with values [1,2,3].
+    So min_key=1, max_key=3.
+    """
+    data = {
+        "id": [1, 2, 3],
+        "name": ["Alice", "Bob", "Charlie"],
+        "age": [25, 30, 35],
+    }
+    table = pa.Table.from_pydict(data)
+    parquet_file = tmp_path / "numeric_data.parquet"
+    pq.write_table(table, parquet_file)
+
+    ds = Dataset.from_parquet(
+        name="numeric_dataset",
+        file_uri=str(parquet_file),
+        metadata_uri=tmp_path,
+        merge_keys="id",
+    )
+    return ds
+
+
+@pytest.fixture
+def sample_string_dataset(tmp_path):
+    """
+    Creates a small Parquet file with a string-based merge key ('name')
+    and initializes a Dataset from it. Merge key has values
+    ['Alice', 'Bob', 'Charlie'] => min_key='Alice', max_key='Charlie'.
+    """
+    data = {
+        "name": ["Alice", "Charlie", "Bob"],  # random order
+        "value": [100, 200, 150],
+    }
+    table = pa.Table.from_pydict(data)
+    parquet_file = tmp_path / "string_data.parquet"
+    pq.write_table(table, parquet_file)
+
+    ds = Dataset.from_parquet(
+        name="string_dataset",
+        file_uri=str(parquet_file),
+        metadata_uri=tmp_path,
+        merge_keys="name",
+    )
+    return ds
+
+
+def test_shards(sample_numeric_dataset, sample_string_dataset):
+    shards = sample_numeric_dataset.shards(num_shards=2)
+
+    num_shards = len(list(shards))
+    assert num_shards == 2
+
+    shard = shards[0]
+    records = list(sample_numeric_dataset.scan(shard=shard).to_pydict())
+    num_records = len(records)
+    assert num_records == 2
+
+    assert records[0]["id"] == 1
+    assert records[0]["name"] == "Alice"
+
+    assert records[1]["id"] == 2
+    assert records[1]["name"] == "Bob"
+
+
+def test_range_shard_repr():
+    shard = RangeShard(min_key=5, max_key=15)
+    assert repr(shard) == "Shard(type=range, min_key=5, max_key=15)"
+
+
+def test_range_shard_split_integers():
+    shards = RangeShard.split(global_min=1, global_max=10, num_shards=2)
+    assert len(shards) == 2
+
+    assert shards[0].min_key == 1
+    assert shards[0].max_key == 5
+    assert shards[1].min_key == 6
+    assert shards[1].max_key == 10
+
+
+def test_range_shard_split_integers_single_shard():
+    shards = RangeShard.split(global_min=1, global_max=10, num_shards=1)
+    assert len(shards) == 1
+    assert shards[0].min_key == 1
+    assert shards[0].max_key == 10
+
+
+def test_range_shard_split_integers_same_value():
+    shards = RangeShard.split(global_min=5, global_max=5, num_shards=3)
+    assert len(shards) == 1
+
+
+def test_range_sharding_strategy_integers(sample_numeric_dataset):
+    strategy = RangeShardingStrategy()
+    shards = list(
+        strategy.shards(num_shards=2, metastore=sample_numeric_dataset._metastore)
+    )
+
+    assert len(shards) == 2, "Expected 2 shards for dataset with keys [1,2,3]"
+
+    shard1, shard2 = shards
+    assert isinstance(shard1, RangeShard)
+    assert isinstance(shard2, RangeShard)
+    assert shard1.min_key == 1
+    assert shard1.max_key == 2
+    assert shard2.min_key == 3
+    assert shard2.max_key == 3
+
+
+def test_range_sharding_strategy_integers_single_shard(sample_numeric_dataset):
+    strategy = RangeShardingStrategy()
+    shards = list(
+        strategy.shards(num_shards=1, metastore=sample_numeric_dataset._metastore)
+    )
+    assert len(shards) == 1
+    shard = shards[0]
+    assert shard.min_key == 1
+    assert shard.max_key == 3
+
+
+def test_range_sharding_strategy_strings(sample_string_dataset):
+    strategy = RangeShardingStrategy()
+    shards = list(
+        strategy.shards(num_shards=2, metastore=sample_string_dataset._metastore)
+    )
+
+    assert len(shards) == 2, "Expected 2 shards for string-based dataset"
+    shard1, shard2 = shards
+    assert isinstance(shard1, RangeShard)
+    assert isinstance(shard2, RangeShard)
+
+    assert shard1.min_key == "Alice"
+    assert shard1.max_key < "Charlie"
+
+    assert shard2.min_key == shard1.max_key
+    assert shard2.max_key == "Charlie"
+
+
+def test_range_sharding_strategy_strings_single_shard(sample_string_dataset):
+    strategy = RangeShardingStrategy()
+    shards = list(
+        strategy.shards(num_shards=1, metastore=sample_string_dataset._metastore)
+    )
+
+    assert len(shards) == 1
+
+    shard = shards[0]
+    assert shard.min_key == "Alice"
+    assert shard.max_key == "Charlie"