deltacat 2.0.0b9__py3-none-any.whl → 2.0.0b11__py3-none-any.whl

deltacat/tests/test_deltacat_api.py

@@ -1,29 +1,39 @@
 import shutil
 import tempfile
+
 import deltacat as dc
+from deltacat.constants import METAFILE_FORMAT_MSGPACK
+from deltacat import Namespace, DeltaCatUrl, DatasetType
+from deltacat.storage import Metafile
+
+from deltacat.io import (
+    METAFILE_TYPE_COLUMN_NAME,
+    METAFILE_DATA_COLUMN_NAME,
+)
 
 
 class TestDeltaCAT:
     @classmethod
-    def setup_class(cls):
+    def setup_method(cls):
         cls.temp_dir_1 = tempfile.mkdtemp()
         cls.temp_dir_2 = tempfile.mkdtemp()
         # Initialize DeltaCAT with two local catalogs.
-        dc.put("test_catalog_1", root=cls.temp_dir_1)
-        dc.put("test_catalog_2", root=cls.temp_dir_2)
+        dc.init()
+        dc.put(DeltaCatUrl("dc://test_catalog_1"), root=cls.temp_dir_1)
+        dc.put(DeltaCatUrl("dc://test_catalog_2"), root=cls.temp_dir_2)
 
     @classmethod
-    def teardown_class(cls):
+    def teardown_method(cls):
         shutil.rmtree(cls.temp_dir_1)
         shutil.rmtree(cls.temp_dir_2)
 
     def test_cross_catalog_namespace_copy(self):
         # Given two empty DeltaCAT catalogs.
         # When a namespace is copied across catalogs.
-        namespace_src = dc.put("test_catalog_1/test_namespace")
+        namespace_src = dc.put(DeltaCatUrl("dc://test_catalog_1/test_namespace"))
         namespace_dst = dc.copy(
-            "test_catalog_1/test_namespace",
-            "test_catalog_2",
+            DeltaCatUrl("dc://test_catalog_1/test_namespace"),
+            DeltaCatUrl("dc://test_catalog_2/test_namespace"),
         )
         # Expect the catalog namespace created in each catalog
         # method to be equivalent and equal to the source namespace.
@@ -33,7 +43,38 @@ class TestDeltaCAT:
         # When each catalog namespace is fetched explicitly
         # Expect them to be equivalent but not equal
         # (due to different metafile IDs).
-        actual_namespace_src = dc.get("test_catalog_1/test_namespace")
-        actual_namespace_dst = dc.get("test_catalog_2/test_namespace")
+        actual_namespace_src = dc.get(DeltaCatUrl("dc://test_catalog_1/test_namespace"))
+        actual_namespace_dst = dc.get(DeltaCatUrl("dc://test_catalog_2/test_namespace"))
         assert actual_namespace_src.equivalent_to(actual_namespace_dst)
         assert not actual_namespace_src == actual_namespace_dst
+
+    def test_catalog_listing_shallow_local_metafiles(self):
+        # Given two empty DeltaCAT catalogs.
+        # When a namespace is put in the catalog.
+        namespace_src: Namespace = dc.put(
+            DeltaCatUrl("dc://test_catalog_1/test_namespace")
+        )
+        # Expect the namespace to be listed.
+        assert any(
+            namespace_src.equivalent_to(other)
+            for other in dc.list(DeltaCatUrl("dc://test_catalog_1"))
+        )
+
+    def test_catalog_listing_shallow_ray_dataset(self):
+        # Given two empty DeltaCAT catalogs.
+        # When a namespace is put in the catalog.
+        namespace_src: Namespace = dc.put(
+            DeltaCatUrl("dc://test_catalog_1/test_namespace")
+        )
+        # Expect the namespace to be listed.
+        dataset = dc.list(
+            DeltaCatUrl("dc://test_catalog_1"),
+            dataset_type=DatasetType.RAY_DATASET,
+        )
+        actual_namespace = Metafile.deserialize(
+            serialized=dataset.take(1)[0][METAFILE_DATA_COLUMN_NAME],
+            meta_format=METAFILE_FORMAT_MSGPACK,
+        )
+        assert actual_namespace.equivalent_to(namespace_src)
+        namespace_type = dataset.take(1)[0][METAFILE_TYPE_COLUMN_NAME]
+        assert namespace_type == "Namespace"

deltacat/types/media.py

@@ -1,30 +1,48 @@
 from enum import Enum
-from typing import Dict, Set
+from typing import Set
 
 
 class ContentType(str, Enum):
-    # See also:
-    # https://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.17
-    # https://www.iana.org/assignments/media-types/media-types.xhtml
+    """
+    Enumeration used to resolve the entity-body Media Type (formerly known as
+    MIME type) in an HTTP request.
+
+    https://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.17
+
+    https://www.iana.org/assignments/media-types/media-types.xhtml
+    """
 
     # IANA registered types
+    AVRO = "application/avro"
+    BINARY = "application/octet-stream"
     CSV = "text/csv"
+    HDF = "application/x-hdf"
+    HTML = "text/html"
     JSON = "application/json"
+    TEXT = "text/plain"
+    WEBDATASET = "application/x-web-dataset"
+    XML = "text/xml"
 
     # unregistered types
-    TSV = "text/tsv"
-    PSV = "text/psv"
-    PARQUET = "application/parquet"
-    ORC = "application/orc"
     FEATHER = "application/feather"
-    UNESCAPED_TSV = "application/x-amzn-unescaped-tsv"
     ION = "application/x-amzn-ion"
+    ORC = "application/orc"
+    PARQUET = "application/parquet"
+    PSV = "text/psv"
+    TSV = "text/tsv"
+    UNESCAPED_TSV = "application/x-amzn-unescaped-tsv"
 
 
 class ContentEncoding(str, Enum):
-    # See also:
-    # https://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.11
-    # http://www.iana.org/assignments/http-parameters/http-parameters.xhtml#content-coding
+    """
+    Enumeration used as a modifier for :class:`deltacat.types.media.ContentType`
+    to indicate that additional encodings have been applied to the entity-body
+    Media Type in an HTTP request.
+
+    https://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.11
+
+    http://www.iana.org/assignments/http-parameters/http-parameters.xhtml#content-coding
+    """
 
     # IANA registered encodings
     GZIP = "gzip"
@@ -37,27 +55,6 @@ class ContentEncoding(str, Enum):
     SNAPPY = "snappy"
 
 
-class TableType(str, Enum):
-    PYARROW = "pyarrow"
-    PANDAS = "pandas"
-    NUMPY = "numpy"
-    PYARROW_PARQUET = "pyarrow_parquet"
-
-
-class DistributedDatasetType(str, Enum):
-    DAFT = "daft"
-    RAY_DATASET = "ray_dataset"
-
-
-class SchemaType(str, Enum):
-    ARROW = "arrow"
-
-
-class StorageType(str, Enum):
-    LOCAL = "local"
-    DISTRIBUTED = "distributed"
-
-
 DELIMITED_TEXT_CONTENT_TYPES: Set[str] = {
     ContentType.UNESCAPED_TSV.value,
     ContentType.TSV.value,
@@ -73,6 +70,7 @@ TABULAR_CONTENT_TYPES: Set[str] = {
     ContentType.PARQUET.value,
     ContentType.ORC.value,
     ContentType.FEATHER.value,
+    ContentType.AVRO.value,
 }
 
 EXPLICIT_COMPRESSION_CONTENT_TYPES: Set[str] = {
@@ -83,13 +81,113 @@ EXPLICIT_COMPRESSION_CONTENT_TYPES: Set[str] = {
     ContentType.JSON.value,
 }
 
-CONTENT_TYPE_TO_USER_KWARGS_KEY: Dict[str, str] = {
-    ContentType.UNESCAPED_TSV.value: "unescaped_tsv",
-    ContentType.TSV.value: "csv",
-    ContentType.CSV.value: "csv",
-    ContentType.PSV.value: "csv",
-    ContentType.PARQUET.value: "parquet",
-    ContentType.FEATHER.value: "feather",
-    ContentType.ORC.value: "orc",
-    ContentType.JSON.value: "json",
-}
+
+class DatasetType(str, Enum):
+    """
+    Enumeration used to identify the in-memory local or distributed dataset
+    to be used for file IO, queries, and data transformation. Typically used
+    together with :class:`deltacat.types.media.DatastoreType` to resolve the
+    compute layer that will be responsible for reading, transforming, and
+    writing data to a given datastore.
+    """
+
+    # local
+    NUMPY = "numpy"  # numpy.ndarray
+    PANDAS = "pandas"  # pandas.DataFrame
+    POLARS = "polars"  # polars.DataFrame
+    PYARROW = "pyarrow"  # pyarrow.Table
+    PYARROW_PARQUET = "pyarrow_parquet"  # pyarrow.parquet.ParquetFile
+
+    # distributed
+    DAFT = "daft"  # daft.DataFrame
+    RAY_DATASET = "ray_dataset"  # ray.data.Dataset
+
+    @staticmethod
+    def distributed():
+        return {
+            DatasetType.DAFT,
+            DatasetType.RAY_DATASET,
+        }
+
+    @staticmethod
+    def local():
+        return {
+            DatasetType.NUMPY,
+            DatasetType.PANDAS,
+            DatasetType.POLARS,
+            DatasetType.PYARROW,
+            DatasetType.PYARROW_PARQUET,
+        }
+
+
+# deprecated by DatasetType - populated dynamically for backwards compatibility
+TableType = Enum(
+    "TableType",
+    {d.name: d.value for d in DatasetType.local()},
+)
+
+# deprecated by DatasetType - populated dynamically for backwards compatibility
+DistributedDatasetType = Enum(
+    "DistributedDatasetType",
+    {d.name: d.value for d in DatasetType.distributed()},
+)
+
+
+# deprecated by DatasetType.local() and DatasetType.distributed()
+# kept for backwards compatibility
+class StorageType(str, Enum):
+    LOCAL = "local"
+    DISTRIBUTED = "distributed"
+
+
+class DatastoreType(str, Enum):
+    """
+    Enumeration used to identify the type of reader required to connect to and
+    correctly interpret data stored at a given path. Typically used together
+    with :class:`deltacat.types.media.DatasetType` to resolve a reader or
+    writer for that data store. Note that, although some overlap exists between
+    enum values here and in :class:`deltacat.types.media.ContentType`, each
+    enum serve a different purpose. The purpose of
+    :class:`deltacat.types.media.ContentType` is to resolve a file's MIME type,
+    and may be used together with datastores that support storing different
+    file types to describe the specific file type read/written from/to that
+    datastore (e.g., Iceberg, Hudi, Delta Lake, Audio, Images, Video, etc.)
+    """
+
+    # DeltaCAT Catalog Datasets
+    DELTACAT = "dc"
+    DELTACAT_NAMESPACE = "namespace"
+    DELTACAT_TABLE = "table"
+    DELTACAT_TABLE_VERSION = "tableversion"
+    DELTACAT_STREAM = "stream"
+    DELTACAT_PARTITION = "partition"
+    DELTACAT_DELTA = "delta"
+
+    # External Datasets
+    AUDIO = "audio"
+    AVRO = "avro"
+    BIGQUERY = "bigquery"
+    BINARY = "binary"
+    CSV = "csv"
+    CLICKHOUSE = "clickhouse"
+    DATABRICKS_TABLES = "databricks"
+    DELTA_LAKE = "deltalake"
+    DELTA_SHARING = "deltasharing"
+    FEATHER = "feather"
+    HDF = "hdf"
+    HTML = "html"
+    HUDI = "hudi"
+    ICEBERG = "iceberg"
+    IMAGES = "images"
+    JSON = "json"
+    LANCE = "lance"
+    MONGO = "mongodb"
+    NUMPY = "numpy"
+    ORC = "orc"
+    PARQUET = "parquet"
+    TEXT = "text"
+    TFRECORDS = "tfrecords"
+    VIDEOS = "videos"
+    WARC = "warc"
+    WEBDATASET = "webdataset"
+    XML = "xml"

deltacat/types/tables.py

@@ -3,9 +3,10 @@ from typing import Callable, Dict, Type, Union
 
 import numpy as np
 import pandas as pd
+import polars as pl
 import pyarrow as pa
 import pyarrow.parquet as papq
-from ray.data.dataset import Dataset
+from ray.data.dataset import Dataset as RayDataset
 from ray.data.read_api import (
     from_arrow,
     from_arrow_refs,
@@ -18,11 +19,12 @@ import deltacat.storage as dcs
 from deltacat.types.media import TableType, DistributedDatasetType
 from deltacat.utils import numpy as np_utils
 from deltacat.utils import pandas as pd_utils
+from deltacat.utils import polars as pl_utils
 from deltacat.utils import pyarrow as pa_utils
 from deltacat.utils import daft as daft_utils
 from deltacat.utils.ray_utils import dataset as ds_utils
 
-TABLE_TYPE_TO_READER_FUNC: Dict[int, Callable] = {
+TABLE_TYPE_TO_S3_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,
@@ -34,8 +36,9 @@ TABLE_CLASS_TO_WRITER_FUNC: Dict[
 ] = {
     pa.Table: pa_utils.table_to_file,
     pd.DataFrame: pd_utils.dataframe_to_file,
+    pl.DataFrame: pl_utils.dataframe_to_file,
     np.ndarray: np_utils.ndarray_to_file,
-    Dataset: ds_utils.dataset_to_file,
+    RayDataset: ds_utils.dataset_to_file,
 }
 
 TABLE_CLASS_TO_SLICER_FUNC: Dict[
@@ -43,8 +46,9 @@ TABLE_CLASS_TO_SLICER_FUNC: Dict[
 ] = {
     pa.Table: pa_utils.slice_table,
     pd.DataFrame: pd_utils.slice_dataframe,
+    pl.DataFrame: pl_utils.slice_table,
     np.ndarray: np_utils.slice_ndarray,
-    Dataset: ds_utils.slice_dataset,
+    RayDataset: ds_utils.slice_dataset,
 }
 
 TABLE_CLASS_TO_SIZE_FUNC: Dict[
@@ -53,13 +57,27 @@ TABLE_CLASS_TO_SIZE_FUNC: Dict[
     pa.Table: pa_utils.table_size,
     papq.ParquetFile: pa_utils.parquet_file_size,
     pd.DataFrame: pd_utils.dataframe_size,
+    pl.DataFrame: pl_utils.dataframe_size,
     np.ndarray: np_utils.ndarray_size,
-    Dataset: ds_utils.dataset_size,
+    RayDataset: ds_utils.dataset_size,
+}
+
+TABLE_CLASS_TO_PYARROW_FUNC: Dict[
+    Type[Union[dcs.LocalTable, dcs.DistributedDataset]], Callable
+] = {
+    pa.Table: lambda table, **kwargs: table,
+    papq.ParquetFile: lambda table, **kwargs: table.read(**kwargs),
+    pd.DataFrame: lambda table, **kwargs: pa.Table.from_pandas(table, **kwargs),
+    pl.DataFrame: lambda table, **kwargs: pl.DataFrame.to_arrow(table, **kwargs),
+    np.ndarray: lambda table, **kwargs: pa.Table.from_arrays(
+        [pa.array(table[:, i]) for i in range(table.shape[1])]
+    ),
 }
 
 TABLE_CLASS_TO_TABLE_TYPE: Dict[Type[dcs.LocalTable], str] = {
     pa.Table: TableType.PYARROW.value,
     papq.ParquetFile: TableType.PYARROW_PARQUET.value,
+    pl.DataFrame: TableType.POLARS.value,
     pd.DataFrame: TableType.PANDAS.value,
     np.ndarray: TableType.NUMPY.value,
 }
@@ -78,7 +96,6 @@ TABLE_TYPE_TO_DATASET_CREATE_FUNC_REFS: Dict[str, Callable] = {
     TableType.PANDAS.value: from_pandas_refs,
 }
 
-
 DISTRIBUTED_DATASET_TYPE_TO_READER_FUNC: Dict[int, Callable] = {
     DistributedDatasetType.DAFT.value: daft_utils.s3_files_to_dataframe
 }
@@ -106,7 +123,18 @@ class TableWriteMode(str, Enum):
 
 
 def get_table_length(table: Union[dcs.LocalTable, dcs.DistributedDataset]) -> int:
-    return len(table) if not isinstance(table, Dataset) else table.count()
+    return len(table) if not isinstance(table, RayDataset) else table.count()
+
+
+def get_table_size(table: Union[dcs.LocalTable, dcs.DistributedDataset]) -> int:
+    table_size_func = TABLE_CLASS_TO_SIZE_FUNC.get(type(table))
+    if table_size_func is None:
+        msg = (
+            f"No size function found for table type: {type(table)}.\n"
+            f"Known table types: {TABLE_CLASS_TO_SIZE_FUNC.keys}"
+        )
+        raise ValueError(msg)
+    return table_size_func(table)
 
 
 def get_table_writer(table: Union[dcs.LocalTable, dcs.DistributedDataset]) -> Callable: