deltacat 0.2.9__py3-none-any.whl → 1.0.0__py3-none-any.whl

deltacat/compute/compactor_v2/utils/merge.py

@@ -0,0 +1,126 @@
+import logging
+
+from deltacat.compute.compactor.model.compact_partition_params import (
+    CompactPartitionParams,
+)
+from deltacat.compute.compactor_v2.model.merge_file_group import (
+    LocalMergeFileGroupsProvider,
+)
+from deltacat.compute.compactor_v2.model.merge_input import MergeInput
+import pyarrow as pa
+from deltacat import logs
+from typing import List, Optional
+
+from deltacat.types.media import DELIMITED_TEXT_CONTENT_TYPES
+from deltacat.compute.compactor.model.materialize_result import MaterializeResult
+from deltacat.compute.compactor.model.pyarrow_write_result import PyArrowWriteResult
+from deltacat.compute.compactor import (
+    RoundCompletionInfo,
+    DeltaAnnotated,
+)
+
+from deltacat.types.tables import TABLE_CLASS_TO_SIZE_FUNC
+
+from deltacat.utils.performance import timed_invocation
+from deltacat.storage import (
+    Partition,
+)
+
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def materialize(
+    input: MergeInput,
+    task_index: int,
+    compacted_tables: List[pa.Table],
+) -> MaterializeResult:
+    compacted_table = pa.concat_tables(compacted_tables)
+    if input.compacted_file_content_type in DELIMITED_TEXT_CONTENT_TYPES:
+        # TODO (rkenmi): Investigate if we still need to convert this table to pandas DataFrame
+        # TODO (pdames): compare performance to pandas-native materialize path
+        df = compacted_table.to_pandas(split_blocks=True, self_destruct=True)
+        compacted_table = df
+    delta, stage_delta_time = timed_invocation(
+        input.deltacat_storage.stage_delta,
+        compacted_table,
+        input.write_to_partition,
+        max_records_per_entry=input.max_records_per_output_file,
+        content_type=input.compacted_file_content_type,
+        s3_table_writer_kwargs=input.s3_table_writer_kwargs,
+        **input.deltacat_storage_kwargs,
+    )
+    compacted_table_size = TABLE_CLASS_TO_SIZE_FUNC[type(compacted_table)](
+        compacted_table
+    )
+    logger.debug(
+        f"Time taken for materialize task"
+        f" to upload {len(compacted_table)} records"
+        f" of size {compacted_table_size} is: {stage_delta_time}s"
+    )
+    manifest = delta.manifest
+    manifest_records = manifest.meta.record_count
+    assert manifest_records == len(compacted_table), (
+        f"Unexpected Error: Materialized delta manifest record count "
+        f"({manifest_records}) does not equal compacted table record count "
+        f"({len(compacted_table)})"
+    )
+    materialize_result = MaterializeResult.of(
+        delta=delta,
+        task_index=task_index,
+        # TODO (pdames): Generalize WriteResult to contain in-memory-table-type
+        #  and in-memory-table-bytes instead of tight coupling to paBytes
+        pyarrow_write_result=PyArrowWriteResult.of(
+            len(manifest.entries),
+            TABLE_CLASS_TO_SIZE_FUNC[type(compacted_table)](compacted_table),
+            manifest.meta.content_length,
+            len(compacted_table),
+        ),
+    )
+    logger.info(f"Materialize result: {materialize_result}")
+    return materialize_result
+
+
+def generate_local_merge_input(
+    params: CompactPartitionParams,
+    annotated_deltas: List[DeltaAnnotated],
+    compacted_partition: Partition,
+    round_completion_info: Optional[RoundCompletionInfo],
+):
+    """
+    Generates a merge input for local deltas that do not reside in the Ray object store and
+    have not been subject to the hash bucketing process.
+
+    Args:
+        params: parameters for compacting a partition
+        annotated_deltas: a list of annotated deltas
+        compacted_partition: the compacted partition to write to
+        round_completion_info: keeps track of high watermarks and other metadata from previous compaction rounds
+
+    Returns:
+        A MergeInput object
+
+    """
+
+    return MergeInput.of(
+        merge_file_groups_provider=LocalMergeFileGroupsProvider(
+            annotated_deltas,
+            read_kwargs_provider=params.read_kwargs_provider,
+            deltacat_storage=params.deltacat_storage,
+            deltacat_storage_kwargs=params.deltacat_storage_kwargs,
+        ),
+        write_to_partition=compacted_partition,
+        compacted_file_content_type=params.compacted_file_content_type,
+        primary_keys=params.primary_keys,
+        sort_keys=params.sort_keys,
+        drop_duplicates=params.drop_duplicates,
+        max_records_per_output_file=params.records_per_compacted_file,
+        enable_profiler=params.enable_profiler,
+        metrics_config=params.metrics_config,
+        s3_table_writer_kwargs=params.s3_table_writer_kwargs,
+        read_kwargs_provider=params.read_kwargs_provider,
+        round_completion_info=round_completion_info,
+        object_store=params.object_store,
+        deltacat_storage=params.deltacat_storage,
+        deltacat_storage_kwargs=params.deltacat_storage_kwargs,
+    )

deltacat/compute/compactor_v2/utils/task_options.py

@@ -1,5 +1,7 @@
 import botocore
+import logging
 from typing import Dict, Optional, List, Tuple
+from deltacat import logs
 from deltacat.types.media import ContentEncoding, ContentType
 from deltacat.types.partial_download import PartialParquetParameters
 from deltacat.storage import (
@@ -17,6 +19,8 @@ from deltacat.compute.compactor_v2.constants import (
     PARQUET_TO_PYARROW_INFLATION,
 )
 
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
 
 def _get_parquet_type_params_if_exist(
     entry: ManifestEntry,
@@ -133,11 +137,15 @@ def hash_bucket_resource_options_provider(
     ray_custom_resources: Optional[Dict] = None,
     **kwargs,
 ) -> Dict:
+    debug_memory_params = {"hash_bucket_task_index": index}
     size_bytes = 0.0
     num_rows = 0
     total_pk_size = 0
 
     if not item.manifest or not item.manifest.entries:
+        logger.debug(
+            f"[Hash bucket task {index}]: No manifest entries, skipping memory allocation calculation"
+        )
         return {"CPU": 0.01}
 
     for entry in item.manifest.entries:
@@ -162,12 +170,30 @@ def hash_bucket_resource_options_provider(
             else:
                 total_pk_size += pk_size
 
-    # total size + pk size + pk hash column + hash bucket index column
+    # total size + pk size + pyarrow-to-numpy conversion + pk hash column + hashlib inefficiency + hash bucket index column
     # Refer to hash_bucket step for more details.
-    total_memory = size_bytes + total_pk_size + num_rows * 20 + num_rows * 4
+    total_memory = (
+        size_bytes
+        + total_pk_size
+        + total_pk_size
+        + num_rows * 20
+        + num_rows * 20
+        + num_rows * 4
+    )
+    debug_memory_params["size_bytes"] = size_bytes
+    debug_memory_params["num_rows"] = num_rows
+    debug_memory_params["total_pk_size"] = total_pk_size
+    debug_memory_params["total_memory"] = total_memory
+
+    debug_memory_params["previous_inflation"] = previous_inflation
+    debug_memory_params["average_record_size_bytes"] = average_record_size_bytes
 
     # Consider buffer
     total_memory = total_memory * (1 + TOTAL_MEMORY_BUFFER_PERCENTAGE / 100.0)
+    debug_memory_params["total_memory_with_buffer"] = total_memory
+    logger.debug(
+        f"[Hash bucket task {index}]: Params used for calculating hash bucketing memory: {debug_memory_params}"
+    )
 
     return get_task_options(0.01, total_memory, ray_custom_resources)
 
@@ -186,10 +212,13 @@ def merge_resource_options_provider(
     deltacat_storage_kwargs: Optional[Dict] = {},
     **kwargs,
 ) -> Dict:
+    debug_memory_params = {"merge_task_index": index}
     hb_group_idx = item[0]
 
     data_size = hash_group_size_bytes.get(hb_group_idx, 0)
     num_rows = hash_group_num_rows.get(hb_group_idx, 0)
+    debug_memory_params["data_size_from_hash_group"] = data_size
+    debug_memory_params["num_rows_from_hash_group"] = num_rows
 
     # upper bound for pk size of incremental
     pk_size_bytes = data_size
@@ -205,10 +234,13 @@ def merge_resource_options_provider(
             round_completion_info.compacted_pyarrow_write_result.pyarrow_bytes
             / round_completion_info.compacted_pyarrow_write_result.file_bytes
         )
+        debug_memory_params["previous_inflation"] = previous_inflation
+
         average_record_size = (
             round_completion_info.compacted_pyarrow_write_result.pyarrow_bytes
             / round_completion_info.compacted_pyarrow_write_result.records
         )
+        debug_memory_params["average_record_size"] = average_record_size
 
         iterable = hash_group_index_to_hash_bucket_indices(
             hb_group_idx, round_completion_info.hash_bucket_count, num_hash_groups
@@ -247,16 +279,27 @@ def merge_resource_options_provider(
                     else:
                         pk_size_bytes += pk_size
 
-    # total data downloaded + primary key hash column + primary key column
-    # + dict size for merge + incremental index array size
+    # total data downloaded + primary key hash column + pyarrow-to-numpy conversion
+    # + primary key column + hashlib inefficiency + dict size for merge + incremental index array size
     total_memory = (
         data_size
         + pk_size_bytes
+        + pk_size_bytes
+        + num_rows * 20
         + num_rows * 20
         + num_rows * 20
         + incremental_index_array_size
     )
+    debug_memory_params["data_size"] = data_size
+    debug_memory_params["num_rows"] = num_rows
+    debug_memory_params["pk_size_bytes"] = pk_size_bytes
+    debug_memory_params["incremental_index_array_size"] = incremental_index_array_size
+    debug_memory_params["total_memory"] = total_memory
 
     total_memory = total_memory * (1 + TOTAL_MEMORY_BUFFER_PERCENTAGE / 100.0)
+    debug_memory_params["total_memory_with_buffer"] = total_memory
+    logger.debug(
+        f"[Merge task {index}]: Params used for calculating merge memory: {debug_memory_params}"
+    )
 
     return get_task_options(0.01, total_memory, ray_custom_resources)

deltacat/compute/merge_on_read/__init__.py

@@ -0,0 +1,4 @@
+from deltacat.types.media import DistributedDatasetType
+from deltacat.compute.merge_on_read.daft import merge as daft_merge
+
+MERGE_FUNC_BY_DISTRIBUTED_DATASET_TYPE = {DistributedDatasetType.DAFT.value: daft_merge}

deltacat/compute/merge_on_read/daft.py

@@ -0,0 +1,40 @@
+import logging
+from deltacat.compute.merge_on_read.model.merge_on_read_params import MergeOnReadParams
+from deltacat.storage.model.types import DistributedDataset
+from deltacat.types.media import TableType, DistributedDatasetType
+from deltacat.compute.merge_on_read.utils.delta import create_df_from_all_deltas
+from deltacat import logs
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def merge(params: MergeOnReadParams, **kwargs) -> DistributedDataset:
+    """
+    Merges the given deltas and returns the result as distributed dataframe.
+    It reads the deltas into the Daft dataframe and leverages operations supported
+    by Daft to perform an efficient merge using Ray cluster.
+
+    TODO(raghumdani): Perform actual merge.
+    """
+
+    delta_dfs = create_df_from_all_deltas(
+        deltas=params.deltas,
+        table_type=TableType.PYARROW,
+        distributed_dataset_type=DistributedDatasetType.DAFT,
+        reader_kwargs=params.reader_kwargs,
+        deltacat_storage=params.deltacat_storage,
+        deltacat_storage_kwargs=params.deltacat_storage_kwargs,
+        **kwargs,
+    )
+
+    logger.info(f"Merging {len(delta_dfs)} delta dfs...")
+
+    # TODO: This code should be optimized from daft side
+    result = None
+    for df in delta_dfs:
+        if result is None:
+            result = df
+        else:
+            result = result.concat(df)
+
+    return result

deltacat/compute/merge_on_read/model/merge_on_read_params.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+from typing import Optional, Dict, List, Union, Any
+from deltacat.storage import (
+    Delta,
+    DeltaLocator,
+    interface as unimplemented_deltacat_storage,
+)
+
+
+class MergeOnReadParams(dict):
+    """
+    This class represents the parameters passed to compact_partition (deltacat/compute/compactor/compaction_session.py)
+    """
+
+    @staticmethod
+    def of(params: Optional[Dict]) -> MergeOnReadParams:
+        params = {} if params is None else params
+
+        result = MergeOnReadParams(params)
+        assert result.deltas is not None, "deltas is a required arg"
+
+        result.deltacat_storage = params.get(
+            "deltacat_storage", unimplemented_deltacat_storage
+        )
+        result.reader_kwargs = params.get("reader_kwargs", {})
+        result.deltacat_storage_kwargs = params.get("deltacat_storage_kwargs", {})
+
+        return result
+
+    @property
+    def deltas(self) -> List[Union[Delta, DeltaLocator]]:
+        """
+        The list of deltas to compact in-memory.
+        """
+        return self["deltas"]
+
+    @deltas.setter
+    def deltas(self, to_set: List[Union[Delta, DeltaLocator]]) -> None:
+        self["deltas"] = to_set
+
+    @property
+    def reader_kwargs(self) -> Dict[Any, Any]:
+        """
+        The key word arguments to be passed to the reader.
+        """
+        return self["reader_kwargs"]
+
+    @reader_kwargs.setter
+    def reader_kwargs(self, kwargs: Dict[Any, Any]) -> None:
+        self["reader_kwargs"] = kwargs
+
+    @property
+    def deltacat_storage(self) -> unimplemented_deltacat_storage:
+        return self["deltacat_storage"]
+
+    @deltacat_storage.setter
+    def deltacat_storage(self, storage: unimplemented_deltacat_storage) -> None:
+        self["deltacat_storage"] = storage
+
+    @property
+    def deltacat_storage_kwargs(self) -> dict:
+        return self["deltacat_storage_kwargs"]
+
+    @deltacat_storage_kwargs.setter
+    def deltacat_storage_kwargs(self, kwargs: dict) -> None:
+        self["deltacat_storage_kwargs"] = kwargs

deltacat/compute/merge_on_read/utils/delta.py

@@ -0,0 +1,42 @@
+from typing import List, Dict, Any, Optional, Union
+from deltacat.storage.model.delta import Delta, DeltaLocator
+from deltacat.storage.model.types import DistributedDataset
+from deltacat.storage import (
+    interface as unimplemented_deltacat_storage,
+)
+from deltacat.types.media import TableType, StorageType, DistributedDatasetType
+
+
+def create_df_from_all_deltas(
+    deltas: List[Union[Delta, DeltaLocator]],
+    table_type: TableType,
+    distributed_dataset_type: DistributedDatasetType,
+    reader_kwargs: Optional[Dict[Any, Any]] = None,
+    deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict[Any, Any]] = None,
+    *args,
+    **kwargs
+) -> List[DistributedDataset]:  # type: ignore
+    """
+    This method creates a distributed dataset for each delta and returns their references.
+    """
+
+    if reader_kwargs is None:
+        reader_kwargs = {}
+    if deltacat_storage_kwargs is None:
+        deltacat_storage_kwargs = {}
+
+    df_list = []
+
+    for delta in deltas:
+        df = deltacat_storage.download_delta(
+            delta_like=delta,
+            table_type=table_type,
+            distributed_dataset_type=distributed_dataset_type,
+            storage_type=StorageType.DISTRIBUTED,
+            **reader_kwargs,
+            **deltacat_storage_kwargs
+        )
+        df_list.append(df)
+
+    return df_list

deltacat/storage/interface.py

@@ -23,7 +23,12 @@ from deltacat.storage import (
     SortKey,
     PartitionLocator,
 )
-from deltacat.types.media import ContentType, StorageType, TableType
+from deltacat.types.media import (
+    ContentType,
+    StorageType,
+    TableType,
+    DistributedDatasetType,
+)
 from deltacat.utils.common import ReadKwargsProvider
 
 
@@ -178,9 +183,10 @@ def download_delta(
     columns: Optional[List[str]] = None,
     file_reader_kwargs_provider: Optional[ReadKwargsProvider] = None,
     ray_options_provider: Callable[[int, Any], Dict[str, Any]] = None,
+    distributed_dataset_type: DistributedDatasetType = DistributedDatasetType.RAY_DATASET,
     *args,
     **kwargs
-) -> Union[LocalDataset, DistributedDataset]:
+) -> Union[LocalDataset, DistributedDataset]:  # type: ignore
     """
     Download the given delta or delta locator into either a list of
     tables resident in the local node's memory, or into a dataset distributed
@@ -205,6 +211,8 @@ def download_delta_manifest_entry(
     given delta or delta locator. If a delta is provided with a non-empty
     manifest, then the entry is downloaded from this manifest. Otherwise, the
     manifest is first retrieved then the given entry index downloaded.
+
+    NOTE: The entry will be downloaded in the current node's memory.
     """
     raise NotImplementedError("download_delta_manifest_entry not implemented")
 

deltacat/storage/model/types.py

@@ -1,24 +1,16 @@
 from enum import Enum
-from typing import List, Union, Any
+from typing import List, Union
 
 from pyarrow.parquet import ParquetFile
 import numpy as np
 import pandas as pd
 import pyarrow as pa
-import pkg_resources
-from ray.data._internal.arrow_block import ArrowRow
 from ray.data.dataset import Dataset
+from daft import DataFrame as DaftDataFrame
 
 LocalTable = Union[pa.Table, pd.DataFrame, np.ndarray, ParquetFile]
 LocalDataset = List[LocalTable]
-# Starting Ray 2.5.0, Dataset follows a strict mode (https://docs.ray.io/en/latest/data/faq.html#migrating-to-strict-mode),
-# and generic annotation is removed. So add a version checker to determine whether to use the old or new definition.
-ray_version = pkg_resources.parse_version(pkg_resources.get_distribution("ray").version)
-change_version = pkg_resources.parse_version("2.5.0")
-if ray_version < change_version:
-    DistributedDataset = Dataset[Union[ArrowRow, np.ndarray, Any]]
-else:
-    DistributedDataset = Dataset
+DistributedDataset = Union[Dataset, DaftDataFrame]
 
 
 class DeltaType(str, Enum):

deltacat/tests/catalog/test_default_catalog_impl.py

@@ -0,0 +1,98 @@
+import unittest
+import sqlite3
+import ray
+import os
+import deltacat.tests.local_deltacat_storage as ds
+from deltacat.utils.common import current_time_ms
+from deltacat.tests.test_utils.pyarrow import (
+    create_delta_from_csv_file,
+    commit_delta_to_partition,
+)
+from deltacat.types.media import DistributedDatasetType, ContentType
+from deltacat.catalog import default_catalog_impl as dc
+
+
+class TestReadTable(unittest.TestCase):
+    READ_TABLE_NAMESPACE = "catalog_read_table_namespace"
+    LOCAL_CATALOG_NAME = "local_catalog"
+    DB_FILE_PATH = f"{current_time_ms()}.db"
+    SAMPLE_FILE_PATH = "deltacat/tests/catalog/data/sample_table.csv"
+
+    @classmethod
+    def setUpClass(cls):
+        ray.init(local_mode=True, ignore_reinit_error=True)
+
+        con = sqlite3.connect(cls.DB_FILE_PATH)
+        cur = con.cursor()
+        cls.kwargs = {
+            ds.SQLITE_CON_ARG: con,
+            ds.SQLITE_CUR_ARG: cur,
+            "supported_content_types": [ContentType.CSV],
+        }
+        cls.deltacat_storage_kwargs = {ds.DB_FILE_PATH_ARG: cls.DB_FILE_PATH}
+
+        super().setUpClass()
+
+    @classmethod
+    def doClassCleanups(cls) -> None:
+        os.remove(cls.DB_FILE_PATH)
+
+    def test_daft_distributed_read_sanity(self):
+        # setup
+        READ_TABLE_TABLE_NAME = "test_read_table"
+        create_delta_from_csv_file(
+            self.READ_TABLE_NAMESPACE,
+            [self.SAMPLE_FILE_PATH],
+            table_name=READ_TABLE_TABLE_NAME,
+            **self.kwargs,
+        )
+
+        dc.initialize(ds=ds)
+        df = dc.read_table(
+            table=READ_TABLE_TABLE_NAME,
+            namespace=self.READ_TABLE_NAMESPACE,
+            catalog=self.LOCAL_CATALOG_NAME,
+            distributed_dataset_type=DistributedDatasetType.DAFT,
+            deltacat_storage_kwargs=self.kwargs,
+        )
+
+        # verify
+        self.assertEqual(df.count_rows(), 6)
+        self.assertEqual(df.column_names, ["pk", "value"])
+
+    def test_daft_distributed_read_multiple_deltas(self):
+        # setup
+        READ_TABLE_TABLE_NAME = "test_read_table_2"
+        delta = create_delta_from_csv_file(
+            self.READ_TABLE_NAMESPACE,
+            [self.SAMPLE_FILE_PATH],
+            table_name=READ_TABLE_TABLE_NAME,
+            **self.kwargs,
+        )
+
+        partition = ds.get_partition(
+            delta.stream_locator, delta.partition_values, **self.kwargs
+        )
+
+        commit_delta_to_partition(
+            partition=partition, file_paths=[self.SAMPLE_FILE_PATH], **self.kwargs
+        )
+
+        # action
+        dc.initialize(ds=ds)
+        df = dc.read_table(
+            table=READ_TABLE_TABLE_NAME,
+            namespace=self.READ_TABLE_NAMESPACE,
+            catalog=self.LOCAL_CATALOG_NAME,
+            distributed_dataset_type=DistributedDatasetType.DAFT,
+            merge_on_read=False,
+            deltacat_storage_kwargs=self.kwargs,
+        )
+
+        # verify
+        self.assertEqual(
+            df.count_rows(),
+            12,
+            "we expect twice as many" " columns as merge on read is disabled",
+        )
+        self.assertEqual(df.column_names, ["pk", "value"])