deltacat 1.1.36__py3-none-any.whl → 2.0__py3-none-any.whl

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

@@ -0,0 +1,241 @@
+import pytest
+import pyarrow as pa
+from deltacat.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/storage/rivulet/test_dataset.py

@@ -0,0 +1,406 @@
+import posixpath
+from deltacat.utils.metafile_locator import _find_partition_path
+import pytest
+
+import pyarrow as pa
+from deltacat.storage.rivulet import Schema, Field, Datatype
+from deltacat.storage.rivulet.dataset import Dataset
+from deltacat.storage.rivulet.reader.query_expression import QueryExpression
+
+
+@pytest.fixture
+def sample_schema():
+    return Schema(
+        fields=[
+            Field("id", Datatype.int32(), is_merge_key=True),
+            Field("name", Datatype.string()),
+            Field("age", Datatype.int32()),
+        ]
+    )
+
+
+@pytest.fixture
+def sample_pydict():
+    return {"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"], "age": [25, 30, 35]}
+
+
+@pytest.fixture
+def sample_parquet_data(tmp_path, sample_pydict):
+    parquet_path = tmp_path / "test.parquet"
+    table = pa.Table.from_pydict(sample_pydict)
+    pa.parquet.write_table(table, parquet_path)
+    return parquet_path
+
+
+# Updated Tests
+
+
+def test_dataset_creation_with_schema(tmp_path, sample_schema):
+    dataset = Dataset(dataset_name="test_dataset", schema=sample_schema)
+    assert len(dataset.fields) == 3
+    assert "id" in dataset.fields
+    assert dataset.fields["id"].is_merge_key
+
+
+def test_dataset_initialization_with_metadata(tmp_path):
+    dataset = Dataset(dataset_name="test_dataset", metadata_uri=str(tmp_path))
+    assert dataset.dataset_name == "test_dataset"
+    assert dataset._metadata_folder.startswith(".riv-meta")
+
+
+def test_invalid_dataset_initialization():
+    with pytest.raises(ValueError, match="Name must be a non-empty string"):
+        Dataset(dataset_name="")
+
+
+def test_dataset_creation_metadata_structure(tmp_path):
+    dataset = Dataset(dataset_name="test_dataset", metadata_uri=str(tmp_path))
+
+    assert dataset._metadata_folder.startswith(".riv-meta")
+    assert dataset._namespace == "DEFAULT"
+    assert dataset.dataset_name == "test_dataset"
+    assert dataset._metadata_path == str(tmp_path / ".riv-meta-test_dataset")
+
+    locator = dataset._locator
+    root_uri = dataset._metadata_path
+
+    partition_path = _find_partition_path(root_uri, locator)
+
+    # Ensures that directory structure for namespace -> table -> table_version -> stream_id -> partition_id exists
+    assert posixpath.exists(partition_path)
+
+
+def test_fields_accessor_add_field(tmp_path, sample_schema):
+    dataset = Dataset(dataset_name="test_dataset", schema=sample_schema)
+    dataset.fields.add("new_field", Datatype.float())
+    assert "new_field" in dataset.fields
+    assert dataset.fields["new_field"].datatype == Datatype.float()
+
+    dataset.fields["new_field2"] = Field("new_field2", Datatype.int32())
+    assert "new_field2" in dataset.fields
+    assert "new_field2" in dataset.schemas["all"]
+    with pytest.raises(TypeError):
+        dataset.fields["new_field3"] = 2
+
+
+def test_field_removal(tmp_path, sample_schema):
+    dataset = Dataset(dataset_name="test_dataset", schema=sample_schema)
+    del dataset.fields["age"]
+    assert "age" not in dataset.fields
+    with pytest.raises(ValueError):
+        del dataset.fields["age"]
+    with pytest.raises(KeyError):
+        _ = dataset.fields["age"]
+
+
+def test_fields_accessor_repr(tmp_path, sample_schema):
+    dataset = Dataset(dataset_name="test_dataset", schema=sample_schema)
+    repr_output = repr(dataset.fields)
+    for field_name in ["id", "name", "age"]:
+        assert field_name in repr_output, f"Field '{field_name}' missing in repr output"
+
+
+def test_schemas_accessor_add_group(tmp_path, sample_schema):
+    dataset = Dataset(dataset_name="test_dataset", schema=sample_schema)
+    dataset.schemas["analytics"] = ["id", "name"]
+    assert "analytics" in dataset.schemas
+    assert len(dataset.schemas["analytics"]) == 2
+
+
+def test_schema_removal(tmp_path, sample_schema):
+    dataset = Dataset(dataset_name="test_dataset", schema=sample_schema)
+    with pytest.raises(ValueError):
+        del dataset.schemas["all"]
+    with pytest.raises(ValueError):
+        del dataset.schemas["does_not_exist"]
+    dataset.schemas["new"] = ["id", "name"]
+    del dataset.schemas["new"]
+    with pytest.raises(KeyError):
+        _ = dataset.schemas["new"]
+
+
+def test_dataset_from_parquet(tmp_path, sample_parquet_data):
+    dataset = Dataset.from_parquet(
+        name="test_dataset",
+        file_uri=str(sample_parquet_data),
+        metadata_uri=str(tmp_path),
+        merge_keys="id",
+    )
+    assert len(dataset.fields) == 3
+    assert "id" in dataset.fields
+    assert dataset.fields["id"].is_merge_key
+
+
+def test_parquet_schema_modes(tmp_path, sample_pydict):
+    # Create two parquet files with overlapping and unique schemas
+    data_1 = {"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"]}
+    data_2 = {"id": [4, 5, 6], "age": [25, 30, 35]}
+
+    path_1 = tmp_path / "data1.parquet"
+    path_2 = tmp_path / "data2.parquet"
+    pa.parquet.write_table(pa.Table.from_pydict(data_1), path_1)
+    pa.parquet.write_table(pa.Table.from_pydict(data_2), path_2)
+
+    dataset_union = Dataset.from_parquet(
+        name="test_dataset_union",
+        file_uri=str(tmp_path),
+        merge_keys="id",
+        schema_mode="union",
+    )
+    assert len(dataset_union.fields) == 3  # id, name, age
+
+    dataset_intersect = Dataset.from_parquet(
+        name="test_dataset_intersect",
+        file_uri=str(tmp_path),
+        merge_keys="id",
+        schema_mode="intersect",
+    )
+    assert len(dataset_intersect.fields) == 1  # Only id
+
+
+def test_merge_all_schemas():
+    schema1 = Schema(
+        fields=[
+            Field("id", Datatype.int32(), is_merge_key=True),
+            Field("name", Datatype.string()),
+        ]
+    )
+    schema2 = Schema(
+        fields=[
+            Field("id", Datatype.int32(), is_merge_key=True),
+            Field("age", Datatype.int32()),
+        ]
+    )
+    merged_schema = Schema.merge_all([schema1, schema2])
+    assert len(merged_schema) == 3
+    assert "id" in merged_schema
+    assert "name" in merged_schema
+    assert "age" in merged_schema
+
+
+def test_writer_creation_with_custom_format(tmp_path, sample_schema):
+    dataset = Dataset(dataset_name="test_dataset", schema=sample_schema)
+    writer = dataset.writer(file_format="feather")
+    assert writer is not None
+
+
+def test_scan_with_query(tmp_path, sample_schema):
+    dataset = Dataset(dataset_name="test_dataset", schema=sample_schema)
+    query = QueryExpression()  # Placeholder query
+    scan = dataset.scan(query)
+    assert scan is not None
+
+
+def test_add_schema_to_new_schemas(tmp_path):
+    """Test adding a schema to a new field group."""
+    base_uri = str(tmp_path / "test_dataset")
+    dataset = Dataset(dataset_name=base_uri)
+
+    schema = Schema(
+        [
+            ("id", Datatype.int32()),
+            ("name", Datatype.string()),
+            ("age", Datatype.int32()),
+        ],
+        merge_keys=["id"],
+    )
+
+    dataset.add_schema(schema, schema_name="new_group")
+
+    # Verify the field group is added
+    assert "new_group" in dataset.schemas
+    assert len(dataset.schemas["new_group"]) == 3
+    assert dataset.schemas["new_group"]["id"].datatype == Datatype.int32()
+    assert dataset.schemas["new_group"]["name"].datatype == Datatype.string()
+    assert dataset.schemas["new_group"]["age"].datatype == Datatype.int32()
+
+
+def test_add_schema_to_existing_schemas(tmp_path):
+    """Test merging a schema into an existing field group."""
+    base_uri = str(tmp_path / "test_dataset")
+    dataset = Dataset(dataset_name=base_uri)
+
+    schema_1 = Schema(
+        [
+            ("id", Datatype.int32()),
+            ("name", Datatype.string()),
+        ],
+        merge_keys=["id"],
+    )
+
+    dataset.add_schema(schema_1, schema_name="existing_group")
+
+    schema_2 = Schema(
+        [
+            ("age", Datatype.int32()),
+            ("email", Datatype.string()),
+        ],
+        merge_keys=["id"],
+    )
+
+    dataset.add_schema(schema_2, schema_name="existing_group")
+
+    # Verify the merged schema
+    assert "existing_group" in dataset.schemas
+    assert len(dataset.schemas["existing_group"]) == 4
+    assert dataset.schemas["existing_group"]["id"].datatype == Datatype.int32()
+    assert dataset.schemas["existing_group"]["name"].datatype == Datatype.string()
+    assert dataset.schemas["existing_group"]["age"].datatype == Datatype.int32()
+    assert dataset.schemas["existing_group"]["email"].datatype == Datatype.string()
+
+
+def test_add_schema_conflicting_fields(tmp_path):
+    """Test adding a schema with conflicting fields."""
+    base_uri = str(tmp_path / "test_dataset")
+    dataset = Dataset(dataset_name=base_uri)
+
+    schema_1 = Schema(
+        [
+            ("id", Datatype.int32()),
+            ("name", Datatype.string()),
+        ],
+        merge_keys=["id"],
+    )
+
+    dataset.add_schema(schema_1, schema_name="conflicting_group")
+
+    schema_2 = Schema(
+        [
+            ("id", Datatype.string()),  # Conflict: datatype mismatch
+            ("age", Datatype.int32()),
+        ],
+        merge_keys=["id"],
+    )
+
+    with pytest.raises(ValueError, match="already exists"):
+        dataset.add_schema(schema_2, schema_name="conflicting_group")
+
+    schema_3 = Schema(
+        [
+            ("id", Datatype.int32()),  # Conflict: datatype mismatch
+            ("age", Datatype.int32()),
+        ],
+        merge_keys=["id"],
+    )
+
+    dataset.add_schema(schema_3, schema_name="conflicting_group")
+    assert "conflicting_group" in dataset.schemas
+    assert len(dataset.schemas["conflicting_group"]) == 3
+    assert dataset.schemas["conflicting_group"]["id"].datatype == Datatype.int32()
+    assert dataset.schemas["conflicting_group"]["name"].datatype == Datatype.string()
+    assert dataset.schemas["conflicting_group"]["age"].datatype == Datatype.int32()
+
+
+def test_add_fields_with_merge_key_field(tmp_path):
+    base_uri = str(tmp_path / "test_dataset")
+    dataset = Dataset(dataset_name=base_uri)
+    dataset.add_fields([Field("my_merge_key", Datatype.string(), True)])
+    assert dataset.schemas["default"].get_merge_key() == "my_merge_key"
+
+
+def test_add_schema_to_nonexistent_schemas(tmp_path):
+    """Test adding a schema to a nonexistent field group."""
+    base_uri = str(tmp_path / "test_dataset")
+    dataset = Dataset(dataset_name=base_uri)
+
+    schema = Schema(
+        [
+            ("id", Datatype.int32()),
+            ("name", Datatype.string()),
+        ],
+        merge_keys=["id"],
+    )
+
+    # Add to a non-existent field group
+    dataset.add_schema(schema, schema_name="nonexistent_group")
+
+    # Verify the field group is created
+    assert "nonexistent_group" in dataset.schemas
+    assert len(dataset.schemas["nonexistent_group"]) == 2
+
+
+def test_add_missing_field_to_schema_raises_error(tmp_path, sample_schema):
+    """
+    Test that attempting to add a missing field to the 'all' schema raises a ValueError.
+    """
+    dataset = Dataset(dataset_name="test_dataset", schema=sample_schema)
+
+    # Attempt to add a non-existent field to the 'all' schema
+    with pytest.raises(
+        ValueError, match="Field 'missing_field' does not exist in the dataset."
+    ):
+        dataset.schemas["all"] = [
+            "missing_field"
+        ]  # Attempt to set a list with a missing field
+
+
+def test_schemas_accessor_methods(tmp_path, sample_schema):
+    """
+    Test the __iter__, __len__, and __repr__ methods of SchemasAccessor.
+    """
+    dataset = Dataset(
+        dataset_name="test_dataset", schema=sample_schema
+    )  # Default schema is defined automatically
+    dataset.schemas["schema_1"] = ["id", "name"]
+    dataset.schemas["schema_2"] = ["age"]
+
+    # Test __iter__
+    schema_names = list(iter(dataset.schemas))
+    assert set(schema_names) == {
+        "schema_1",
+        "schema_2",
+        "all",
+        "default",
+    }, "Schema names do not match expected values"
+
+    # Test __len__
+    assert len(dataset.schemas) == 4, "Length of schemas accessor is incorrect"
+
+    # Test __repr__
+    repr_output = repr(dataset.schemas)
+    for schema_name in ["schema_1", "schema_2", "all"]:
+        assert (
+            schema_name in repr_output
+        ), f"Schema '{schema_name}' missing in repr output"
+
+
+def test_get_merge_keys(tmp_path, sample_schema):
+    """
+    Test the get_merge_keys method to ensure it returns all merge keys in the dataset.
+    """
+    dataset = Dataset(dataset_name="test_dataset", schema=sample_schema)
+
+    # Add fields with additional merge key to the dataset
+    other_schema = Schema(
+        [("id2", Datatype.int32()), ("zip", Datatype.string())], merge_keys=["id2"]
+    )
+
+    dataset.add_schema(other_schema, "id2+zip")
+
+    # Call get_merge_keys and validate the result
+    merge_keys = dataset.get_merge_keys()
+    assert merge_keys == [
+        "id",
+        "id2",
+    ], f"Expected merge keys ['id', 'id2'], got {merge_keys}"
+
+
+def test_add_fields_no_fields_raises_error(tmp_path, sample_schema):
+    dataset = Dataset(dataset_name="test_dataset")
+    with pytest.raises(ValueError):
+        dataset.add_fields(fields=[])
+
+
+def test_add_fields_mismatched_merge_keys_raises_error(tmp_path, sample_schema):
+    dataset = Dataset(dataset_name="test_dataset")
+    with pytest.raises(
+        ValueError,
+        match="The following merge keys were not found in the provided fields: does_not_exist",
+    ):
+        dataset.add_fields(fields=sample_schema.values(), merge_keys=["does_not_exist"])
+
+    with pytest.raises(TypeError, match="Merge key status conflict"):
+        dataset.add_fields(
+            fields=[Field("id", Datatype.int32()), Field("name", Datatype.string())],
+            merge_keys=["id"],
+        )