deltacat 2.0.0b10__py3-none-any.whl → 2.0.0b12__py3-none-any.whl

deltacat/types/media.py

@@ -1,11 +1,14 @@
 from enum import Enum
-from typing import Set
+from typing import Set, Dict
 
 
 class ContentType(str, Enum):
     """
-    Enumeration used to resolve the entity-body Media Type (formerly known as
-    MIME type) in an HTTP request.
+    Enumeration used to resolve a file's entity-body Media Type (formerly known
+    as MIME type). All content types here are writeable by at least one
+    :class:`deltacat.types.media.DatasetType`. The Media Type is used as the
+    content type of each :class:`deltacat.storage.model.manifest.ManifestEntry`
+    written by that dataset type.
 
     https://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.17
 
@@ -16,16 +19,10 @@ class ContentType(str, Enum):
     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
     FEATHER = "application/feather"
-    ION = "application/x-amzn-ion"
     ORC = "application/orc"
     PARQUET = "application/parquet"
     PSV = "text/psv"
@@ -37,7 +34,7 @@ class ContentEncoding(str, Enum):
     """
     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.
+    Media Type.
 
     https://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.11
 
@@ -55,6 +52,46 @@ class ContentEncoding(str, Enum):
     SNAPPY = "snappy"
 
 
+# Map of file extensions to content types
+EXT_TO_CONTENT_TYPE: Dict[str, ContentType] = {
+    ".parquet": ContentType.PARQUET,
+    ".pq": ContentType.PARQUET,
+    ".csv": ContentType.CSV,
+    ".tsv": ContentType.TSV,
+    ".psv": ContentType.PSV,
+    ".json": ContentType.JSON,
+    ".feather": ContentType.FEATHER,
+    ".avro": ContentType.AVRO,
+    ".orc": ContentType.ORC,
+}
+
+# Inverse map of content types to file extensions
+CONTENT_TYPE_TO_EXT: Dict[ContentType, str] = {
+    v: k for k, v in EXT_TO_CONTENT_TYPE.items()
+}
+
+# Map of file extensions to content encodings
+EXT_TO_CONTENT_ENCODING: Dict[str, ContentEncoding] = {
+    ".gz": ContentEncoding.GZIP,
+    ".bz2": ContentEncoding.BZIP2,
+    ".zst": ContentEncoding.ZSTD,
+    ".sz": ContentEncoding.SNAPPY,
+    ".zz": ContentEncoding.DEFLATE,
+    ".zip": ContentEncoding.DEFLATE,
+}
+
+# Inverse map of content encodings to file extensions
+CONTENT_ENCODING_TO_EXT: Dict[ContentEncoding, str] = {
+    v: k for k, v in EXT_TO_CONTENT_ENCODING.items()
+}
+
+SCHEMA_CONTENT_TYPES: Set[str] = {
+    ContentType.PARQUET.value,
+    ContentType.ORC.value,
+    ContentType.FEATHER.value,
+    ContentType.AVRO.value,
+}
+
 DELIMITED_TEXT_CONTENT_TYPES: Set[str] = {
     ContentType.UNESCAPED_TSV.value,
     ContentType.TSV.value,
@@ -119,6 +156,158 @@ class DatasetType(str, Enum):
             DatasetType.PYARROW_PARQUET,
         }
 
+    def readable_content_types(self) -> Set[ContentType]:
+        # if this is DAFT then it can read PARQUET, JSON, and CSV
+        if self == DatasetType.DAFT:
+            return {
+                ContentType.PARQUET,
+                ContentType.JSON,
+                ContentType.CSV,
+                ContentType.PSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+            }
+        if self == DatasetType.RAY_DATASET:
+            return {
+                ContentType.CSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+                ContentType.PSV,
+                ContentType.PARQUET,
+                ContentType.JSON,
+                ContentType.AVRO,
+                ContentType.ORC,
+                ContentType.FEATHER,
+            }
+        if self == DatasetType.PYARROW:
+            return {
+                ContentType.CSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+                ContentType.PSV,
+                ContentType.PARQUET,
+                ContentType.FEATHER,
+                ContentType.JSON,
+                ContentType.AVRO,
+                ContentType.ORC,
+            }
+        if self == DatasetType.PANDAS:
+            return {
+                ContentType.CSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+                ContentType.PSV,
+                ContentType.PARQUET,
+                ContentType.FEATHER,
+                ContentType.JSON,
+                ContentType.AVRO,
+                ContentType.ORC,
+            }
+        if self == DatasetType.POLARS:
+            return {
+                ContentType.CSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+                ContentType.PSV,
+                ContentType.PARQUET,
+                ContentType.FEATHER,
+                ContentType.JSON,
+                ContentType.AVRO,
+                ContentType.ORC,
+            }
+        if self == DatasetType.NUMPY:
+            return {
+                ContentType.CSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+                ContentType.PSV,
+                ContentType.PARQUET,
+                ContentType.FEATHER,
+                ContentType.JSON,
+                ContentType.AVRO,
+                ContentType.ORC,
+            }
+        if self == DatasetType.PYARROW_PARQUET:
+            return {ContentType.PARQUET}
+        raise ValueError(f"No readable content types for {self}")
+
+    def writable_content_types(self) -> Set[ContentType]:
+        if self == DatasetType.PYARROW:
+            return {
+                ContentType.CSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+                ContentType.PSV,
+                ContentType.PARQUET,
+                ContentType.FEATHER,
+                ContentType.JSON,
+                ContentType.AVRO,
+                ContentType.ORC,
+            }
+        if self == DatasetType.PANDAS:
+            return {
+                ContentType.CSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+                ContentType.PSV,
+                ContentType.PARQUET,
+                ContentType.FEATHER,
+                ContentType.JSON,
+                ContentType.AVRO,
+                ContentType.ORC,
+            }
+        if self == DatasetType.POLARS:
+            return {
+                ContentType.CSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+                ContentType.PSV,
+                ContentType.PARQUET,
+                ContentType.FEATHER,
+                ContentType.JSON,
+                ContentType.AVRO,
+                ContentType.ORC,
+            }
+        if self == DatasetType.RAY_DATASET:
+            return {
+                ContentType.CSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+                ContentType.PSV,
+                ContentType.PARQUET,
+                ContentType.JSON,
+            }
+        if self == DatasetType.DAFT:
+            return {
+                ContentType.CSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+                ContentType.PSV,
+                ContentType.PARQUET,
+                ContentType.JSON,
+            }
+        if self == DatasetType.NUMPY:
+            return {
+                ContentType.CSV,
+                ContentType.TSV,
+                ContentType.UNESCAPED_TSV,
+                ContentType.PSV,
+                ContentType.PARQUET,
+                ContentType.FEATHER,
+                ContentType.JSON,
+                ContentType.AVRO,
+                ContentType.ORC,
+            }
+        if self == DatasetType.PYARROW_PARQUET:
+            return {}
+        raise ValueError(f"No writable content types for {self}")
+
+    def can_read(self, content_type: ContentType) -> bool:
+        return content_type in self.readable_content_types()
+
+    def can_write(self, content_type: ContentType) -> bool:
+        return content_type in self.writable_content_types()
+
 
 # deprecated by DatasetType - populated dynamically for backwards compatibility
 TableType = Enum(
@@ -140,6 +329,26 @@ class StorageType(str, Enum):
     DISTRIBUTED = "distributed"
 
 
+DATASET_TYPE_TO_SUPPORTED_READ_CONTENT_TYPES: Dict[DatasetType, Set[str]] = {
+    DatasetType.DAFT: {
+        ContentType.CSV,
+        ContentType.PARQUET,
+        ContentType.JSON,
+    },
+    DatasetType.RAY_DATASET: {
+        ContentType.CSV,
+        ContentType.TSV,
+        ContentType.UNESCAPED_TSV,
+        ContentType.PSV,
+        ContentType.PARQUET,
+        ContentType.JSON,
+        ContentType.AVRO,
+        ContentType.ORC,
+        ContentType.FEATHER,
+    },
+}
+
+
 class DatastoreType(str, Enum):
     """
     Enumeration used to identify the type of reader required to connect to and
@@ -148,10 +357,11 @@ class DatastoreType(str, Enum):
     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 the MIME type for
-    specific types of files, and may be used together with multi-content-type
-    datastore types to describe the specific file types read/written to that
-    datastore (e.g., Iceberg, Hudi, Delta Lake, Audio, Images, Video, etc.)
+    :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., DeltaCAT, Iceberg, Hudi, Delta Lake, Audio, Images, Video,
+    etc.)
     """
 
     # DeltaCAT Catalog Datasets