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

deltacat/tests/utils/test_pyarrow.py

@@ -0,0 +1,133 @@
+from unittest import TestCase
+from deltacat.utils.pyarrow import (
+    s3_parquet_file_to_table,
+    s3_partial_parquet_file_to_table,
+)
+from deltacat.types.media import ContentEncoding, ContentType
+from deltacat.types.partial_download import PartialParquetParameters
+from pyarrow.parquet import ParquetFile
+import pyarrow as pa
+
+PARQUET_FILE_PATH = "deltacat/tests/utils/data/test_file.parquet"
+
+
+class TestS3ParquetFileToTable(TestCase):
+    def test_s3_parquet_file_to_table_sanity(self):
+
+        result = s3_parquet_file_to_table(
+            PARQUET_FILE_PATH,
+            ContentType.PARQUET.value,
+            ContentEncoding.IDENTITY.value,
+            ["n_legs", "animal"],
+            ["n_legs"],
+        )
+
+        self.assertEqual(len(result), 6)
+        self.assertEqual(len(result.column_names), 1)
+        schema = result.schema
+        schema_index = schema.get_field_index("n_legs")
+        self.assertEqual(schema.field(schema_index).type, "int64")
+
+    def test_s3_parquet_file_to_table_when_schema_overridden(self):
+
+        schema = pa.schema(
+            [pa.field("animal", pa.string()), pa.field("n_legs", pa.string())]
+        )
+
+        pa_kwargs_provider = lambda content_type, kwargs: {"schema": schema}
+
+        result = s3_parquet_file_to_table(
+            PARQUET_FILE_PATH,
+            ContentType.PARQUET.value,
+            ContentEncoding.IDENTITY.value,
+            ["n_legs", "animal"],
+            pa_read_func_kwargs_provider=pa_kwargs_provider,
+        )
+
+        self.assertEqual(len(result), 6)
+        self.assertEqual(len(result.column_names), 2)
+
+        result_schema = result.schema
+        for index, field in enumerate(result_schema):
+            self.assertEqual(field.name, schema.field(index).name)
+
+        self.assertEqual(result.schema.field(1).type, "string")
+
+
+class TestS3PartialParquetFileToTable(TestCase):
+    def test_s3_partial_parquet_file_to_table_sanity(self):
+
+        pq_file = ParquetFile(PARQUET_FILE_PATH)
+        partial_parquet_params = PartialParquetParameters.of(
+            pq_metadata=pq_file.metadata
+        )
+
+        self.assertEqual(
+            partial_parquet_params.num_row_groups, 2, "test_file.parquet has changed."
+        )
+
+        # only first row group to be downloaded
+        partial_parquet_params.row_groups_to_download.pop()
+
+        result = s3_partial_parquet_file_to_table(
+            PARQUET_FILE_PATH,
+            include_columns=["n_legs"],
+            content_encoding=ContentEncoding.IDENTITY.value,
+            content_type=ContentType.PARQUET.value,
+            partial_file_download_params=partial_parquet_params,
+        )
+
+        self.assertEqual(len(result), 3)
+        self.assertEqual(len(result.columns), 1)
+
+    def test_s3_partial_parquet_file_to_table_when_schema_passed(self):
+
+        pq_file = ParquetFile(PARQUET_FILE_PATH)
+        partial_parquet_params = PartialParquetParameters.of(
+            pq_metadata=pq_file.metadata
+        )
+        # only first row group to be downloaded
+        partial_parquet_params.row_groups_to_download.pop()
+
+        schema = pa.schema(
+            [pa.field("animal", pa.string()), pa.field("n_legs", pa.string())]
+        )
+
+        pa_kwargs_provider = lambda content_type, kwargs: {"schema": schema}
+
+        result = s3_partial_parquet_file_to_table(
+            PARQUET_FILE_PATH,
+            ContentType.PARQUET.value,
+            ContentEncoding.IDENTITY.value,
+            ["n_legs", "animal"],
+            pa_read_func_kwargs_provider=pa_kwargs_provider,
+            partial_file_download_params=partial_parquet_params,
+        )
+
+        self.assertEqual(len(result), 3)
+        self.assertEqual(len(result.column_names), 2)
+
+        result_schema = result.schema
+        self.assertEqual(result_schema.field(0).type, "string")
+        self.assertEqual(result_schema.field(0).name, "n_legs")  # order doesn't change
+
+    def test_s3_partial_parquet_file_to_table_when_multiple_row_groups(self):
+
+        pq_file = ParquetFile(PARQUET_FILE_PATH)
+        partial_parquet_params = PartialParquetParameters.of(
+            pq_metadata=pq_file.metadata
+        )
+
+        self.assertEqual(
+            partial_parquet_params.num_row_groups, 2, "test_file.parquet has changed."
+        )
+
+        result = s3_partial_parquet_file_to_table(
+            PARQUET_FILE_PATH,
+            content_encoding=ContentEncoding.IDENTITY.value,
+            content_type=ContentType.PARQUET.value,
+            partial_file_download_params=partial_parquet_params,
+        )
+
+        self.assertEqual(len(result), 6)
+        self.assertEqual(len(result.columns), 2)

deltacat/tests/utils/test_resources.py

@@ -1,37 +1,22 @@
 import unittest
 from unittest import mock
-import sys
+import time
 
 
 class TestGetCurrentClusterUtilization(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.ray_mock = mock.MagicMock()
-        cls.ray_mock.cluster_resources.return_value = {
+    @mock.patch("deltacat.utils.resources.ray")
+    def test_sanity(self, ray_mock):
+        ray_mock.cluster_resources.return_value = {
             "CPU": 10,
             "memory": 10,
             "object_store_memory": 5,
         }
-        cls.ray_mock.available_resources.return_value = {
+        ray_mock.available_resources.return_value = {
             "CPU": 6,
             "memory": 4,
             "object_store_memory": 5,
         }
 
-        cls.module_patcher = mock.patch.dict("sys.modules", {"ray": cls.ray_mock})
-        cls.module_patcher.start()
-
-        # delete reference to reload from mocked ray
-        if "deltacat.utils.resources" in sys.modules:
-            del sys.modules["deltacat.utils.resources"]
-
-        super().setUpClass()
-
-    @classmethod
-    def tearDownClass(cls) -> None:
-        cls.module_patcher.stop()
-
-    def test_sanity(self):
         from deltacat.utils.resources import ClusterUtilization
 
         result = ClusterUtilization.get_current_cluster_utilization()
@@ -43,3 +28,21 @@ class TestGetCurrentClusterUtilization(unittest.TestCase):
         self.assertEqual(0, result.used_object_store_memory_bytes)
         self.assertEqual(6, result.used_memory_bytes)
         self.assertIsNotNone(result.used_resources)
+
+
+class TestClusterUtilizationOverTimeRange(unittest.TestCase):
+    @mock.patch("deltacat.utils.resources.ray")
+    def test_sanity(self, ray_mock):
+        from deltacat.utils.resources import ClusterUtilizationOverTimeRange
+
+        ray_mock.cluster_resources.side_effect = [{"CPU": 32} for _ in range(5)]
+        ray_mock.available_resources.side_effect = [
+            {"CPU": 2 ** (i + 1)} for i in range(5)
+        ]
+
+        with ClusterUtilizationOverTimeRange() as cu:
+            time.sleep(3)
+            self.assertTrue(cu.used_vcpu_seconds <= 82)  # 30 + 28 + 24
+            self.assertTrue(
+                cu.total_vcpu_seconds >= cu.used_vcpu_seconds
+            )  # total is greater than used

deltacat/types/media.py

@@ -41,6 +41,7 @@ class TableType(str, Enum):
     PYARROW = "pyarrow"
     PANDAS = "pandas"
     NUMPY = "numpy"
+    PYARROW_PARQUET = "pyarrow_parquet"
 
 
 class SchemaType(str, Enum):

deltacat/types/partial_download.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from typing import Dict, Any, Optional, List
+from pyarrow.parquet import FileMetaData
+
+
+class PartialFileDownloadParams(Dict[str, Any]):
+    """
+    A content type params class used to represent arguments required
+    to down the file partially. This is useful specifically in cases
+    where you'd like to instruct downloader to partially download a
+    manifest entry.
+    """
+
+    pass
+
+
+class PartialParquetParameters(PartialFileDownloadParams):
+    @staticmethod
+    def of(
+        row_groups_to_download: Optional[List[int]] = None,
+        num_row_groups: Optional[int] = None,
+        num_rows: Optional[int] = None,
+        in_memory_size_bytes: Optional[float] = None,
+        pq_metadata: Optional[FileMetaData] = None,
+    ) -> PartialParquetParameters:
+
+        if (
+            row_groups_to_download is None
+            or num_row_groups is None
+            or num_rows is None
+            or in_memory_size_bytes is None
+        ):
+            assert (
+                pq_metadata is not None
+            ), "Parquet file metadata must be passed explicitly"
+
+            num_row_groups = pq_metadata.num_row_groups
+            row_groups_to_download = [rg for rg in range(num_row_groups)]
+            in_memory_size_bytes = 0.0
+
+            for rg in row_groups_to_download:
+                row_group_meta = pq_metadata.row_group(rg)
+                in_memory_size_bytes += row_group_meta.total_byte_size
+
+        result = PartialParquetParameters(
+            {
+                "row_groups_to_download": row_groups_to_download,
+                "num_row_groups": num_row_groups,
+                "num_rows": num_rows,
+                "in_memory_size_bytes": in_memory_size_bytes,
+            }
+        )
+
+        if pq_metadata:
+            result["pq_metadata"] = pq_metadata
+
+        return result
+
+    @property
+    def row_groups_to_download(self) -> List[int]:
+        return self["row_groups_to_download"]
+
+    @property
+    def num_row_groups(self) -> List[int]:
+        return self["num_row_groups"]
+
+    @property
+    def num_rows(self) -> int:
+        return self["num_rows"]
+
+    @property
+    def in_memory_size_bytes(self) -> float:
+        return self["in_memory_size_bytes"]
+
+    @property
+    def pq_metadata(self) -> Optional[FileMetaData]:
+        return self.get("pq_metadata")
+
+    @pq_metadata.setter
+    def pq_metadata(self, metadata: FileMetaData) -> None:
+        self["pq_metadata"] = metadata

deltacat/types/tables.py

@@ -21,6 +21,7 @@ from deltacat.utils import pyarrow as pa_utils
 from deltacat.utils.ray_utils import dataset as ds_utils
 
 TABLE_TYPE_TO_READER_FUNC: Dict[int, Callable] = {
+    TableType.PYARROW_PARQUET.value: pa_utils.s3_file_to_parquet,
     TableType.PYARROW.value: pa_utils.s3_file_to_table,
     TableType.PANDAS.value: pd_utils.s3_file_to_dataframe,
     TableType.NUMPY.value: np_utils.s3_file_to_ndarray,

deltacat/utils/arguments.py

@@ -0,0 +1,26 @@
+import inspect
+import copy
+from typing import Any, Dict
+
+
+def sanitize_kwargs_to_callable(callable: Any, kwargs: Dict) -> Dict:
+    """
+    This method removes any upsupported key word arguments if variable
+    kwargs are not allowed in the method signature.
+
+    Returns: a sanitized dict of kwargs.
+    """
+    signature = inspect.signature(callable)
+    params = signature.parameters
+
+    new_kwargs = copy.copy(kwargs)
+
+    for key in params:
+        if params[key].kind == inspect.Parameter.VAR_KEYWORD:
+            return kwargs
+
+    for key in kwargs.keys():
+        if key not in params:
+            new_kwargs.pop(key)
+
+    return new_kwargs

deltacat/utils/daft.py

@@ -0,0 +1,87 @@
+import logging
+from typing import Optional, List
+
+from daft.table import Table
+from daft.logical.schema import Schema
+from daft import TimeUnit
+from daft.io import IOConfig, S3Config
+import pyarrow as pa
+
+from deltacat import logs
+from deltacat.utils.common import ReadKwargsProvider
+
+from deltacat.types.media import ContentType, ContentEncoding
+from deltacat.aws.constants import BOTO_MAX_RETRIES
+from deltacat.utils.performance import timed_invocation
+
+from deltacat.types.partial_download import (
+    PartialFileDownloadParams,
+)
+
+
+logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
+
+
+def daft_s3_file_to_table(
+    s3_url: str,
+    content_type: str,
+    content_encoding: str,
+    column_names: Optional[List[str]] = None,
+    include_columns: Optional[List[str]] = None,
+    pa_read_func_kwargs_provider: Optional[ReadKwargsProvider] = None,
+    partial_file_download_params: Optional[PartialFileDownloadParams] = None,
+    **s3_client_kwargs,
+):
+    assert (
+        content_type == ContentType.PARQUET.value
+    ), f"daft native reader currently only supports parquet, got {content_type}"
+
+    assert (
+        content_encoding == ContentEncoding.IDENTITY.value
+    ), f"daft native reader currently only supports identity encoding, got {content_encoding}"
+
+    kwargs = {}
+    if pa_read_func_kwargs_provider is not None:
+        kwargs = pa_read_func_kwargs_provider(content_type, kwargs)
+
+    coerce_int96_timestamp_unit = TimeUnit.from_str(
+        kwargs.get("coerce_int96_timestamp_unit", "ms")
+    )
+
+    row_groups = None
+    if (
+        partial_file_download_params is not None
+        and partial_file_download_params.row_groups_to_download is not None
+    ):
+        row_groups = partial_file_download_params.row_groups_to_download
+
+    io_config = IOConfig(
+        s3=S3Config(
+            key_id=s3_client_kwargs.get("aws_access_key_id"),
+            access_key=s3_client_kwargs.get("aws_secret_access_key"),
+            session_token=s3_client_kwargs.get("aws_session_token"),
+            retry_mode="adaptive",
+            num_tries=BOTO_MAX_RETRIES,
+        )
+    )
+
+    table, latency = timed_invocation(
+        Table.read_parquet,
+        path=s3_url,
+        columns=include_columns or column_names,
+        row_groups=row_groups,
+        io_config=io_config,
+        coerce_int96_timestamp_unit=coerce_int96_timestamp_unit,
+        multithreaded_io=False,
+    )
+
+    logger.debug(f"Time to read S3 object from {s3_url} into daft table: {latency}s")
+
+    if kwargs.get("schema") is not None:
+        schema = kwargs["schema"]
+        if include_columns is not None:
+            schema = pa.schema([schema.field(col) for col in include_columns])
+        daft_schema = Schema.from_pyarrow_schema(schema)
+        return table.cast_to_schema(daft_schema).to_arrow()
+    else:
+        return table.to_arrow()

deltacat/utils/placement.py

@@ -21,9 +21,10 @@ logger = logs.configure_deltacat_logger(logging.getLogger(__name__))
 
 @dataclass
 class PlacementGroupConfig:
-    def __init__(self, opts, resource):
+    def __init__(self, opts, resource, node_ips):
         self.opts = opts
         self.resource = resource
+        self.node_ips = node_ips
 
 
 class NodeGroupManager:
@@ -207,6 +208,7 @@ class PlacementGroupManager:
         cpu_per_bundle: int,
         strategy="SPREAD",
         capture_child_tasks=True,
+        memory_per_bundle=None,
     ):
         head_res_key = self.get_current_node_resource_key()
         # run the task on head and consume a fractional cpu, so that pg can be created on non-head node
@@ -216,7 +218,11 @@ class PlacementGroupManager:
         self._pg_configs = ray.get(
             [
                 _config.options(resources={head_res_key: 0.01}).remote(
-                    total_cpus_per_pg, cpu_per_bundle, strategy, capture_child_tasks
+                    total_cpus_per_pg,
+                    cpu_per_bundle,
+                    strategy,
+                    capture_child_tasks,
+                    memory_per_bundle=memory_per_bundle,
                 )
                 for i in range(num_pgs)
             ]
@@ -251,12 +257,18 @@ def _config(
     strategy="SPREAD",
     capture_child_tasks=True,
     time_out: Optional[float] = None,
+    memory_per_bundle: Optional[float] = None,
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     pg_config = None
     opts = {}
     cluster_resources = {}
     num_bundles = (int)(total_cpus_per_pg / cpu_per_node)
     bundles = [{"CPU": cpu_per_node} for i in range(num_bundles)]
+
+    if memory_per_bundle:
+        for bundle in bundles:
+            bundle["memory"] = memory_per_bundle
+
     pg = placement_group(bundles, strategy=strategy)
     ray.get(pg.ready(), timeout=time_out)
     if not pg:
@@ -275,6 +287,7 @@ def _config(
     # query available resources given list of node id
     all_nodes_available_res = ray._private.state.state._available_resources_per_node()
     pg_res = {"CPU": 0, "memory": 0, "object_store_memory": 0}
+    node_ips = []
     for node_id in node_ids:
         if node_id in all_nodes_available_res:
             v = all_nodes_available_res[node_id]
@@ -282,10 +295,14 @@ def _config(
             pg_res["CPU"] += node_detail["resources_total"]["CPU"]
             pg_res["memory"] += v["memory"]
             pg_res["object_store_memory"] += v["object_store_memory"]
+            node_ips.append(node_detail["node_ip"])
     cluster_resources["CPU"] = int(pg_res["CPU"])
     cluster_resources["memory"] = float(pg_res["memory"])
     cluster_resources["object_store_memory"] = float(pg_res["object_store_memory"])
-    pg_config = PlacementGroupConfig(opts, cluster_resources)
+    pg_config = PlacementGroupConfig(
+        opts=opts, resource=cluster_resources, node_ips=node_ips
+    )
     logger.info(f"pg has resources:{cluster_resources}")
+    logger.debug(f"pg has node ips:{node_ips}")
 
     return pg_config