tracdap-runtime 0.8.0rc2__py3-none-any.whl → 0.9.0b1__py3-none-any.whl

tracdap/rt/_impl/core/data.py

@@ -16,10 +16,12 @@
 import abc
 import copy
 import dataclasses as dc
+import pathlib
 import typing as tp
 import datetime as dt
 import decimal
 import platform
+import random
 
 import pyarrow as pa
 import pyarrow.compute as pc
@@ -36,8 +38,20 @@ except ModuleNotFoundError:
 
 import tracdap.rt.api.experimental as _api
 import tracdap.rt.metadata as _meta
+import tracdap.rt.config as _cfg
 import tracdap.rt.exceptions as _ex
 import tracdap.rt._impl.core.logging as _log
+import tracdap.rt._impl.core.util as _util
+
+
+@dc.dataclass(frozen=True)
+class DataPartKey:
+
+    @classmethod
+    def for_root(cls) -> "DataPartKey":
+        return DataPartKey(opaque_key='part-root')
+
+    opaque_key: str
 
 
 @dc.dataclass(frozen=True)
@@ -47,61 +61,470 @@ class DataSpec:
     schema_type: _meta.SchemaType
     data_item: str
 
-    data_def: _meta.DataDefinition
-    file_def: _meta.FileDefinition
-    storage_def: _meta.StorageDefinition
-    schema_def: tp.Optional[_meta.SchemaDefinition]
+    definition: tp.Union[_meta.DataDefinition, _meta.FileDefinition]
+    storage: _meta.StorageDefinition
+    schema: tp.Optional[_meta.SchemaDefinition] = None
+
+    primary_id: _meta.TagHeader = None
+    storage_id: _meta.TagHeader = None
+    schema_id: tp.Optional[_meta.TagHeader] = None
+    context_key: tp.Optional[str] = None
 
     @staticmethod
     def create_data_spec(
             data_item: str,
-            data_def: _meta.DataDefinition,
-            storage_def: _meta.StorageDefinition,
-            schema_def: tp.Optional[_meta.SchemaDefinition] = None) -> "DataSpec":
-
-        if schema_def:
-            schema_type = schema_def.schemaType
-        elif data_def.schema:
-            schema_type = data_def.schema.schemaType
+            definition: _meta.DataDefinition,
+            storage: _meta.StorageDefinition,
+            schema: tp.Optional[_meta.SchemaDefinition] = None) -> "DataSpec":
+
+        if schema:
+            schema_type = schema.schemaType
+        elif definition.schema:
+            schema_type = definition.schema.schemaType
         else:
             schema_type = _meta.SchemaType.SCHEMA_TYPE_NOT_SET
 
         return DataSpec(
             _meta.ObjectType.DATA, schema_type, data_item,
-            data_def,
-            storage_def=storage_def,
-            schema_def=schema_def,
-            file_def=None)
+            definition, storage, schema)
 
     @staticmethod
     def create_file_spec(
             data_item: str,
-            file_def: _meta.FileDefinition,
-            storage_def: _meta.StorageDefinition) -> "DataSpec":
+            definition: _meta.FileDefinition,
+            storage: _meta.StorageDefinition) -> "DataSpec":
 
         return DataSpec(
             _meta.ObjectType.FILE, _meta.SchemaType.SCHEMA_TYPE_NOT_SET, data_item,
-            file_def=file_def,
-            storage_def=storage_def,
-            data_def=None,
-            schema_def=None)
+            definition, storage)
 
     @staticmethod
     def create_empty_spec(object_type: _meta.ObjectType, schema_type: _meta.SchemaType):
-        return DataSpec(object_type, schema_type, None, None, None, None, None)
+        return DataSpec(object_type, schema_type, "", None, None, None)  # noqa
+
+    def with_ids(
+            self, primary_id: _meta.TagHeader,
+            storage_id: _meta.TagHeader,
+            schema_id: tp.Optional[_meta.TagHeader] = None,
+            context_key: tp.Optional[str] = None):
+
+        return DataSpec(
+            self.object_type, self.schema_type, self.data_item,
+            self.definition, self.storage, self.schema,
+            primary_id, storage_id, schema_id, context_key)
 
     def is_empty(self):
         return self.data_item is None or len(self.data_item) == 0
 
 
-@dc.dataclass(frozen=True)
-class DataPartKey:
+class StorageLayout:
+
+    __LAYOUTS: "tp.Dict[_meta.StorageLayout, StorageLayout]" = dict()
 
     @classmethod
-    def for_root(cls) -> "DataPartKey":
-        return DataPartKey(opaque_key='part_root')
+    def select(cls, layout_key: _meta.StorageLayout) -> "StorageLayout":
 
-    opaque_key: str
+        # Legacy compatibility - layout key not set in storage definition
+        if not layout_key or layout_key.value == 0:
+            layout_key = _meta.StorageLayout.OBJECT_ID_LAYOUT
+
+        layout = cls.__LAYOUTS.get(layout_key)
+
+        if layout is not None:
+            return layout
+
+        if layout_key == _meta.StorageLayout.OBJECT_ID_LAYOUT:
+            layout = ObjectIdLayout()
+        elif layout_key == _meta.StorageLayout.DEVELOPER_LAYOUT:
+            layout = DevelopmentLayout()
+        else:
+            raise _ex.ETracInternal(f"Unknown storage layout [{layout_key.name}]")
+
+        cls.__LAYOUTS[layout_key] = layout
+
+        return layout
+
+    @abc.abstractmethod
+    def layout_key(self) -> _meta.StorageLayout:
+        pass
+
+    @abc.abstractmethod
+    def new_data_spec(
+            self, data_id: _meta.TagHeader, storage_id: _meta.TagHeader,
+            context_key: tp.Optional[str], trac_schema: _meta.SchemaDefinition,
+            storage_config: _cfg.StorageConfig) -> DataSpec:
+        pass
+
+    @abc.abstractmethod
+    def new_data_version(
+            self, data_id: _meta.TagHeader, storage_id: _meta.TagHeader,
+            context_key: tp.Optional[str], trac_schema: _meta.SchemaDefinition,
+            prior_spec: DataSpec) -> DataSpec:
+        pass
+
+    @abc.abstractmethod
+    def new_file_spec(
+            self, file_id: _meta.TagHeader, storage_id: _meta.TagHeader,
+            context_key: str, file_type: _meta.FileType,
+            storage_config: _cfg.StorageConfig) -> DataSpec:
+        pass
+
+    @abc.abstractmethod
+    def new_file_version(
+            self, file_id: _meta.TagHeader, storage_id: _meta.TagHeader,
+            context_key: str, file_type: _meta.FileType,
+            prior_spec: DataSpec) -> DataSpec:
+        pass
+
+
+class BaseLayout(StorageLayout):
+
+    __DATA_ITEM_TEMPLATE = "data/{}/{}/{}/snap-{:d}/delta-{:d}"
+    __FILE_ITEM_TEMPLATE = "file/{}/version-{}"
+
+    @abc.abstractmethod
+    def _data_storage_path(
+            self, data_id: _meta.TagHeader, context_key: str, trac_schema: _meta.SchemaDefinition,
+            part_key: _meta.PartKey, snap_index: int, delta_index: int, storage_format: str,
+            prior_copy: tp.Optional[_meta.StorageCopy]):
+        pass
+
+    @abc.abstractmethod
+    def _file_storage_path(
+            self, file_id: _meta.TagHeader, file_def: _meta.FileDefinition,
+            prior_copy: tp.Optional[_meta.StorageCopy]):
+        pass
+
+    def new_data_spec(
+            self, data_id: _meta.TagHeader, storage_id: _meta.TagHeader,
+            context_key: tp.Optional[str], trac_schema: _meta.SchemaDefinition,
+            storage_config: _cfg.StorageConfig) -> DataSpec:
+
+        part_key = _meta.PartKey("part-root", _meta.PartType.PART_ROOT)
+        snap_index = 0
+
+        data_item = self.__DATA_ITEM_TEMPLATE.format(
+            trac_schema.schemaType.name.lower(), data_id.objectId,
+            part_key.opaqueKey, snap_index, 0)
+
+        # Blank data definition with no parts
+        new_data_def = _meta.DataDefinition(
+            schema=trac_schema,
+            storageId=_util.selector_for_latest(storage_id))
+
+        data_def = self._add_new_snap(new_data_def, data_item, part_key, snap_index)
+
+        # Take default location from the storage config
+        storage_key = storage_config.defaultBucket
+        storage_format = "JSON" if trac_schema.schemaType == _meta.SchemaType.STRUCT else storage_config.defaultFormat
+        storage_path = self._data_storage_path(data_id, context_key, trac_schema, part_key, snap_index, 0, storage_format, prior_copy=None)
+
+        storage_copy = _meta.StorageCopy(
+            storageKey=storage_key,
+            storagePath=storage_path,
+            storageFormat=storage_format,
+            copyStatus=_meta.CopyStatus.COPY_AVAILABLE,
+            copyTimestamp=data_id.objectTimestamp)
+
+        new_storage_def = _meta.StorageDefinition()
+
+        storage_def = self._add_storage_copy(new_storage_def, data_item, storage_copy)
+
+        # Dynamic data def will always use an embedded schema (this is no ID for an external schema)
+
+        return DataSpec \
+            .create_data_spec(data_item, data_def, storage_def, schema=None) \
+            .with_ids(data_id, storage_id)
+
+    def new_data_version(
+            self, data_id: _meta.TagHeader, storage_id: _meta.TagHeader,
+            context_key: tp.Optional[str], trac_schema: _meta.SchemaDefinition,
+            prior_spec: DataSpec) -> DataSpec:
+
+        part_key = _meta.PartKey("part-root", _meta.PartType.PART_ROOT)
+        snap_index = prior_spec.primary_id.objectVersion  # snap index is zero-based
+
+        data_item = self.__DATA_ITEM_TEMPLATE.format(
+            trac_schema.schemaType.name.lower(), data_id.objectId,
+            part_key.opaqueKey, snap_index, 0)
+
+        data_def = self._add_new_snap(prior_spec.definition, data_item, part_key, snap_index)
+
+        prior_item = next(iter(prior_spec.storage.dataItems.keys()), None)
+        prior_copy = self._find_storage_copy(prior_item, prior_spec.storage)
+
+        if prior_copy is None:
+            raise _ex.ETracInternal(f"Missing prior metadata for [{_util.object_key(data_id)}]")
+
+        storage_key = prior_copy.storageKey
+        storage_format = prior_copy.storageFormat
+        storage_path = self._data_storage_path(data_id, context_key, trac_schema, part_key, snap_index, 0, storage_format, prior_copy)
+
+        storage_copy = _meta.StorageCopy(
+            storageKey=storage_key,
+            storagePath=storage_path,
+            storageFormat=storage_format,
+            copyStatus=_meta.CopyStatus.COPY_AVAILABLE,
+            copyTimestamp=data_id.objectTimestamp)
+
+        storage_def = self._add_storage_copy(prior_spec.storage, data_item, storage_copy)
+
+        return DataSpec \
+            .create_data_spec(data_item, data_def, storage_def, schema=None) \
+            .with_ids(data_id, storage_id)
+
+    def new_file_spec(
+            self, file_id: _meta.TagHeader, storage_id: _meta.TagHeader,
+            context_key: str, file_type: _meta.FileType,
+            storage_config: _cfg.StorageConfig) -> DataSpec:
+
+        data_item = self.__FILE_ITEM_TEMPLATE.format(file_id.objectId, file_id.objectVersion)
+
+        file_def = _meta.FileDefinition(
+            name=f"{context_key}.{file_type.extension}",
+            extension=file_type.extension,
+            mimeType=file_type.mimeType,
+            dataItem=data_item,
+            storageId=_util.selector_for_latest(storage_id),
+            size=0)
+
+        storage_key = storage_config.defaultBucket
+        storage_format = file_def.mimeType
+        storage_path = self._file_storage_path(file_id, file_def, prior_copy=None)
+
+        storage_copy = _meta.StorageCopy(
+            storageKey=storage_key,
+            storagePath=storage_path,
+            storageFormat=storage_format,
+            copyStatus=_meta.CopyStatus.COPY_AVAILABLE,
+            copyTimestamp=file_id.objectTimestamp)
+
+        new_storage_def = _meta.StorageDefinition()
+        new_storage_def.layout = self.layout_key()
+
+        storage_def = self._add_storage_copy(new_storage_def, data_item, storage_copy)
+
+        return DataSpec \
+            .create_file_spec(data_item, file_def, storage_def) \
+            .with_ids(file_id, storage_id)
+
+    def new_file_version(
+            self, file_id: _meta.TagHeader, storage_id: _meta.TagHeader,
+            context_key: str, file_type: _meta.FileType, prior_spec: DataSpec) -> DataSpec:
+
+        data_item = self.__FILE_ITEM_TEMPLATE.format(file_id.objectId, file_id.objectVersion)
+
+        file_def = _meta.FileDefinition(
+            name=f"{context_key}.{file_type.extension}",
+            extension=file_type.extension,
+            mimeType=file_type.mimeType,
+            dataItem=data_item,
+            storageId=_util.selector_for_latest(storage_id),
+            size=0)
+
+        prior_copy = self._find_storage_copy(prior_spec.definition.dataItem, prior_spec.storage)
+
+        if prior_copy is None:
+            raise _ex.ETracInternal(f"Missing prior metadata for [{_util.object_key(file_id)}]")
+
+        storage_key = prior_copy.storageKey
+        storage_format = file_def.mimeType
+        storage_path = self._file_storage_path(file_id, file_def, prior_copy=None)
+
+        storage_copy = _meta.StorageCopy(
+            storageKey=storage_key,
+            storagePath=storage_path,
+            storageFormat=storage_format,
+            copyStatus=_meta.CopyStatus.COPY_AVAILABLE,
+            copyTimestamp=file_id.objectTimestamp)
+
+        storage_def = self._add_storage_copy(prior_spec.storage, data_item, storage_copy)
+
+        return DataSpec \
+            .create_file_spec(data_item, file_def, storage_def) \
+            .with_ids(file_id, storage_id)
+
+    @classmethod
+    def _add_new_snap(
+            cls, data_def: _meta.DataDefinition,data_item: str,
+            part_key: _meta.PartKey, snap_index: int):
+
+        delta = _meta.DataDefinition.Delta(
+            deltaIndex=0,
+            dataItem=data_item)
+
+        snap = _meta.DataDefinition.Snap(
+            snapIndex=snap_index,
+            deltas=[delta])
+
+        part = _meta.DataDefinition.Part(
+            partKey=part_key,
+            snap=snap)
+
+        data_def = copy.copy(data_def)
+        data_def.parts = copy.copy(data_def.parts)
+        data_def.parts[part_key.opaqueKey] = part
+
+        return data_def
+
+    @classmethod
+    def _add_storage_copy(cls, storage_def: _meta.StorageDefinition, data_item: str, storage_copy: _meta.StorageCopy):
+
+        new_incarnation = _meta.StorageIncarnation(
+            copies=[storage_copy],
+            incarnationIndex=0,
+            incarnationTimestamp=storage_copy.copyTimestamp,
+            incarnationStatus=_meta.IncarnationStatus.INCARNATION_AVAILABLE)
+
+        new_item = _meta.StorageItem(incarnations=[new_incarnation])
+
+        storage_def = copy.copy(storage_def)
+        storage_def.dataItems = copy.copy(storage_def.dataItems)
+        storage_def.dataItems[data_item] = new_item
+
+        return storage_def
+
+    @classmethod
+    def _find_storage_copy(cls, data_item: str, storage_def: _meta.StorageDefinition) -> tp.Optional[_meta.StorageCopy]:
+
+        if data_item is None:
+            return None
+
+        storage_item = storage_def.dataItems.get(data_item)
+
+        if storage_item is None:
+            return None
+
+        # Latest available incarnation
+        incarnation = next(filter(
+            lambda i: i.incarnationStatus == _meta.IncarnationStatus.INCARNATION_AVAILABLE,
+            reversed(storage_item.incarnations)), None)
+
+        if incarnation is None:
+            return None
+
+        # Use any available copy (currently there is no location preference)
+        return next(filter(
+            lambda c: c.copyStatus == _meta.CopyStatus.COPY_AVAILABLE,
+            incarnation.copies), None)
+
+
+class ObjectIdLayout(BaseLayout):
+
+    __DATA_STORAGE_TEMPLATE = "data/{}/{}/{}/snap-{:d}/delta-{:d}-x{:0>6x}"
+    __FILE_STORAGE_TEMPLATE = "file/{}/version-{:d}-x{:0>6x}/{}.{}"
+
+    def __init__(self):
+        self.__random = random.Random()
+        self.__random.seed()
+
+    def layout_key(self) -> _meta.StorageLayout:
+        return _meta.StorageLayout.OBJECT_ID_LAYOUT
+
+    def _data_storage_path(
+            self, data_id, context_key, trac_schema,
+            part_key, snap_index, delta_index,
+            storage_format, prior_copy):
+
+        schema_type = trac_schema.schemaType.name.lower()
+        version_suffix = self.__random.randint(0, 1 << 24)
+
+        return self.__DATA_STORAGE_TEMPLATE.format(
+            schema_type, data_id.objectId,
+            part_key.opaqueKey, snap_index, delta_index,
+            version_suffix)
+
+    def _file_storage_path(self, file_id, file_def, prior_copy):
+
+        version_suffix = self.__random.randint(0, 1 << 24)
+
+        return self.__FILE_STORAGE_TEMPLATE.format(
+            file_id.objectId, file_id.objectVersion, version_suffix,
+            file_def.name, file_def.extension.lower())
+
+
+class DevelopmentLayout(BaseLayout):
+
+    __DEFAULT_OUTPUT_DIR = "Dev Outputs"
+
+    __DATA_STORAGE_PATH = "{}/{}{}.{}"
+    __FILE_STORAGE_PATH = "{}/{}{}.{}"
+
+    def layout_key(self) -> _meta.StorageLayout:
+        return _meta.StorageLayout.DEVELOPER_LAYOUT
+
+    def _data_storage_path(
+            self, data_id, context_key, trac_schema,
+            part_key, snap_index, delta_index,
+            storage_format, prior_copy):
+
+        storage_dir = self._dev_storage_dir(prior_copy)
+        suffix = f"-{data_id.objectVersion}" if data_id.objectVersion > 1 else ""
+
+        if prior_copy is not None:
+            prior_path = pathlib.Path(prior_copy.storagePath)
+            file_name = prior_path.stem
+            if data_id.objectVersion > 2 and "-" in file_name:
+                file_name = file_name[:file_name.rfind("-")]
+        else:
+            file_name = context_key
+
+        return self.__DATA_STORAGE_PATH.format(storage_dir, file_name, suffix, storage_format.lower())
+
+    def _file_storage_path(self, file_id, file_def, prior_copy):
+
+        storage_dir = self._dev_storage_dir(prior_copy)
+        suffix = f"-{file_id.objectVersion}" if file_id.objectVersion > 1 else ""
+
+        return self.__FILE_STORAGE_PATH.format(storage_dir, file_def.name, suffix, file_def.extension.lower())
+
+    def _dev_storage_dir(self, prior_copy: _meta.StorageCopy):
+
+        if prior_copy is None:
+            return self.__DEFAULT_OUTPUT_DIR
+
+        prior_path = pathlib.Path(prior_copy.storagePath)
+
+        if len(prior_path.parts) > 1:
+            return prior_path.parent
+        else:
+            return self.__DEFAULT_OUTPUT_DIR
+
+
+def build_data_spec(
+        data_id: _meta.TagHeader, storage_id: _meta.TagHeader,
+        context_key: tp.Optional[str], trac_schema: _meta.SchemaDefinition,
+        storage_config: _cfg.StorageConfig,
+        prior_spec: tp.Optional[DataSpec] = None) \
+        -> DataSpec:
+
+    if prior_spec is None:
+        layout_key = storage_config.defaultLayout
+        layout = StorageLayout.select(layout_key)
+        return layout.new_data_spec(data_id, storage_id, context_key, trac_schema, storage_config)
+
+    else:
+        layout_key = prior_spec.storage.layout
+        layout = StorageLayout.select(layout_key)
+        return layout.new_data_version(data_id, storage_id, context_key, trac_schema, prior_spec)
+
+
+def build_file_spec(
+        file_id: _meta.TagHeader, storage_id: _meta.TagHeader,
+        context_key: tp.Optional[str],  file_type: _meta.FileType,
+        storage_config: _cfg.StorageConfig,
+        prior_spec: tp.Optional[DataSpec] = None) \
+        -> DataSpec:
+
+    if prior_spec is None:
+        layout_key = storage_config.defaultLayout
+        layout = StorageLayout.select(layout_key)
+        return layout.new_file_spec(file_id, storage_id, context_key, file_type, storage_config)
+
+    else:
+        layout_key = prior_spec.storage.layout
+        layout = StorageLayout.select(layout_key)
+        return layout.new_file_version(file_id, storage_id, context_key, file_type, prior_spec)
 
 
 @dc.dataclass(frozen=True)
@@ -184,6 +607,11 @@ class DataView:
         else:
             return DataView(_meta.ObjectType.DATA, trac_schema, parts = dict())
 
+    @staticmethod
+    def for_arrow_schema(arrow_schema: pa.Schema):
+        trac_schema = DataMapping.arrow_to_trac_schema(arrow_schema)
+        return DataView(_meta.ObjectType.DATA, trac_schema, arrow_schema, dict())
+
     @staticmethod
     def for_file_item(file_item: DataItem):
         return DataView(file_item.object_type, file_item=file_item)
@@ -259,8 +687,17 @@ class DataMapping:
         pa.date64(): _meta.BasicType.DATE
     }
 
-    @staticmethod
-    def arrow_to_python_type(arrow_type: pa.DataType) -> type:
+    # For now, categorical handling is disabled by default and enabled by this setting
+    # The default will change to "true" for the 0.9 release
+    CATEGORICAL_CONFIG_KEY = "trac.runtime.categorical"
+    __categorical_enabled = False
+
+    @classmethod
+    def enable_categorical(cls, enabled: bool):
+        cls.__categorical_enabled = enabled
+
+    @classmethod
+    def arrow_to_python_type(cls, arrow_type: pa.DataType) -> type:
 
         if pa.types.is_boolean(arrow_type):
             return bool
@@ -283,6 +720,11 @@ class DataMapping:
         if pa.types.is_timestamp(arrow_type):
             return dt.datetime
 
+        # The python type for a dictionary-encoded field is its value type
+        if pa.types.is_dictionary(arrow_type):
+            if isinstance(arrow_type, pa.DictionaryType):
+                return cls.arrow_to_python_type(arrow_type.value_type)
+
         raise _ex.ETracInternal(f"No Python type mapping available for Arrow type [{arrow_type}]")
 
     @classmethod
@@ -340,7 +782,13 @@ class DataMapping:
     def trac_to_arrow_field(cls, trac_field: _meta.FieldSchema):
 
         arrow_type = cls.trac_to_arrow_basic_type(trac_field.fieldType)
-        nullable = not trac_field.notNull if trac_field.notNull is not None else not trac_field.businessKey
+
+        # Categorical data uses an unordered dictionary with int32 index, ordered encoding not (currently) supported
+        # For legacy compatability, only use dictionary encoding if the categorical feature is enabled
+        if trac_field.categorical and cls.__categorical_enabled:
+            arrow_type = pa.dictionary(pa.int32(), arrow_type, False)
+
+        nullable = not (trac_field.notNull or trac_field.businessKey)
 
         return pa.field(trac_field.fieldName, arrow_type, nullable)
 
@@ -369,12 +817,15 @@ class DataMapping:
         field_type = cls.arrow_to_trac_type(field.type)
         label = field.metadata["label"] if field.metadata and "label" in field.metadata else field.name
 
+        # When converting Arrow -> TRAC, always set the categorical flag for dictionary encoded fields
+        # This affects dynamic imports and is informational only (physical layout is controlled by Arrow schemas)
+
         return _meta.FieldSchema(
             field.name, field_index, field_type,
             label=label,
             businessKey=False,
             notNull=not field.nullable,
-            categorical=False)
+            categorical=pa.types.is_dictionary(field.type))
 
     @classmethod
     def arrow_to_trac_type(cls, arrow_type: pa.DataType) -> _meta.BasicType:
@@ -390,6 +841,11 @@ class DataMapping:
         if pa.types.is_timestamp(arrow_type):
             return _meta.BasicType.DATETIME
 
+        # The basic type for a dictionary-encoded field is its value type
+        if pa.types.is_dictionary(arrow_type):
+            if isinstance(arrow_type, pa.DictionaryType):
+                return cls.arrow_to_trac_type(arrow_type.value_type)
+
         raise _ex.ETracInternal(f"No data type mapping available for Arrow type [{arrow_type}]")
 
     @classmethod
@@ -766,6 +1222,10 @@ class DataConformance:
         "Field [{field_name}] cannot be converted from {vector_type} to {field_type}, " + \
         "source and target have different time zones"
 
+    __E_WRONG_CATEGORICAL_TYPE = \
+        "Field [{field_name}] categorical types do not match" + \
+        "(expected {field_type}, got {vector_type})"
+
     @classmethod
     def column_filter(cls, columns: tp.List[str], schema: tp.Optional[pa.Schema]) -> tp.Optional[tp.List[str]]:
 
@@ -918,13 +1378,19 @@ class DataConformance:
     @classmethod
     def _coerce_vector(cls, vector: pa.Array, field: pa.Field, pandas_type=None) -> pa.Array:
 
+        # Handle null vector
         if pa.types.is_null(vector.type):
-
             if field.nullable:
                 return pa.nulls(size=len(vector), type=field.type)
             else:
                 raise _ex.EDataConformance(f"All null values in non-null field [{field.name}]")
 
+        # If the vector is dict-encoded but the expected result is not, decode the dictionary
+        if pa.types.is_dictionary(vector.type) and not pa.types.is_dictionary(field.type):
+            if isinstance(vector, pa.DictionaryArray):
+                dict_vector: pa.DictionaryArray = vector
+                vector = dict_vector.dictionary_decode()
+
         if pa.types.is_boolean(field.type):
             return cls._coerce_boolean(vector, field)
 
@@ -946,6 +1412,9 @@ class DataConformance:
         if pa.types.is_timestamp(field.type):
             return cls._coerce_timestamp(vector, field)
 
+        if pa.types.is_dictionary(field.type):
+            return cls._coerce_dictionary(vector, field)
+
         error_message = cls._format_error(cls.__E_WRONG_DATA_TYPE, vector, field)
         cls.__log.error(error_message)
         raise _ex.EDataConformance(error_message)
@@ -1188,6 +1657,82 @@ class DataConformance:
 
         return pc.cast(scaled_vector, field.type)
 
+    @classmethod
+    def _coerce_dictionary(cls, vector: pa.Array, field: pa.Field):
+
+        try:
+
+            if not isinstance(field.type, pa.DictionaryType):
+                raise _ex.EUnexpected()
+
+            field_type: pa.DictionaryType = field.type
+
+            # Supplied vector is a dictionary (but the dictionary type is not an exact match)
+            if pa.types.is_dictionary(vector.type):
+
+                if not isinstance(vector.type, pa.DictionaryType):
+                    raise _ex.EUnexpected()
+
+                vector_type: pa.DictionaryType = vector.type
+
+                # Do not allow coercion to a smaller index type or from unordered to ordered
+                if (vector_type.index_type.bit_width > field_type.index_type.bit_width) or \
+                   (field_type.ordered and not vector_type.ordered):
+
+                    error_message = cls._format_error(cls.__E_WRONG_DATA_TYPE, vector, field)
+                    cls.__log.error(error_message)
+                    raise _ex.EDataConformance(error_message)
+
+                # Value types are the same - basic cast should succeed
+                if vector_type.value_type == field_type.value_type:
+                    return pc.cast(vector, field.type)
+
+                # Value types differ - try to coerce the underlying dictionary
+                elif isinstance(vector, pa.DictionaryArray):
+                    try:
+                        values_field = pa.field(field.name, field_type.value_type, field.nullable)
+                        values_vector = cls._coerce_vector(vector.dictionary, values_field)
+                        dict_vector = pa.DictionaryArray.from_arrays(vector.indices, values_vector, ordered=field_type.ordered)  # noqa
+                        return pc.cast(dict_vector, field.type)
+                    # Handle errors converting the value type
+                    except _ex.EDataConformance as e:
+                        error_message = cls._format_error(cls.__E_WRONG_CATEGORICAL_TYPE, vector, field)
+                        cls.__log.error(error_message)
+                        raise _ex.EDataConformance(error_message) from e
+
+                # Special handling for chunked dictionaries
+                elif isinstance(vector, pa.ChunkedArray):
+                    chunks = [cls._coerce_dictionary(chunk, field) for chunk in vector.chunks]
+                    return pa.chunked_array(chunks)
+
+                # Vector type not recognized, coercion is not possible
+                else:
+                    error_message = cls._format_error(cls.__E_WRONG_DATA_TYPE, vector, field)
+                    cls.__log.error(error_message)
+                    raise _ex.EDataConformance(error_message)
+
+            # Supplied vector matches the dictionary value type - perform dictionary encoding
+            elif vector.type == field_type.value_type and not field_type.ordered:
+                return vector.dictionary_encode().cast(field.type)
+
+            # Fallback option - try to coerce the value type first, then perform dictionary encoding
+            else:
+                try:
+                    values_field = pa.field(field.name, field_type.value_type, field.nullable)
+                    values_vector = cls._coerce_vector(vector, values_field)
+                    return values_vector.dictionary_encode().cast(field.type)
+                # Handle errors converting the value type
+                except _ex.EDataConformance as e:
+                    error_message = cls._format_error(cls.__E_WRONG_CATEGORICAL_TYPE, vector, field)
+                    cls.__log.error(error_message)
+                    raise _ex.EDataConformance(error_message) from e
+
+        except pa.ArrowInvalid as e:
+
+            error_message = cls._format_error(cls.__E_DATA_LOSS_DID_OCCUR, vector, field, e)
+            cls.__log.error(error_message)
+            raise _ex.EDataConformance(error_message) from e
+
     @classmethod
     def _format_error(cls, error_template: str, vector: pa.Array, field: pa.Field, e: Exception = None):