open-edge-platform
diff --git a/‎src/datumaro/experimental/categories.py‎
Lines changed: 189 additions & 1 deletion b/‎src/datumaro/experimental/categories.py‎
Lines changed: 189 additions & 1 deletion
diff --git a/‎src/datumaro/experimental/fields.py‎
Lines changed: 1 addition & 9 deletions b/‎src/datumaro/experimental/fields.py‎
Lines changed: 1 addition & 9 deletions
@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+from collections import defaultdict
 from dataclasses import dataclass, field
 from enum import IntEnum
 from functools import cache
@@ -33,6 +34,7 @@ class GroupType(IntEnum):
 
     EXCLUSIVE = 0  # Only one label from the group can be assigned
     INCLUSIVE = 1  # Multiple labels from the group can be assigned
+    RESTRICTED = 2  # For empty labels
 
     def to_str(self) -> str:
         return self.name.lower()
@@ -58,7 +60,10 @@ class Categories:
 
 @dataclass(frozen=True)
 class LabelCategories(Categories):
-    """Represents a group of labels with a specific group type and semantics."""
+    """
+    Represents a group of labels with a specific group type and semantics.
+    Use this for simple, non-hierarchical tasks.
+    """
 
     labels: Tuple[str, ...] = field(default_factory=tuple)
     group_type: GroupType = GroupType.EXCLUSIVE
@@ -131,6 +136,189 @@ def __hash__(self):
         return hash((self.labels, self.group_type, frozenset(self.label_semantics.items())))
 
 
+@dataclass(frozen=True)
+class HierarchicalLabelCategory:
+    """Represents a single label category with hierarchical support."""
+
+    name: str
+    parent: str = field(default="")
+    label_semantics: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+        """Validate that name is not empty."""
+        if not self.name or not isinstance(self.name, str):
+            raise ValueError("Label name cannot be empty and must be a string")
+
+    def __hash__(self):
+        return hash((self.name, self.parent, frozenset(self.label_semantics.items())))
+
+
+@dataclass(frozen=True)
+class LabelGroup:
+    """Represents a group of labels with a specific group type."""
+
+    name: str
+    labels: Tuple[str, ...] = field(default_factory=tuple)
+    group_type: GroupType = GroupType.EXCLUSIVE
+
+    def __post_init__(self):
+        """Validate that name is not empty and labels is a tuple."""
+        if not self.name or not isinstance(self.name, str):
+            raise ValueError("Label group name cannot be empty and must be a string")
+
+
+@dataclass(frozen=True)
+class HierarchicalLabelCategories(Categories):
+    """
+    Represents hierarchical label categories with groups and parent-child relationships.
+    Use this for complex hierarchical classification tasks.
+    """
+
+    items: Tuple[HierarchicalLabelCategory, ...] = field(default_factory=tuple)
+    label_groups: Tuple[LabelGroup, ...] = field(default_factory=tuple)
+    label_semantics: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+        if not isinstance(self.items, tuple):
+            raise TypeError("items must be a tuple of HierarchicalLabelCategory")
+        if not isinstance(self.label_groups, tuple):
+            raise TypeError("label_groups must be a tuple of LabelGroup")
+
+        # Validate no duplicate names
+        seen_names = set()
+        for item in self.items:
+            if item.name in seen_names:
+                raise ValueError(f"Duplicate label name: {item.name}")
+            seen_names.add(item.name)
+
+            # Also check for name label_semantics
+            # TODO: Remove this check after migrating completely to new system
+            if "name" in item.label_semantics:
+                name = item.label_semantics["name"]
+                if name not in seen_names:
+                    seen_names.add(name)
+
+        # Validate that all parents exist
+        for item in self.items:
+            if item.parent and item.parent not in seen_names:
+                raise ValueError(f"Parent '{item.parent}' not found for label '{item.name}'")
+
+        # Validate that all labels in groups exist
+        for group in self.label_groups:
+            for label_name in group.labels:
+                if label_name not in seen_names:
+                    raise ValueError(
+                        f"Label '{label_name}' in group '{group.name}' not found in items"
+                    )
+
+    @property
+    @cache
+    def _index_map(self) -> Dict[str, int]:
+        """Cached mapping from label names to indices."""
+        return {item.name: idx for idx, item in enumerate(self.items)}
+
+    @property
+    @cache
+    def _children_map(self) -> Dict[str, Tuple[str, ...]]:
+        """Cached mapping from parent names to child names."""
+        children_map: defaultdict[str, List[str]] = defaultdict(list)
+        for item in self.items:
+            if item.parent:
+                children_map[item.parent].append(item.name)
+        return {parent: tuple(children) for parent, children in children_map.items()}
+
+    @property
+    def labels(self) -> Tuple[str, ...]:
+        """Get all label names for compatibility."""
+        return tuple(item.name for item in self.items)
+
+    def find(self, name: str) -> Tuple[Optional[int], Optional[HierarchicalLabelCategory]]:
+        """
+        Find a label by name.
+
+        Args:
+            name: The label name to find
+
+        Returns:
+            A tuple of (index, category) or (None, None) if not found
+        """
+        index = self._index_map.get(name)
+        if index is not None:
+            return index, self.items[index]
+        return None, None
+
+    def get_children(self, parent_name: str) -> Tuple[str, ...]:
+        """
+        Get all children of a parent label.
+
+        Args:
+            parent_name: The name of the parent label
+
+        Returns:
+            Tuple of child label names
+        """
+        return self._children_map.get(parent_name, ())
+
+    def get_parent(self, label_name: str) -> Optional[str]:
+        """
+        Get the parent of a label.
+
+        Args:
+            label_name: The name of the label
+
+        Returns:
+            Parent name or None if no parent
+        """
+        index = self._index_map.get(label_name)
+        if index is not None:
+            return self.items[index].parent
+        return None
+
+    def get_hierarchy_level(self, label_name: str) -> int:
+        """
+        Get the hierarchy level of a label (0 for root, 1 for first level children, etc.)
+
+        Args:
+            label_name: The name of the label
+
+        Returns:
+            Hierarchy level
+        """
+        level = 0
+        current = label_name
+        while True:
+            parent = self.get_parent(current)
+            if not parent:
+                break
+            level += 1
+            current = parent
+        return level
+
+    def __getitem__(self, idx: int) -> HierarchicalLabelCategory:
+        """Get category by index."""
+        return self.items[idx]
+
+    def __contains__(self, value: Union[int, str]) -> bool:
+        """Check if a label exists by name or index."""
+        if isinstance(value, str):
+            return value in self._index_map
+        else:
+            return 0 <= value < len(self.items)
+
+    def __len__(self) -> int:
+        """Get the number of labels."""
+        return len(self.items)
+
+    def __iter__(self) -> Iterator[HierarchicalLabelCategory]:
+        """Iterate over label categories."""
+        return iter(self.items)
+
+    def __hash__(self):
+        # Hash label_groups via value-based representation to avoid relying on their own hash implementation.
+        lg_repr = tuple((lg.name, tuple(lg.labels), lg.group_type) for lg in self.label_groups)
+        return hash((self.items, lg_repr, frozenset(self.label_semantics.items())))
+
+
 class RgbColor(NamedTuple):
     """RGB color representation with named fields."""
 
 
@@ -532,15 +532,7 @@ def to_polars_schema(self, name: str) -> dict[str, pl.DataType]:
 
     def to_polars(self, name: str, value: Any) -> dict[str, pl.Series]:
         """Convert label(s) to Polars format for single or multi-label cases."""
-        pl_type = self._pl_type
-
-        if value is None:
-            return {name: pl.Series(name, [None], dtype=pl_type)}
-
-        if self.multi_label:
-            return {name: pl.Series(name, [to_numpy(value)], dtype=pl.List(self.dtype))}
-
-        return {name: pl.Series(name, [value], dtype=pl_type)}
+        return {name: pl.Series(name, [value], dtype=self._pl_type)}
 
     def from_polars(self, name: str, row_index: int, df: pl.DataFrame, target_type: type[T]) -> T:
         """Reconstruct label(s) from Polars data."""