deltacat 0.1.18b14__py3-none-any.whl → 0.1.18b15__py3-none-any.whl

deltacat/tests/compute/testcases.py

@@ -0,0 +1,390 @@
+import pyarrow as pa
+from typing import Dict, List
+from deltacat.tests.compute.common import (
+    MAX_RECORDS_PER_FILE,
+    offer_iso8601_timestamp_list,
+)
+from deltacat.tests.compute.common import (
+    BASE_TEST_SOURCE_TABLE_VERSION,
+    BASE_TEST_DESTINATION_TABLE_VERSION,
+    HASH_BUCKET_COUNT,
+)
+from deltacat.compute.compactor.compaction_session import (
+    compact_partition_from_request as compact_partition_v1,
+)
+
+
+def create_tests_cases_for_all_compactor_versions(test_cases: Dict[str, List]):
+    final_cases = {}
+    for version, compact_partition_func in enumerate([compact_partition_v1]):
+        for case_name, case_value in test_cases.items():
+            final_cases[f"{case_name}_v{version}"] = [
+                *case_value,
+                compact_partition_func,
+            ]
+
+    return final_cases
+
+
+"""
+TODO Test Cases:
+1. incremental w/wout round completion file
+2. Backfill w/wout round completion
+3. Rebase w/wout round completion file
+4. Rebase then incremental (use same round completion file)
+"""
+
+
+INCREMENTAL_INDEPENDENT_TEST_CASES = {
+    "1-incremental-pkstr-sknone-norcf": [
+        BASE_TEST_SOURCE_TABLE_VERSION,
+        BASE_TEST_DESTINATION_TABLE_VERSION,
+        {"pk_col_1"},  # Primary key columns
+        [],  # Sort key columns
+        [{"key_name": "region_id", "key_type": "int"}],  # Partition keys
+        ["pk_col_1"],  # column_names
+        [pa.array([str(i) for i in range(10)])],  # arrow arrays
+        None,  # rebase_source_partition_locator_param
+        ["1"],  # partition_values_param
+        pa.Table.from_arrays(
+            [pa.array([str(i) for i in range(10)])],
+            names=["pk_col_1"],  # expected_result
+        ),
+        None,  # validation_callback_func
+        None,  # validation_callback_func_kwargs
+        True,  # teardown_local_deltacat_storage_db
+        False,  # use_prev_compacted
+        True,  # create_placement_group_param
+        MAX_RECORDS_PER_FILE,  # records_per_compacted_file_param
+        HASH_BUCKET_COUNT,  # hash_bucket_count_param
+    ],
+    "2-incremental-pkstr-skstr-norcf": [
+        BASE_TEST_SOURCE_TABLE_VERSION,
+        BASE_TEST_DESTINATION_TABLE_VERSION,
+        ["pk_col_1"],
+        [
+            {
+                "key_name": "sk_col_1",
+            }
+        ],
+        [],
+        ["pk_col_1", "sk_col_1"],
+        [pa.array([str(i) for i in range(10)]), pa.array(["test"] * 10)],
+        None,
+        ["1"],
+        pa.Table.from_arrays(
+            [pa.array([str(i) for i in range(10)]), pa.array(["test"] * 10)],
+            names=["pk_col_1", "sk_col_1"],
+        ),
+        None,
+        None,
+        True,
+        False,
+        True,
+        MAX_RECORDS_PER_FILE,
+        HASH_BUCKET_COUNT,
+    ],
+    "3-incremental-pkstr-multiskstr-norcf": [
+        BASE_TEST_SOURCE_TABLE_VERSION,
+        BASE_TEST_DESTINATION_TABLE_VERSION,
+        ["pk_col_1"],
+        [
+            {
+                "key_name": "sk_col_1",
+            },
+            {
+                "key_name": "sk_col_2",
+            },
+        ],
+        [],
+        ["pk_col_1", "sk_col_1", "sk_col_2"],
+        [
+            pa.array([str(i) for i in range(10)]),
+            pa.array(["test"] * 10),
+            pa.array(["foo"] * 10),
+        ],
+        None,
+        ["1"],
+        pa.Table.from_arrays(
+            [
+                pa.array([str(i) for i in range(10)]),
+                pa.array(["test"] * 10),
+                pa.array(["foo"] * 10),
+            ],
+            names=["pk_col_1", "sk_col_1", "sk_col_2"],
+        ),
+        None,
+        None,
+        True,
+        False,
+        True,
+        MAX_RECORDS_PER_FILE,
+        HASH_BUCKET_COUNT,
+    ],
+    "4-incremental-duplicate-pk": [
+        BASE_TEST_SOURCE_TABLE_VERSION,
+        BASE_TEST_DESTINATION_TABLE_VERSION,
+        ["pk_col_1"],
+        [
+            {
+                "key_name": "sk_col_1",
+            },
+            {
+                "key_name": "sk_col_2",
+            },
+        ],
+        [],
+        ["pk_col_1", "sk_col_1", "sk_col_2"],
+        [
+            pa.array([str(i) for i in range(5)] + ["6", "6", "6", "6", "6"]),
+            pa.array([str(i) for i in range(10)]),
+            pa.array(["foo"] * 10),
+        ],
+        None,
+        ["1"],
+        pa.Table.from_arrays(
+            [
+                pa.array([str(i) for i in range(5)] + ["6"]),
+                pa.array([str(i) for i in range(5)] + ["9"]),
+                pa.array(["foo"] * 6),
+            ],
+            names=["pk_col_1", "sk_col_1", "sk_col_2"],
+        ),
+        None,
+        None,
+        True,
+        False,
+        True,
+        MAX_RECORDS_PER_FILE,
+        HASH_BUCKET_COUNT,
+    ],
+    "5-incremental-decimal-pk-simple": [
+        BASE_TEST_SOURCE_TABLE_VERSION,
+        BASE_TEST_DESTINATION_TABLE_VERSION,
+        ["pk_col_1"],
+        [
+            {
+                "key_name": "sk_col_1",
+            },
+        ],
+        [],
+        ["pk_col_1", "sk_col_1"],
+        [
+            pa.array([i / 10 for i in range(0, 10)]),
+            pa.array([str(i) for i in range(10)]),
+        ],
+        None,
+        ["1"],
+        pa.Table.from_arrays(
+            [
+                pa.array([i / 10 for i in range(0, 10)]),
+                pa.array([str(i) for i in range(10)]),
+            ],
+            names=["pk_col_1", "sk_col_1"],
+        ),
+        None,
+        None,
+        True,
+        False,
+        True,
+        MAX_RECORDS_PER_FILE,
+        HASH_BUCKET_COUNT,
+    ],
+    "7-incremental-integer-pk-simple": [
+        BASE_TEST_SOURCE_TABLE_VERSION,
+        BASE_TEST_DESTINATION_TABLE_VERSION,
+        ["pk_col_1"],
+        [
+            {
+                "key_name": "sk_col_1",
+            },
+        ],
+        [],
+        ["pk_col_1", "sk_col_1"],
+        [
+            pa.array([i for i in range(0, 10)]),
+            pa.array([str(i) for i in range(10)]),
+        ],
+        None,
+        ["1"],
+        pa.Table.from_arrays(
+            [
+                pa.array([i for i in range(0, 10)]),
+                pa.array([str(i) for i in range(10)]),
+            ],
+            names=["pk_col_1", "sk_col_1"],
+        ),
+        None,
+        None,
+        True,
+        False,
+        True,
+        MAX_RECORDS_PER_FILE,
+        HASH_BUCKET_COUNT,
+    ],
+    "8-incremental-timestamp-pk-simple": [
+        BASE_TEST_SOURCE_TABLE_VERSION,
+        BASE_TEST_DESTINATION_TABLE_VERSION,
+        ["pk_col_1"],
+        [
+            {
+                "key_name": "sk_col_1",
+            },
+        ],
+        [],
+        ["pk_col_1", "sk_col_1"],
+        [
+            pa.array(offer_iso8601_timestamp_list(10, "minutes")),
+            pa.array([str(i) for i in range(10)]),
+        ],
+        None,
+        ["1"],
+        pa.Table.from_arrays(
+            [
+                pa.array(offer_iso8601_timestamp_list(10, "minutes")),
+                pa.array([str(i) for i in range(10)]),
+            ],
+            names=["pk_col_1", "sk_col_1"],
+        ),
+        None,
+        None,
+        True,
+        False,
+        True,
+        MAX_RECORDS_PER_FILE,
+        HASH_BUCKET_COUNT,
+    ],
+    "9-incremental-decimal-timestamp-pk-multi": [
+        BASE_TEST_SOURCE_TABLE_VERSION,
+        BASE_TEST_DESTINATION_TABLE_VERSION,
+        ["pk_col_1", "pk_col_2"],
+        [
+            {
+                "key_name": "sk_col_1",
+            },
+        ],
+        [],
+        ["pk_col_1", "pk_col_2", "sk_col_1"],
+        [
+            pa.array([i / 10 for i in range(0, 20)]),
+            pa.array(offer_iso8601_timestamp_list(20, "minutes")),
+            pa.array([0.1] * 4 + [0.2] * 4 + [0.3] * 4 + [0.4] * 4 + [0.5] * 4),
+        ],
+        None,
+        ["1"],
+        pa.Table.from_arrays(
+            [
+                pa.array([i / 10 for i in range(0, 20)]),
+                pa.array(offer_iso8601_timestamp_list(20, "minutes")),
+                pa.array([0.1] * 4 + [0.2] * 4 + [0.3] * 4 + [0.4] * 4 + [0.5] * 4),
+            ],
+            names=["pk_col_1", "pk_col_2", "sk_col_1"],
+        ),
+        None,
+        None,
+        True,
+        False,
+        True,
+        MAX_RECORDS_PER_FILE,
+        HASH_BUCKET_COUNT,
+    ],
+    "10-incremental-decimal-pk-multi-dup": [
+        BASE_TEST_SOURCE_TABLE_VERSION,
+        BASE_TEST_DESTINATION_TABLE_VERSION,
+        ["pk_col_1"],
+        [
+            {
+                "key_name": "sk_col_1",
+            },
+        ],
+        [],
+        ["pk_col_1", "sk_col_1"],
+        [
+            pa.array([0.1] * 4 + [0.2] * 4 + [0.3] * 4 + [0.4] * 4 + [0.5] * 4),
+            pa.array(reversed([i for i in range(20)])),
+        ],
+        None,
+        ["1"],
+        pa.Table.from_arrays(
+            [
+                pa.array([0.1, 0.2, 0.3, 0.4, 0.5]),
+                pa.array([19, 15, 11, 7, 3]),
+            ],
+            names=["pk_col_1", "sk_col_1"],
+        ),
+        None,
+        None,
+        True,
+        False,
+        True,
+        MAX_RECORDS_PER_FILE,
+        HASH_BUCKET_COUNT,
+    ],
+}
+
+"""
+for test_name, (
+            source_table_version,
+            destination_table_version,
+            primary_keys_param,
+            sort_keys_param,
+            partition_keys_param,
+            column_names_param,
+            arrow_arrays_param,
+            rebase_source_partition_locator_param,
+            partition_values_param,
+            expected_result,
+            validation_callback_func,
+            validation_callback_func_kwargs,
+            do_teardown_local_deltacat_storage_db,
+            use_prev_compacted,
+            create_placement_group_param,
+            records_per_compacted_file_param,
+            hash_bucket_count_param,
+        ) in INCREMENTAL_TEST_CASES.items()
+"""
+
+# TODO: Add test cases where next tc is dependent on the previous compacted table existing
+INCREMENTAL_DEPENDENT_TEST_CASES = {
+    "11-incremental-multi-dup-retain-table": (
+        BASE_TEST_SOURCE_TABLE_VERSION,
+        BASE_TEST_DESTINATION_TABLE_VERSION,
+        ["pk_col_1", "pk_col_2"],
+        [
+            {
+                "key_name": "sk_col_1",
+            },
+        ],
+        [],
+        ["pk_col_1", "pk_col_2", "sk_col_1"],
+        [
+            pa.array([i / 10 for i in range(0, 20)]),
+            pa.array(offer_iso8601_timestamp_list(20, "minutes")),
+            pa.array([0.1] * 4 + [0.2] * 4 + [0.3] * 4 + [0.4] * 4 + [0.5] * 4),
+        ],
+        None,
+        ["1"],
+        pa.Table.from_arrays(
+            [
+                pa.array([i / 10 for i in range(0, 20)]),
+                pa.array(offer_iso8601_timestamp_list(20, "minutes")),
+                pa.array([0.1] * 4 + [0.2] * 4 + [0.3] * 4 + [0.4] * 4 + [0.5] * 4),
+            ],
+            names=["pk_col_1", "pk_col_2", "sk_col_1"],
+        ),
+        None,
+        None,
+        False,
+        False,
+        True,
+        MAX_RECORDS_PER_FILE,
+        HASH_BUCKET_COUNT,
+    ),
+}
+
+INCREMENTAL_TEST_CASES = create_tests_cases_for_all_compactor_versions(
+    {
+        **INCREMENTAL_INDEPENDENT_TEST_CASES,
+        **INCREMENTAL_DEPENDENT_TEST_CASES,
+    }
+)

deltacat/tests/io/test_memcached_object_store.py

@@ -11,7 +11,9 @@ class TestMemcachedObjectStore(unittest.TestCase):
     def setUp(self):
         from deltacat.io.memcached_object_store import MemcachedObjectStore
 
-        self.object_store = MemcachedObjectStore()
+        self.object_store = MemcachedObjectStore(
+            storage_node_ips=["172.1.1.1", "172.2.2.2", "172.3.3.3"]
+        )
 
     @mock.patch("deltacat.io.memcached_object_store.Client")
     @mock.patch("deltacat.io.memcached_object_store.RetryingClient")
@@ -29,11 +31,10 @@ class TestMemcachedObjectStore(unittest.TestCase):
         mock_retrying_client.return_value = mock_client.return_value
         mock_client.return_value.set_many.return_value = []
 
-        result = self.object_store.put_many(["a", "b"])
+        result = self.object_store.put_many(["a", "b", "c"])
 
-        self.assertEqual(2, len(result))
+        self.assertEqual(3, len(result))
         self.assertRegex(result[0], ".*_.*")
-        self.assertEqual(1, mock_client.return_value.set_many.call_count)
 
     @mock.patch("deltacat.io.memcached_object_store.Client")
     @mock.patch("deltacat.io.memcached_object_store.RetryingClient")

deltacat/tests/local_deltacat_storage/__init__.py

@@ -181,16 +181,22 @@ def list_deltas(
     partition = get_partition(stream.locator, partition_values, *args, **kwargs)
 
     all_deltas = list_partition_deltas(
-        partition, include_manifest, *args, **kwargs
+        partition,
+        first_stream_position=first_stream_position,
+        last_stream_position=last_stream_position,
+        ascending_order=ascending_order,
+        include_manifest=include_manifest,
+        *args,
+        **kwargs,
     ).all_items()
 
     result = []
 
     for delta in all_deltas:
         if (
-            not first_stream_position or first_stream_position <= delta.stream_position
+            not first_stream_position or first_stream_position < delta.stream_position
         ) and (
-            not last_stream_position or last_stream_position > delta.stream_position
+            not last_stream_position or delta.stream_position <= last_stream_position
         ):
             result.append(delta)
 
@@ -202,16 +208,38 @@ def list_deltas(
 
 
 def list_partition_deltas(
-    partition: Partition, include_manifest: bool = False, *args, **kwargs
+    partition_like: Union[Partition, PartitionLocator],
+    first_stream_position: Optional[int] = None,
+    last_stream_position: Optional[int] = None,
+    ascending_order: bool = False,
+    include_manifest: bool = False,
+    *args,
+    **kwargs,
 ) -> ListResult[Delta]:
     cur, con = _get_sqlite3_cursor_con(kwargs)
 
-    if partition is None:
+    if partition_like is None:
         return ListResult.of([], None, None)
 
+    if first_stream_position is None:
+        first_stream_position = 0
+
+    if last_stream_position is None:
+        last_stream_position = float("inf")
+
+    assert isinstance(partition_like, Partition) or isinstance(
+        partition_like, PartitionLocator
+    ), f"Expected a Partition or PartitionLocator as an input argument but found {partition_like}"
+
+    partition_locator = None
+    if isinstance(partition_like, Partition):
+        partition_locator = partition_like.locator
+    else:
+        partition_locator = partition_like
+
     res = cur.execute(
         "SELECT * FROM deltas WHERE partition_locator = ?",
-        (partition.locator.canonical_string(),),
+        (partition_locator.canonical_string(),),
     )
 
     serialized_items = res.fetchall()
@@ -222,12 +250,19 @@ def list_partition_deltas(
     result = []
     for item in serialized_items:
         current_delta = Delta(json.loads(item[2]))
-        result.append(current_delta)
+        if (
+            first_stream_position
+            <= current_delta.stream_position
+            <= last_stream_position
+        ):
+            result.append(current_delta)
 
         if not include_manifest:
             current_delta.manifest = None
 
-    result.sort(reverse=True, key=lambda d: d.stream_position)
+    result.sort(
+        reverse=True if not ascending_order else False, key=lambda d: d.stream_position
+    )
     return ListResult.of(result, None, None)
 
 
@@ -334,7 +369,6 @@ def download_delta_manifest_entry(
     cur, con = _get_sqlite3_cursor_con(kwargs)
 
     manifest = get_delta_manifest(delta_like, *args, **kwargs)
-
     if entry_index >= len(manifest.entries):
         raise IndexError(
             f"Manifest entry index {entry_index} does not exist. "
@@ -352,7 +386,6 @@ def download_delta_manifest_entry(
         )
 
     serialized_data = serialized_data[0]
-
     if entry.meta.content_type == ContentType.PARQUET:
         if table_type == TableType.PYARROW_PARQUET:
             table = pa.parquet.ParquetFile(io.BytesIO(serialized_data))
@@ -388,18 +421,17 @@ def download_delta_manifest_entry(
 
 def get_delta_manifest(
     delta_like: Union[Delta, DeltaLocator], *args, **kwargs
-) -> Manifest:
+) -> Optional[Manifest]:
     delta = get_delta(
-        delta_like.namespace,
-        delta_like.table_name,
-        delta_like.stream_position,
-        delta_like.partition_values,
-        delta_like.table_version,
-        True,
+        namespace=delta_like.namespace,
+        table_name=delta_like.table_name,
+        stream_position=delta_like.stream_position,
+        partition_values=delta_like.partition_values,
+        table_version=delta_like.table_version,
+        include_manifest=True,
         *args,
         **kwargs,
     )
-
     if not delta:
         return None
 
@@ -462,7 +494,6 @@ def create_table_version(
     cur, con = _get_sqlite3_cursor_con(kwargs)
 
     latest_version = get_latest_table_version(namespace, table_name, *args, **kwargs)
-
     if (
         table_version is not None
         and latest_version
@@ -762,7 +793,18 @@ def commit_partition(partition: Partition, *args, **kwargs) -> Partition:
         params = (json.dumps(pv_partition), pv_partition.locator.canonical_string())
         cur.execute("UPDATE partitions SET value = ? WHERE locator = ?", params)
 
+    deltas = list_partition_deltas(partition, *args, **kwargs).all_items()
+    deltas.sort(reverse=True, key=lambda x: x.stream_position)
+
+    stream_position = partition.stream_position
+    if deltas:
+        stream_position = deltas[0].stream_position
+
     partition.state = CommitState.COMMITTED
+    partition.stream_position = stream_position
+    partition.previous_stream_position = (
+        pv_partition.stream_position if pv_partition else None
+    )
     params = (json.dumps(partition), partition.locator.canonical_string())
     cur.execute("UPDATE partitions SET value = ? WHERE locator = ?", params)
     con.commit()
@@ -1032,6 +1074,7 @@ def get_stream(
     *args,
     **kwargs,
 ) -> Optional[Stream]:
+    assert not isinstance(table_version, int), f"Passed an integer as the table version"
     obj = get_table_version(namespace, table_name, table_version, *args, **kwargs)
 
     if obj is None:

deltacat/tests/test_utils/pyarrow.py

@@ -0,0 +1,32 @@
+from typing import List
+import pyarrow as pa
+from deltacat.storage import Delta
+import deltacat.tests.local_deltacat_storage as ds
+
+
+def create_delta_from_csv_file(
+    namespace: str, file_paths: List[str], *args, **kwargs
+) -> Delta:
+    tables = []
+
+    for file_path in file_paths:
+        table = pa.csv.read_csv(file_path)
+        tables.append(table)
+
+    ds.create_namespace(namespace, {}, **kwargs)
+    table_name = "-".join(file_paths).replace("/", "_")
+    ds.create_table_version(namespace, table_name, "1", **kwargs)
+    stream = ds.get_stream(namespace, table_name, "1", **kwargs)
+    staged_partition = ds.stage_partition(stream, [], **kwargs)
+
+    deltas = []
+
+    for table in tables:
+        delta = ds.stage_delta(table, staged_partition, **kwargs)
+        deltas.append(delta)
+
+    merged_delta = Delta.merge_deltas(deltas=deltas)
+    committed_delta = ds.commit_delta(merged_delta, **kwargs)
+    ds.commit_partition(staged_partition, **kwargs)
+
+    return committed_delta

deltacat/tests/test_utils/utils.py

@@ -0,0 +1,13 @@
+# Allow classes to use self-referencing Type hints in Python 3.7.
+from __future__ import annotations
+import json
+from typing import Any, Dict
+from boto3.resources.base import ServiceResource
+
+
+def read_s3_contents(
+    s3_resource: ServiceResource, bucket_name: str, key: str
+) -> Dict[str, Any]:
+    response = s3_resource.Object(bucket_name, key).get()
+    file_content: str = response["Body"].read().decode("utf-8")
+    return json.loads(file_content)

deltacat/tests/utils/test_daft.py

@@ -0,0 +1,76 @@
+import unittest
+from deltacat.types.media import ContentEncoding, ContentType
+from deltacat.utils.daft import daft_s3_file_to_table
+
+from deltacat.utils.pyarrow import ReadKwargsProviderPyArrowSchemaOverride
+from deltacat.types.partial_download import PartialParquetParameters
+import pyarrow as pa
+
+from pyarrow import parquet as pq
+
+
+class TestDaftParquetReader(unittest.TestCase):
+    MVP_PATH = "deltacat/tests/utils/data/mvp.parquet"
+
+    def test_read_from_s3_all_columns(self):
+        table = daft_s3_file_to_table(
+            self.MVP_PATH,
+            content_encoding=ContentEncoding.IDENTITY.value,
+            content_type=ContentType.PARQUET.value,
+        )
+        self.assertEqual(table.schema.names, ["a", "b"])
+        self.assertEqual(table.num_rows, 100)
+
+    def test_read_from_s3_single_column_via_include_columns(self):
+        table = daft_s3_file_to_table(
+            self.MVP_PATH,
+            content_encoding=ContentEncoding.IDENTITY.value,
+            content_type=ContentType.PARQUET.value,
+            include_columns=["b"],
+        )
+        self.assertEqual(table.schema.names, ["b"])
+        self.assertEqual(table.num_rows, 100)
+
+    def test_read_from_s3_single_column_via_column_names(self):
+        table = daft_s3_file_to_table(
+            self.MVP_PATH,
+            content_encoding=ContentEncoding.IDENTITY.value,
+            content_type=ContentType.PARQUET.value,
+            column_names=["b"],
+        )
+        self.assertEqual(table.schema.names, ["b"])
+        self.assertEqual(table.num_rows, 100)
+
+    def test_read_from_s3_single_column_with_schema(self):
+        schema = pa.schema([("a", pa.int64()), ("b", pa.string())])
+        pa_read_func_kwargs_provider = ReadKwargsProviderPyArrowSchemaOverride(
+            schema=schema
+        )
+        table = daft_s3_file_to_table(
+            self.MVP_PATH,
+            content_encoding=ContentEncoding.IDENTITY.value,
+            content_type=ContentType.PARQUET.value,
+            include_columns=["b"],
+            pa_read_func_kwargs_provider=pa_read_func_kwargs_provider,
+        )
+        self.assertEqual(table.schema.names, ["b"])
+        self.assertEqual(table.num_rows, 100)
+
+    def test_read_from_s3_single_column_with_row_groups(self):
+
+        metadata = pq.read_metadata(self.MVP_PATH)
+        ppp = PartialParquetParameters.of(pq_metadata=metadata)
+        ppp["row_groups_to_download"] = ppp.row_groups_to_download[1:2]
+        table = daft_s3_file_to_table(
+            self.MVP_PATH,
+            content_encoding=ContentEncoding.IDENTITY.value,
+            content_type=ContentType.PARQUET.value,
+            column_names=["b"],
+            partial_file_download_params=ppp,
+        )
+        self.assertEqual(table.schema.names, ["b"])
+        self.assertEqual(table.num_rows, 10)
+
+
+if __name__ == "__main__":
+    unittest.main()