Skip to content

Graph Fields

Mode

Bases: str, Enum

Source code in src/mifrost/graph_fields.py 
 8
 9
10
11
12
class Mode(str, Enum):
    STACK = "stack"
    CAT = "cat"
    RAGGED_CAT = "ragged_cat"
    CONST = "const"

DType

Bases: str, Enum

Source code in src/mifrost/graph_fields.py 
15
16
17
18
19
class DType(str, Enum):
    F32 = "f32"
    I64 = "i64"
    STR = "str"
    PYOBJ = "pyobj"

Inc dataclass

Source code in src/mifrost/graph_fields.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
@dataclass(frozen=True)
class Inc:
    kind: str
    node_type: str = ""
    field_key: str = ""

    @staticmethod
    def none() -> "Inc":
        return Inc(kind="none")

    @staticmethod
    def node_offset(node_type: str) -> "Inc":
        if not node_type:
            raise ValueError("Inc.node_offset requires a non-empty node_type")
        return Inc(kind="node_offset", node_type=node_type)

    @staticmethod
    def field_offset(field_key: str) -> "Inc":
        if not field_key:
            raise ValueError("Inc.field_offset requires a non-empty field_key")
        return Inc(kind="field_offset", field_key=field_key)

    def to_core_dict(self) -> dict[str, str]:
        payload = {"kind": self.kind}
        if self.kind == "node_offset":
            if not self.node_type:
                raise ValueError("Inc.node_offset requires a non-empty node_type")
            payload["node_type"] = self.node_type
        elif self.kind == "field_offset":
            if not self.field_key:
                raise ValueError("Inc.field_offset requires a non-empty field_key")
            payload["field_key"] = self.field_key
        return payload

GraphFieldSpec dataclass

Source code in src/mifrost/graph_fields.py 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
@dataclass(frozen=True)
class GraphFieldSpec:
    mode: Mode
    dtype: DType
    dim: int = 1
    cat_dim: int = 0
    inc: Inc = field(default_factory=Inc.none)

    def __post_init__(self) -> None:
        if self.dim <= 0:
            raise ValueError("GraphFieldSpec.dim must be > 0")
        if self.dtype in (DType.PYOBJ, DType.STR):
            raise ValueError(
                "GraphFieldSpec dtype is only supported for Python-side "
                "batch_encodings collate_spec, not native HGraph dynamic fields"
            )
        normalized_cat_dim = 1 if self.cat_dim == -1 else int(self.cat_dim)
        if self.mode in (Mode.CAT, Mode.RAGGED_CAT):
            if normalized_cat_dim not in (0, 1):
                raise ValueError(
                    "GraphFieldSpec.cat_dim must be 0, 1, or -1 for CAT/RAGGED_CAT"
                )
        elif normalized_cat_dim != 0:
            raise ValueError("GraphFieldSpec.cat_dim must be 0 for STACK/CONST")
        object.__setattr__(self, "cat_dim", normalized_cat_dim)
        if self.inc.kind == "node_offset" and self.dtype is not DType.I64:
            raise ValueError("NODE_OFFSET increment requires dtype=DType.I64")
        if self.inc.kind == "field_offset" and self.dtype is not DType.I64:
            raise ValueError("FIELD_OFFSET increment requires dtype=DType.I64")

    def to_core_dict(self) -> dict[str, Any]:
        return {
            "mode": self.mode.value,
            "dtype": self.dtype.value,
            "dim": int(self.dim),
            "cat_dim": int(self.cat_dim),
            "inc": self.inc.to_core_dict(),
        }

    @classmethod
    def from_spec(cls, spec: "GraphFieldSpec | Mapping[str, Any]") -> "GraphFieldSpec":
        if isinstance(spec, cls):
            return spec
        if not isinstance(spec, Mapping):
            raise TypeError(
                f"Graph field spec must be GraphFieldSpec or mapping, got {type(spec)!r}"
            )

        mode_raw = spec.get("mode")
        dtype_raw = spec.get("dtype")
        if mode_raw is None or dtype_raw is None:
            raise ValueError("Graph field spec mapping requires 'mode' and 'dtype'")

        mode = _normalize_mode(mode_raw)
        dtype = _normalize_dtype(dtype_raw)
        dim = int(spec.get("dim", 1))
        cat_dim = int(spec.get("cat_dim", 0))
        inc = _normalize_inc(spec.get("inc", {"kind": "none"}))
        return cls(mode=mode, dtype=dtype, dim=dim, cat_dim=cat_dim, inc=inc)