deltacat 1.0.2__py3-none-any.whl → 1.1.1__py3-none-any.whl

deltacat/compute/compactor_v2/steps/merge.py

@@ -4,13 +4,15 @@ from deltacat.compute.compactor_v2.model.merge_input import MergeInput
 import numpy as np
 import pyarrow as pa
 import ray
+import itertools
 import time
 import pyarrow.compute as pc
 import deltacat.compute.compactor_v2.utils.merge as merge_utils
 from uuid import uuid4
 from deltacat import logs
-from typing import List, Optional, Tuple
+from typing import Callable, Iterator, List, Optional, Tuple
 from deltacat.compute.compactor_v2.model.merge_result import MergeResult
+from deltacat.compute.compactor_v2.model.merge_file_group import MergeFileGroup
 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, DeltaFileEnvelope
@@ -69,17 +71,11 @@ def _drop_delta_type_rows(table: pa.Table, delta_type: DeltaType) -> pa.Table:
 
 
 def _build_incremental_table(
-    df_envelopes_list: List[List[DeltaFileEnvelope]],
+    df_envelopes: List[DeltaFileEnvelope],
 ) -> pa.Table:
 
     hb_tables = []
     # sort by delta file stream position now instead of sorting every row later
-    df_envelopes = [d for dfe_list in df_envelopes_list for d in dfe_list]
-    df_envelopes = sorted(
-        df_envelopes,
-        key=lambda df: (df.stream_position, df.file_index),
-        reverse=False,  # ascending
-    )
     is_delete = False
     for df_envelope in df_envelopes:
         assert (
@@ -96,9 +92,7 @@ def _build_incremental_table(
             )
 
         hb_tables.append(table)
-
     result = pa.concat_tables(hb_tables)
-
     return result
 
 
@@ -111,7 +105,7 @@ def _merge_tables(
     """
     Merges the table with compacted table dropping duplicates where necessary.
 
-    This method ensures the appropriate deltas of types DELETE/UPSERT are correctly
+    This method ensures the appropriate deltas of types [UPSERT] are correctly
     appended to the table.
     """
 
@@ -214,7 +208,7 @@ def _copy_all_manifest_files_from_old_hash_buckets(
     hb_index_to_indices = round_completion_info.hb_index_to_entry_range
 
     if hb_index_to_indices is None:
-        logger.info(f"Nothing to copy by reference. Skipping...")
+        logger.info("Nothing to copy by reference. Skipping...")
         return []
 
     for hb_index in hb_index_copy_by_reference:
@@ -257,62 +251,209 @@ def _copy_all_manifest_files_from_old_hash_buckets(
     return materialize_result_list
 
 
+def _has_previous_compacted_table(input: MergeInput, hb_idx: int) -> bool:
+    """
+    Checks if the given hash bucket index has a compacted table available from the previous compaction round.
+
+    Args:
+        input (MergeInput): The input for the merge operation.
+        hb_idx (int): The hash bucket index to check.
+
+    Returns:
+        bool: True if the hash bucket index has a compacted table available, False otherwise.
+    """
+    return (
+        input.round_completion_info
+        and input.round_completion_info.hb_index_to_entry_range
+        and input.round_completion_info.hb_index_to_entry_range.get(str(hb_idx))
+        is not None
+    )
+
+
+def _can_copy_by_reference(
+    has_delete: bool, merge_file_group: MergeFileGroup, input: MergeInput
+) -> bool:
+    """
+    Can copy by reference only if there are no deletes to merge in
+    and previous compacted stream id matches that of new stream
+    """
+    return (
+        not has_delete
+        and not merge_file_group.dfe_groups
+        and input.round_completion_info is not None
+        and (
+            input.write_to_partition.stream_id
+            == input.round_completion_info.compacted_delta_locator.stream_id
+        )
+    )
+
+
+def _flatten_dfe_list(
+    df_envelopes_list: List[List[DeltaFileEnvelope]],
+) -> List[DeltaFileEnvelope]:
+    """
+    Flattens a list of lists of DeltaFileEnvelope objects into a single list of DeltaFileEnvelope objects.
+
+    Args:
+        df_envelopes_list (List[List[DeltaFileEnvelope]]): A list of lists of DeltaFileEnvelope objects.
+
+    Returns:
+        List[DeltaFileEnvelope]: A flattened list of DeltaFileEnvelope objects.
+    """
+    if not df_envelopes_list:
+        return []
+    return [d for dfe_list in df_envelopes_list for d in dfe_list]
+
+
+def _sort_df_envelopes(
+    df_envelopes: List[DeltaFileEnvelope],
+    key: Callable = lambda df: (df.stream_position, df.file_index),
+) -> List[DeltaFileEnvelope]:
+    """
+    Sorts a list of DeltaFileEnvelope objects based on a specified key function.
+
+    Args:
+        df_envelopes (List[DeltaFileEnvelope]): A list of DeltaFileEnvelope objects.
+        key (Callable, optional): A function that takes a DeltaFileEnvelope object and returns a key for sorting.
+            Defaults to lambda df: (df.stream_position, df.file_index).
+
+    Returns:
+        List[DeltaFileEnvelope]: A sorted list of DeltaFileEnvelope objects.
+    """
+    if not df_envelopes:
+        return []
+    return sorted(
+        df_envelopes,
+        key=key,
+        reverse=False,  # ascending
+    )
+
+
+def _group_sequence_by_delta_type(
+    df_envelopes: List[DeltaFileEnvelope],
+) -> Iterator[Tuple[List, List]]:
+    """
+    Groups a list of DeltaFileEnvelope objects by their delta_type.
+
+    Args:
+        df_envelopes (List[DeltaFileEnvelope]): A list of DeltaFileEnvelope objects.
+
+    Yields:
+        Iterator[Tuple[DeltaType, List[DeltaFileEnvelope]]]: A tuple containing the delta_type
+        and a list of DeltaFileEnvelope objects that share the same delta_type.
+    """
+    iter_df_envelopes = iter(df_envelopes)
+    for delta_type, delta_type_sequence in itertools.groupby(
+        iter_df_envelopes, lambda x: x.delta_type
+    ):
+        yield delta_type, list(delta_type_sequence)
+
+
 def _compact_tables(
-    input: MergeInput, dfe_list: List[List[DeltaFileEnvelope]], hb_idx: int
-) -> Tuple[pa.Table, int, int]:
+    input: MergeInput,
+    dfe_list: Optional[List[List[DeltaFileEnvelope]]],
+    hb_idx: int,
+    compacted_table: Optional[pa.Table] = None,
+) -> Tuple[pa.Table, int, int, int]:
+    """
+    Compacts a list of DeltaFileEnvelope objects into a single PyArrow table.
+
+    Args:
+        input (MergeInput): The input for the merge operation.
+        dfe_list (List[List[DeltaFileEnvelope]]): A list of lists of DeltaFileEnvelope objects.
+        hb_idx (int): The hash bucket index for the compaction.
+
+    Returns:
+        Tuple[pa.Table, int, int, int]: A tuple containing:
+            1. The compacted PyArrow table.
+            2. The total number of records in the incremental data.
+            3. The total number of deduplicated records.
+            4. The total number of deleted records due to DELETE operations.
+    """
+    df_envelopes: List[DeltaFileEnvelope] = _flatten_dfe_list(dfe_list)
+    delete_file_envelopes = input.delete_file_envelopes or []
+    reordered_all_dfes: List[DeltaFileEnvelope] = _sort_df_envelopes(
+        delete_file_envelopes + df_envelopes
+    )
+    assert all(
+        dfe.delta_type in (DeltaType.UPSERT, DeltaType.DELETE)
+        for dfe in reordered_all_dfes
+    ), "All reordered delta file envelopes must be of the UPSERT or DELETE"
+    table = compacted_table
+    aggregated_incremental_len = 0
+    aggregated_deduped_records = 0
+    aggregated_dropped_records = 0
+    for i, (delta_type, delta_type_sequence) in enumerate(
+        _group_sequence_by_delta_type(reordered_all_dfes)
+    ):
+        if delta_type is DeltaType.UPSERT:
+            (
+                table,
+                incremental_len,
+                deduped_records,
+                merge_time,
+            ) = _apply_upserts(input, delta_type_sequence, hb_idx, table)
+            logger.info(
+                f" [Merge task index {input.merge_task_index}] Merged"
+                f" record count: {len(table)}, size={table.nbytes} took: {merge_time}s"
+            )
+            aggregated_incremental_len += incremental_len
+            aggregated_deduped_records += deduped_records
+        elif delta_type is DeltaType.DELETE:
+            table_size_before_delete = len(table) if table else 0
+            (table, dropped_rows), delete_time = timed_invocation(
+                func=input.delete_strategy.apply_many_deletes,
+                table=table,
+                delete_file_envelopes=delta_type_sequence,
+            )
+            logger.info(
+                f" [Merge task index {input.merge_task_index}]"
+                + f" Dropped record count: {dropped_rows} from table"
+                + f" of record count {table_size_before_delete} took: {delete_time}s"
+            )
+            aggregated_dropped_records += dropped_rows
+    return (
+        table,
+        aggregated_incremental_len,
+        aggregated_deduped_records,
+        aggregated_dropped_records,
+    )
+
+
+def _apply_upserts(
+    input: MergeInput,
+    dfe_list: List[DeltaFileEnvelope],
+    hb_idx,
+    prev_table=None,
+) -> Tuple[pa.Table, int, int, int]:
+    assert all(
+        dfe.delta_type is DeltaType.UPSERT for dfe in dfe_list
+    ), "All incoming delta file envelopes must of the DeltaType.UPSERT"
     logger.info(
         f"[Hash bucket index {hb_idx}] Reading dedupe input for "
         f"{len(dfe_list)} delta file envelope lists..."
     )
     table = _build_incremental_table(dfe_list)
-
     incremental_len = len(table)
     logger.info(
         f"[Hash bucket index {hb_idx}] Got the incremental table of length {incremental_len}"
     )
-
     if input.sort_keys:
         # Incremental is sorted and merged, as sorting
         # on non event based sort key does not produce consistent
         # compaction results. E.g., compaction(delta1, delta2, delta3)
         # will not be equal to compaction(compaction(delta1, delta2), delta3).
         table = table.sort_by(input.sort_keys)
-
-    compacted_table = None
-
-    if (
-        input.round_completion_info
-        and input.round_completion_info.hb_index_to_entry_range
-        and input.round_completion_info.hb_index_to_entry_range.get(str(hb_idx))
-        is not None
-    ):
-        compacted_table = _download_compacted_table(
-            hb_index=hb_idx,
-            rcf=input.round_completion_info,
-            read_kwargs_provider=input.read_kwargs_provider,
-            deltacat_storage=input.deltacat_storage,
-            deltacat_storage_kwargs=input.deltacat_storage_kwargs,
-        )
-
-    hb_table_record_count = len(table) + (
-        len(compacted_table) if compacted_table else 0
-    )
-
+    hb_table_record_count = len(table) + (len(prev_table) if prev_table else 0)
     table, merge_time = timed_invocation(
         func=_merge_tables,
         table=table,
         primary_keys=input.primary_keys,
         can_drop_duplicates=input.drop_duplicates,
-        compacted_table=compacted_table,
+        compacted_table=prev_table,
     )
-    total_deduped_records = hb_table_record_count - len(table)
-
-    logger.info(
-        f"[Merge task index {input.merge_task_index}] Merged "
-        f"record count: {len(table)}, size={table.nbytes} took: {merge_time}s"
-    )
-
-    return table, incremental_len, total_deduped_records
+    deduped_records = hb_table_record_count - len(table)
+    return table, incremental_len, deduped_records, merge_time
 
 
 def _copy_manifests_from_hash_bucketing(
@@ -345,20 +486,47 @@ def _timed_merge(input: MergeInput) -> MergeResult:
         f"merge_{worker_id}_{task_id}.bin"
     ) if input.enable_profiler else nullcontext():
         total_input_records, total_deduped_records = 0, 0
+        total_dropped_records = 0
         materialized_results: List[MaterializeResult] = []
         merge_file_groups = input.merge_file_groups_provider.create()
         hb_index_copy_by_ref_ids = []
 
         for merge_file_group in merge_file_groups:
-            if not merge_file_group.dfe_groups:
+            compacted_table = None
+            has_delete = input.delete_file_envelopes is not None
+            if has_delete:
+                assert (
+                    input.delete_strategy is not None
+                ), "Merge input missing delete_strategy"
+            if _can_copy_by_reference(
+                has_delete=has_delete, merge_file_group=merge_file_group, input=input
+            ):
                 hb_index_copy_by_ref_ids.append(merge_file_group.hb_index)
                 continue
 
-            table, input_records, deduped_records = _compact_tables(
-                input, merge_file_group.dfe_groups, merge_file_group.hb_index
+            if _has_previous_compacted_table(input, merge_file_group.hb_index):
+                compacted_table = _download_compacted_table(
+                    hb_index=merge_file_group.hb_index,
+                    rcf=input.round_completion_info,
+                    read_kwargs_provider=input.read_kwargs_provider,
+                    deltacat_storage=input.deltacat_storage,
+                    deltacat_storage_kwargs=input.deltacat_storage_kwargs,
+                )
+            if not merge_file_group.dfe_groups and compacted_table is None:
+                logger.warning(
+                    f" [Hash bucket index {merge_file_group.hb_index}]"
+                    + f" No new deltas and no compacted table found. Skipping compaction for {merge_file_group.hb_index}"
+                )
+                continue
+            table, input_records, deduped_records, dropped_records = _compact_tables(
+                input,
+                merge_file_group.dfe_groups,
+                merge_file_group.hb_index,
+                compacted_table,
             )
             total_input_records += input_records
             total_deduped_records += deduped_records
+            total_dropped_records += dropped_records
             materialized_results.append(
                 merge_utils.materialize(input, merge_file_group.hb_index, [table])
             )
@@ -382,6 +550,7 @@ def _timed_merge(input: MergeInput) -> MergeResult:
             materialized_results,
             np.int64(total_input_records),
             np.int64(total_deduped_records),
+            np.int64(total_dropped_records),
             np.double(peak_memory_usage_bytes),
             np.double(0.0),
             np.double(time.time()),
@@ -400,7 +569,8 @@ def merge(input: MergeInput) -> MergeResult:
                 f"({process_util.max_memory/BYTES_PER_GIBIBYTE} GB)"
             )
 
-        process_util.schedule_callback(log_peak_memory, 10)
+        if input.memory_logs_enabled:
+            process_util.schedule_callback(log_peak_memory, 10)
 
         merge_result, duration = timed_invocation(func=_timed_merge, input=input)
 
@@ -420,6 +590,7 @@ def merge(input: MergeInput) -> MergeResult:
             merge_result[1],
             merge_result[2],
             merge_result[3],
+            merge_result[4],
             np.double(emit_metrics_time),
             merge_result[4],
         )

deltacat/compute/compactor_v2/utils/delta.py

@@ -5,7 +5,10 @@ from deltacat.compute.compactor import (
     DeltaAnnotated,
     DeltaFileEnvelope,
 )
-
+from deltacat.storage import (
+    Delta,
+)
+from deltacat.storage.model.delta import DeltaType
 from deltacat.storage import interface as unimplemented_deltacat_storage
 from deltacat.types.media import StorageType
 from deltacat.utils.common import ReadKwargsProvider
@@ -18,6 +21,13 @@ import logging
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
 
+def contains_delete_deltas(deltas: List[Delta]) -> bool:
+    for delta in deltas:
+        if delta.type is DeltaType.DELETE:
+            return True
+    return False
+
+
 def read_delta_file_envelopes(
     annotated_delta: DeltaAnnotated,
     read_kwargs_provider: Optional[ReadKwargsProvider],

deltacat/compute/compactor_v2/utils/merge.py

@@ -25,6 +25,12 @@ from deltacat.utils.performance import timed_invocation
 from deltacat.storage import (
     Partition,
 )
+from deltacat.compute.compactor_v2.deletes.delete_strategy import (
+    DeleteStrategy,
+)
+from deltacat.compute.compactor_v2.deletes.delete_file_envelope import (
+    DeleteFileEnvelope,
+)
 
 
 logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
@@ -86,6 +92,8 @@ def generate_local_merge_input(
     annotated_deltas: List[DeltaAnnotated],
     compacted_partition: Partition,
     round_completion_info: Optional[RoundCompletionInfo],
+    delete_strategy: Optional[DeleteStrategy] = None,
+    delete_file_envelopes: Optional[DeleteFileEnvelope] = None,
 ):
     """
     Generates a merge input for local deltas that do not reside in the Ray object store and
@@ -123,4 +131,6 @@ def generate_local_merge_input(
         object_store=params.object_store,
         deltacat_storage=params.deltacat_storage,
         deltacat_storage_kwargs=params.deltacat_storage_kwargs,
+        delete_strategy=delete_strategy,
+        delete_file_envelopes=delete_file_envelopes,
     )

deltacat/compute/compactor_v2/utils/task_options.py

@@ -1,7 +1,10 @@
 import botocore
 import logging
-from typing import Dict, Optional, List, Tuple
+from typing import Dict, Optional, List, Tuple, Any
 from deltacat import logs
+from deltacat.compute.compactor_v2.model.merge_file_group import (
+    LocalMergeFileGroupsProvider,
+)
 from deltacat.types.media import ContentEncoding, ContentType
 from deltacat.types.partial_download import PartialParquetParameters
 from deltacat.storage import (
@@ -15,7 +18,6 @@ from deltacat.compute.compactor_v2.utils.primary_key_index import (
     hash_group_index_to_hash_bucket_indices,
 )
 from deltacat.compute.compactor_v2.constants import (
-    TOTAL_MEMORY_BUFFER_PERCENTAGE,
     PARQUET_TO_PYARROW_INFLATION,
 )
 
@@ -133,8 +135,10 @@ def hash_bucket_resource_options_provider(
     item: DeltaAnnotated,
     previous_inflation: float,
     average_record_size_bytes: float,
+    total_memory_buffer_percentage: int,
     primary_keys: List[str] = None,
     ray_custom_resources: Optional[Dict] = None,
+    memory_logs_enabled: Optional[bool] = None,
     **kwargs,
 ) -> Dict:
     debug_memory_params = {"hash_bucket_task_index": index}
@@ -189,10 +193,11 @@ def hash_bucket_resource_options_provider(
     debug_memory_params["average_record_size_bytes"] = average_record_size_bytes
 
     # Consider buffer
-    total_memory = total_memory * (1 + TOTAL_MEMORY_BUFFER_PERCENTAGE / 100.0)
+    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}"
+    logger.debug_conditional(
+        f"[Hash bucket task {index}]: Params used for calculating hash bucketing memory: {debug_memory_params}",
+        memory_logs_enabled,
     )
 
     return get_task_options(0.01, total_memory, ray_custom_resources)
@@ -204,12 +209,14 @@ def merge_resource_options_provider(
     num_hash_groups: int,
     hash_group_size_bytes: Dict[int, int],
     hash_group_num_rows: Dict[int, int],
+    total_memory_buffer_percentage: int,
     round_completion_info: Optional[RoundCompletionInfo] = None,
     compacted_delta_manifest: Optional[Manifest] = None,
     ray_custom_resources: Optional[Dict] = None,
     primary_keys: Optional[List[str]] = None,
     deltacat_storage=unimplemented_deltacat_storage,
     deltacat_storage_kwargs: Optional[Dict] = {},
+    memory_logs_enabled: Optional[bool] = None,
     **kwargs,
 ) -> Dict:
     debug_memory_params = {"merge_task_index": index}
@@ -224,6 +231,84 @@ def merge_resource_options_provider(
     pk_size_bytes = data_size
     incremental_index_array_size = num_rows * 4
 
+    return get_merge_task_options(
+        index,
+        hb_group_idx,
+        data_size,
+        pk_size_bytes,
+        num_rows,
+        num_hash_groups,
+        total_memory_buffer_percentage,
+        incremental_index_array_size,
+        debug_memory_params,
+        ray_custom_resources,
+        round_completion_info=round_completion_info,
+        compacted_delta_manifest=compacted_delta_manifest,
+        primary_keys=primary_keys,
+        deltacat_storage=deltacat_storage,
+        deltacat_storage_kwargs=deltacat_storage_kwargs,
+        memory_logs_enabled=memory_logs_enabled,
+    )
+
+
+def local_merge_resource_options_provider(
+    estimated_da_size: float,
+    estimated_num_rows: int,
+    total_memory_buffer_percentage: int,
+    round_completion_info: Optional[RoundCompletionInfo] = None,
+    compacted_delta_manifest: Optional[Manifest] = None,
+    ray_custom_resources: Optional[Dict] = None,
+    primary_keys: Optional[List[str]] = None,
+    deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict] = {},
+    memory_logs_enabled: Optional[bool] = None,
+    **kwargs,
+) -> Dict:
+    index = hb_group_idx = LocalMergeFileGroupsProvider.LOCAL_HASH_BUCKET_INDEX
+    debug_memory_params = {"merge_task_index": index}
+
+    # upper bound for pk size of incremental
+    pk_size_bytes = estimated_da_size
+    incremental_index_array_size = estimated_num_rows * 4
+
+    return get_merge_task_options(
+        index=index,
+        hb_group_idx=hb_group_idx,
+        data_size=estimated_da_size,
+        pk_size_bytes=pk_size_bytes,
+        num_rows=estimated_num_rows,
+        num_hash_groups=1,
+        incremental_index_array_size=incremental_index_array_size,
+        total_memory_buffer_percentage=total_memory_buffer_percentage,
+        debug_memory_params=debug_memory_params,
+        ray_custom_resources=ray_custom_resources,
+        round_completion_info=round_completion_info,
+        compacted_delta_manifest=compacted_delta_manifest,
+        primary_keys=primary_keys,
+        deltacat_storage=deltacat_storage,
+        deltacat_storage_kwargs=deltacat_storage_kwargs,
+        memory_logs_enabled=memory_logs_enabled,
+    )
+
+
+def get_merge_task_options(
+    index: int,
+    hb_group_idx: int,
+    data_size: float,
+    pk_size_bytes: float,
+    num_rows: int,
+    num_hash_groups: int,
+    total_memory_buffer_percentage: int,
+    incremental_index_array_size: int,
+    debug_memory_params: Dict[str, Any],
+    ray_custom_resources: Optional[Dict],
+    round_completion_info: Optional[RoundCompletionInfo] = None,
+    compacted_delta_manifest: Optional[Manifest] = None,
+    primary_keys: Optional[List[str]] = None,
+    deltacat_storage=unimplemented_deltacat_storage,
+    deltacat_storage_kwargs: Optional[Dict] = {},
+    memory_logs_enabled: Optional[bool] = None,
+) -> Dict[str, Any]:
     if (
         round_completion_info
         and compacted_delta_manifest
@@ -296,10 +381,11 @@ def merge_resource_options_provider(
     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)
+    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}"
+    logger.debug_conditional(
+        f"[Merge task {index}]: Params used for calculating merge memory: {debug_memory_params}",
+        memory_logs_enabled,
     )
 
     return get_task_options(0.01, total_memory, ray_custom_resources)

deltacat/io/memcached_object_store.py

@@ -181,15 +181,35 @@ class MemcachedObjectStore(IObjectStore):
         for chunk_index in range(chunk_count):
             ref = self._create_ref(uid, ip, chunk_index)
             chunk = client.get(ref)
+            if chunk is None:
+                raise ValueError(
+                    f"Expected uid: {uid}, chunk index: {chunk_index} from client ip: {ip}"
+                    f" to be non-empty."
+                )
             serialized.extend(chunk)
 
         return cloudpickle.loads(serialized)
 
+    def clear(self) -> bool:
+        flushed = all(
+            [
+                self._get_client_by_ip(ip).flush_all(noreply=False)
+                for ip in self.storage_node_ips
+            ]
+        )
+        self.client_cache.clear()
+
+        if flushed:
+            logger.info("Successfully cleared cache contents.")
+
+        return flushed
+
     def close(self) -> None:
         for client in self.client_cache.values():
             client.close()
 
         self.client_cache.clear()
+        logger.info("Successfully closed object store clients.")
 
     def _create_ref(self, uid, ip, chunk_index) -> str:
         return f"{uid}{self.SEPARATOR}{ip}{self.SEPARATOR}{chunk_index}"

deltacat/io/ray_plasma_object_store.py

@@ -2,6 +2,7 @@ import ray
 from ray import cloudpickle
 from deltacat.io.object_store import IObjectStore
 from typing import Any, List
+from ray.types import ObjectRef
 
 
 class RayPlasmaObjectStore(IObjectStore):
@@ -21,3 +22,8 @@ class RayPlasmaObjectStore(IObjectStore):
     def get_many(self, refs: List[Any], *args, **kwargs) -> List[object]:
         loaded_refs = [cloudpickle.loads(obj_id) for obj_id in refs]
         return ray.get(loaded_refs)
+
+    def deserialize_references(
+        self, refs: List[Any], *args, **kwargs
+    ) -> List[ObjectRef]:
+        return [cloudpickle.loads(obj_id) for obj_id in refs]