k3-node 1.0.0__py3-none-any.whl

k3_node/data/hetero_data.py

@@ -0,0 +1,374 @@
+import copy
+from collections import defaultdict
+from typing import Any, Callable, Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple, Union
+
+import numpy as np
+from keras import ops
+
+from k3_node.data.data import BaseData, Data, size_repr
+from k3_node.data.storage import (
+    BaseStorage,
+    EdgeStorage,
+    NodeStorage,
+    get_shape,
+    is_tensor_like,
+)
+
+NodeType = str
+EdgeType = Tuple[str, str, str]
+
+
+class HeteroData(BaseData):
+    """A data object describing a heterogeneous graph."""
+
+    def __init__(self, _mapping: Optional[Dict[str, Any]] = None, **kwargs):
+        self.__dict__["_node_store_dict"] = {}
+        self.__dict__["_edge_store_dict"] = {}
+
+        for key, value in (_mapping or {}).items():
+            setattr(self, key, value)
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def _to_edge_type(self, key: Any) -> Optional[EdgeType]:
+        if isinstance(key, tuple):
+            if len(key) == 3:
+                return (str(key[0]), str(key[1]), str(key[2]))
+            if len(key) == 2:
+                # Find matching edge type
+                matches = [k for k in self.edge_types if k[0] == key[0] and k[-1] == key[1]]
+                if len(matches) == 1:
+                    return matches[0]
+                elif len(matches) > 1:
+                    raise KeyError(f"Ambiguous edge type for '{key}': {matches}")
+                return (str(key[0]), "to", str(key[1]))
+        elif isinstance(key, str):
+            matches = [k for k in self.edge_types if k[1] == key]
+            if len(matches) == 1:
+                return matches[0]
+            elif len(matches) > 1:
+                raise KeyError(f"Ambiguous edge type for rel '{key}': {matches}")
+        return None
+
+    def __getitem__(self, key: Any) -> Any:
+        edge_key = self._to_edge_type(key)
+        if edge_key is not None and (edge_key in self._edge_store_dict or isinstance(key, tuple)):
+            if edge_key not in self._edge_store_dict:
+                store = EdgeStorage(_parent=self)
+                store.__dict__["_key"] = edge_key
+                self._edge_store_dict[edge_key] = store
+            return self._edge_store_dict[edge_key]
+
+        if isinstance(key, str):
+            if key not in self._node_store_dict:
+                store = NodeStorage(_parent=self)
+                store.__dict__["_key"] = key
+                self._node_store_dict[key] = store
+            return self._node_store_dict[key]
+
+        raise KeyError(f"Invalid key '{key}' for {self.__class__.__name__}")
+
+    def __setitem__(self, key: Any, value: Any):
+        store = self[key]
+        if isinstance(value, BaseStorage):
+            for k, v in value.items():
+                store[k] = v
+        elif isinstance(value, dict):
+            for k, v in value.items():
+                store[k] = v
+        else:
+            raise ValueError(f"Value for key '{key}' must be a Storage or dict, got {type(value)}")
+
+    def __delitem__(self, key: Any):
+        edge_key = self._to_edge_type(key)
+        if edge_key is not None and edge_key in self._edge_store_dict:
+            del self._edge_store_dict[edge_key]
+        elif isinstance(key, str) and key in self._node_store_dict:
+            del self._node_store_dict[key]
+        else:
+            raise KeyError(key)
+
+    def __getattr__(self, key: str) -> Any:
+        if key in self.__dict__:
+            return self.__dict__[key]
+        if "_node_store_dict" in self.__dict__ and key in self._node_store_dict:
+            return self._node_store_dict[key]
+        if key.endswith("_dict") and "_node_store_dict" in self.__dict__:
+            return self.collect(key[:-5])
+        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{key}'")
+
+    def collect(self, key: str) -> Dict[Any, Any]:
+        r"""Returns the attribute ``key`` of every node and edge type that has it, e.g.
+        ``data.collect("x")`` (also available as ``data.x_dict``) or ``data.edge_index_dict``."""
+        out = {}
+        for stores in (self._node_store_dict, self._edge_store_dict):
+            for type_, store in stores.items():
+                value = store.get(key) if hasattr(store, "get") else getattr(store, key, None)
+                if value is not None:
+                    out[type_] = value
+        return out
+
+    def __setattr__(self, key: str, value: Any):
+        if key in ("_node_store_dict", "_edge_store_dict"):
+            self.__dict__[key] = value
+        elif isinstance(value, BaseStorage):
+            if isinstance(value, NodeStorage):
+                self._node_store_dict[key] = value
+            elif isinstance(value, EdgeStorage):
+                edge_key = self._to_edge_type(key)
+                self._edge_store_dict[edge_key or key] = value
+        else:
+            self.__dict__[key] = value
+
+    @property
+    def node_types(self) -> List[NodeType]:
+        return list(self._node_store_dict.keys())
+
+    @property
+    def edge_types(self) -> List[EdgeType]:
+        return list(self._edge_store_dict.keys())
+
+    def metadata(self) -> Tuple[List[NodeType], List[EdgeType]]:
+        return self.node_types, self.edge_types
+
+    def get_node_store(self, key: NodeType) -> NodeStorage:
+        r"""Gets the NodeStorage object of a particular node type."""
+        out = self._node_store_dict.get(key, None)
+        if out is None:
+            out = NodeStorage(_parent=self)
+            out.__dict__["_key"] = key
+            self._node_store_dict[key] = out
+        return out
+
+    def get_edge_store(self, src: str, rel: str, dst: str) -> EdgeStorage:
+        r"""Gets the EdgeStorage object of a particular edge type given by (src, rel, dst)."""
+        key = (src, rel, dst)
+        out = self._edge_store_dict.get(key, None)
+        if out is None:
+            out = EdgeStorage(_parent=self)
+            out.__dict__["_key"] = key
+            self._edge_store_dict[key] = out
+        return out
+
+    def stores_as(self, data: "HeteroData") -> "HeteroData":
+        for node_type in data.node_types:
+            self.get_node_store(node_type)
+        for edge_type in data.edge_types:
+            self.get_edge_store(*edge_type)
+        return self
+
+    @property
+    def stores(self) -> List[BaseStorage]:
+        return list(self._node_store_dict.values()) + list(self._edge_store_dict.values())
+
+    @property
+    def node_stores(self) -> List[NodeStorage]:
+        return list(self._node_store_dict.values())
+
+    @property
+    def edge_stores(self) -> List[EdgeStorage]:
+        return list(self._edge_store_dict.values())
+
+    @property
+    def num_nodes_dict(self) -> Dict[NodeType, int]:
+        return {k: v.num_nodes for k, v in self._node_store_dict.items()}
+
+    @property
+    def num_edges_dict(self) -> Dict[EdgeType, int]:
+        return {k: v.num_edges for k, v in self._edge_store_dict.items()}
+
+    def set_value_dict(self, key: str, value_dict: Optional[Dict[Any, Any]]):
+        for k, v in (value_dict or {}).items():
+            self[k][key] = v
+        return self
+
+    def __copy__(self):
+        out = self.__class__.__new__(self.__class__)
+        out.__dict__["_node_store_dict"] = {}
+        out.__dict__["_edge_store_dict"] = {}
+        for k, v in self._node_store_dict.items():
+            store = copy.copy(v)
+            setattr(store, "_parent", out)
+            out._node_store_dict[k] = store
+        for k, v in self._edge_store_dict.items():
+            store = copy.copy(v)
+            setattr(store, "_parent", out)
+            out._edge_store_dict[k] = store
+        return out
+
+    def __deepcopy__(self, memo=None):
+        out = self.__class__.__new__(self.__class__)
+        out.__dict__["_node_store_dict"] = {}
+        out.__dict__["_edge_store_dict"] = {}
+        for k, v in self._node_store_dict.items():
+            store = copy.deepcopy(v, memo)
+            setattr(store, "_parent", out)
+            out._node_store_dict[k] = store
+        for k, v in self._edge_store_dict.items():
+            store = copy.deepcopy(v, memo)
+            setattr(store, "_parent", out)
+            out._edge_store_dict[k] = store
+        return out
+
+    def __getstate__(self) -> Dict[str, Any]:
+        return self.__dict__.copy()
+
+    def __setstate__(self, mapping: Dict[str, Any]):
+        import weakref
+
+        for key, value in mapping.items():
+            self.__dict__[key] = value
+        for store in self.stores:
+            store.__dict__["_parent"] = weakref.ref(self)
+
+    def clone(self) -> "HeteroData":
+        return copy.deepcopy(self)
+
+    def collect(self, key: str, allow_missing: bool = True) -> Dict[Any, Any]:
+        mapping = {}
+        for k, store in list(self._node_store_dict.items()) + list(self._edge_store_dict.items()):
+            if key in store:
+                mapping[k] = store[key]
+            elif not allow_missing:
+                raise KeyError(f"Key '{key}' not found in store '{k}'")
+        return mapping
+
+    def __inc__(self, key: str, value: Any, store: Optional[BaseStorage] = None, *args, **kwargs) -> Any:
+        if "batch" in key:
+            return int(value.max()) + 1 if is_tensor_like(value) and value.ndim > 0 and value.shape[0] > 0 else 0
+        if "index" in key and store is not None and isinstance(store, EdgeStorage):
+            edge_type = store._key
+            src, _, dst = edge_type
+            src_num = self[src].num_nodes
+            dst_num = self[dst].num_nodes
+            return np.array([[src_num], [dst_num]])
+        return 0
+
+    def __cat_dim__(self, key: str, value: Any, store: Optional[BaseStorage] = None, *args, **kwargs) -> int:
+        if key in ("edge_index", "adj_t"):
+            return -1
+        if is_tensor_like(value) and len(get_shape(value)) == 2 and get_shape(value)[0] == 2 and "index" in key:
+            return -1
+        return 0
+
+    def to_dict(self) -> Dict[str, Any]:
+        out = {}
+        for k, store in self._node_store_dict.items():
+            out[k] = store.to_dict()
+        for k, store in self._edge_store_dict.items():
+            out[k] = store.to_dict()
+        return out
+
+    def to_namedtuple(self) -> NamedTuple:
+        # Build nested namedtuple
+        node_fields = sorted(list(self._node_store_dict.keys()))
+        node_dict = {k: self._node_store_dict[k].to_dict() for k in node_fields}
+        edge_fields = [f"{k[0]}__{k[1]}__{k[2]}" for k in sorted(list(self._edge_store_dict.keys()))]
+        edge_dict = {f"{k[0]}__{k[1]}__{k[2]}": self._edge_store_dict[k].to_dict() for k in sorted(list(self._edge_store_dict.keys()))}
+        all_fields = node_fields + edge_fields
+        HeteroTuple = collections.namedtuple("HeteroTuple", all_fields)
+        return HeteroTuple(**node_dict, **edge_dict)
+
+    def edge_type_subgraph(self, edge_types: List[EdgeType]) -> "HeteroData":
+        out = copy.deepcopy(self)
+        for et in list(out._edge_store_dict.keys()):
+            if et not in edge_types:
+                del out._edge_store_dict[et]
+        return out
+
+    def subgraph(self, subset_dict: Dict[NodeType, Any]) -> "HeteroData":
+        out = copy.deepcopy(self)
+        for node_type, subset in subset_dict.items():
+            store = out[node_type]
+            subset_np = ops.convert_to_numpy(subset)
+            indices = np.where(subset_np)[0] if subset_np.dtype == bool else subset_np
+            for key in store.node_attrs():
+                val = store[key]
+                if is_tensor_like(val):
+                    store[key] = ops.take(val, indices, axis=self.__cat_dim__(key, val, store))
+        return out
+
+    def to_homogeneous(
+        self,
+        node_attrs: Optional[List[str]] = None,
+        edge_attrs: Optional[List[str]] = None,
+        add_node_type: bool = True,
+        add_edge_type: bool = True,
+        dummy_values: bool = True,
+    ) -> Data:
+        data = Data()
+
+        # Compute node offsets and slices
+        node_slices = {}
+        curr_offset = 0
+        node_type_list = []
+        for i, node_type in enumerate(self.node_types):
+            num_nodes = self[node_type].num_nodes
+            node_slices[node_type] = curr_offset
+            if add_node_type:
+                node_type_list.append(np.full((num_nodes,), i, dtype=np.int64))
+            curr_offset += num_nodes
+
+        data.num_nodes = curr_offset
+        if add_node_type and len(node_type_list) > 0:
+            node_type_arr = np.concatenate(node_type_list, axis=0)
+            data.node_type = ops.convert_to_tensor(node_type_arr, dtype="int64")
+
+        # Concat node features
+        if node_attrs is None:
+            # find common node attrs across node types
+            all_node_attrs = set()
+            for store in self.node_stores:
+                all_node_attrs.update(store.node_attrs())
+            node_attrs = list(all_node_attrs)
+
+        for attr in node_attrs:
+            attr_vals = []
+            for node_type in self.node_types:
+                val = self[node_type].get(attr)
+                if val is not None:
+                    attr_vals.append(ops.convert_to_numpy(val))
+                elif dummy_values:
+                    num_nodes = self[node_type].num_nodes
+                    attr_vals.append(np.zeros((num_nodes, 0), dtype=np.float32))
+            if len(attr_vals) > 0:
+                concat_val = np.concatenate(attr_vals, axis=0)
+                data[attr] = ops.convert_to_tensor(concat_val)
+
+        # Offsetting edge indices
+        edge_indices = []
+        edge_type_list = []
+        for i, edge_type in enumerate(self.edge_types):
+            store = self[edge_type]
+            edge_index = store.get("edge_index")
+            if edge_index is not None:
+                ei_np = ops.convert_to_numpy(edge_index).copy()
+                src_offset = node_slices[edge_type[0]]
+                dst_offset = node_slices[edge_type[-1]]
+                ei_np[0] += src_offset
+                ei_np[1] += dst_offset
+                edge_indices.append(ei_np)
+                if add_edge_type:
+                    edge_type_list.append(np.full((ei_np.shape[1],), i, dtype=np.int64))
+
+        if len(edge_indices) > 0:
+            concat_ei = np.concatenate(edge_indices, axis=1)
+            data.edge_index = ops.convert_to_tensor(concat_ei, dtype="int64")
+        if add_edge_type and len(edge_type_list) > 0:
+            concat_et = np.concatenate(edge_type_list, axis=0)
+            data.edge_type = ops.convert_to_tensor(concat_et, dtype="int64")
+
+        return data
+
+    def __repr__(self) -> str:
+        cls = self.__class__.__name__
+        info_lines = []
+        for k, store in self._node_store_dict.items():
+            attrs = [size_repr(attr, store[attr]) for attr in store.keys()]
+            info_lines.append(f"  {k}={{{', '.join(attrs)}}}")
+        for k, store in self._edge_store_dict.items():
+            attrs = [size_repr(attr, store[attr]) for attr in store.keys()]
+            edge_name = f"('{k[0]}', '{k[1]}', '{k[2]}')"
+            info_lines.append(f"  {edge_name}={{{', '.join(attrs)}}}")
+        info = ",\n".join(info_lines)
+        return f"{cls}(\n{info}\n)" if info else f"{cls}()"

k3_node/data/hypergraph_data.py

@@ -0,0 +1,59 @@
+from typing import Any, List, Optional
+import numpy as np
+from keras import ops
+
+from k3_node.data.data import Data
+from k3_node.data.storage import is_tensor_like
+
+
+class HypergraphData(Data):
+    """A data object describing a hypergraph."""
+
+    def __init__(
+        self,
+        x=None,
+        edge_index=None,
+        edge_attr=None,
+        y=None,
+        pos=None,
+        **kwargs,
+    ):
+        super().__init__(
+            x=x,
+            edge_index=edge_index,
+            edge_attr=edge_attr,
+            y=y,
+            pos=pos,
+            **kwargs,
+        )
+
+    @property
+    def num_edges(self) -> int:
+        if self.edge_index is None:
+            return 0
+        ei_np = ops.convert_to_numpy(self.edge_index)
+        if ei_np.size == 0 or ei_np.shape[1] == 0:
+            return 0
+        return int(np.max(ei_np[1])) + 1
+
+    @property
+    def num_nodes(self) -> Optional[int]:
+        num = super().num_nodes
+        if self.edge_index is not None and num == self.num_edges:
+            ei_np = ops.convert_to_numpy(self.edge_index)
+            if ei_np.size > 0 and ei_np.shape[1] > 0:
+                return int(np.max(ei_np[0])) + 1
+        return num
+
+    @num_nodes.setter
+    def num_nodes(self, num_nodes: Optional[int]) -> None:
+        self._store.num_nodes = num_nodes
+
+    def __inc__(self, key: str, value: Any, *args, **kwargs) -> Any:
+        if key == "edge_index":
+            return np.array([[self.num_nodes], [self.num_edges]])
+        return super().__inc__(key, value, *args, **kwargs)
+
+
+HyperGraphData = HypergraphData
+

k3_node/data/in_memory_dataset.py

@@ -0,0 +1,177 @@
+import copy
+import os
+import os.path as osp
+import pickle
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+from k3_node.data.collate import collate
+from k3_node.data.data import BaseData, Data
+from k3_node.data.dataset import Dataset
+from k3_node.data.separate import separate
+
+
+class InMemoryDataset(Dataset):
+    """Dataset base class for in-memory graph collections."""
+
+    def __init__(
+        self,
+        root: Optional[str] = None,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        log: bool = True,
+        force_reload: bool = False,
+    ):
+        # Must be set before `super().__init__()`, which triggers
+        # `self.process()` for datasets processed for the first time --
+        # `process()` implementations (e.g. `TUDataset`) commonly call
+        # `len(self)` / `self.get(idx)`, both of which read these attributes.
+        self._data: Optional[BaseData] = None
+        self.slices: Optional[Dict[str, Any]] = None
+        self.sizes: Dict[str, Any] = {}
+        self._data_list: Optional[List[BaseData]] = None
+        super().__init__(root, transform, pre_transform, pre_filter, log, force_reload)
+
+    @property
+    def data(self) -> Optional[BaseData]:
+        return self._data
+
+    @data.setter
+    def data(self, value: Optional[BaseData]):
+        self._data = value
+
+    def len(self) -> int:
+        if self._data_list is not None:
+            return len(self._data_list)
+        if self.slices is None:
+            return 1 if self._data is not None else 0
+        for key, value in self.slices.items():
+            if isinstance(value, dict):
+                for _, sub_val in value.items():
+                    return len(sub_val) - 1
+            return len(value) - 1
+        return 0
+
+    def get(self, idx: int) -> BaseData:
+        if self._data_list is not None:
+            return self._data_list[idx]
+        if self._data is None:
+            raise RuntimeError("Dataset does not contain data. Call 'load' first.")
+        if self.slices is None:
+            if idx == 0:
+                return copy.copy(self._data)
+            raise IndexError(f"Index {idx} out of bounds for single graph dataset.")
+        return separate(self._data.__class__, self._data, idx, self.slices)
+
+    @classmethod
+    def collate(cls, data_list: List[BaseData]) -> Tuple[BaseData, Optional[Dict[str, Any]]]:
+        if len(data_list) == 1:
+            return data_list[0], None
+        base_cls = data_list[0].__class__
+        data, slices, _ = collate(
+            base_cls,
+            data_list,
+            increment=False,
+            add_batch=False,
+        )
+        return data, slices
+
+    def save(self, data_list: List[BaseData], path: str):
+        os.makedirs(osp.dirname(path), exist_ok=True)
+        data, slices = self.collate(data_list)
+        if path.endswith((".pt", ".pth")):
+            try:
+                import torch
+                data_dict = data.to_dict() if hasattr(data, "to_dict") else dict(data)
+                torch.save((data_dict, slices), path)
+                return
+            except Exception:
+                pass
+        with open(path, "wb") as f:
+            pickle.dump((data, slices), f)
+
+    def load(self, path: str):
+        obj = None
+        if path.endswith((".pt", ".pth")):
+            try:
+                import torch
+                obj = torch.load(path, map_location="cpu", weights_only=False)
+            except Exception:
+                pass
+
+        if obj is None:
+            try:
+                with open(path, "rb") as f:
+                    obj = pickle.load(f)
+            except Exception:
+                try:
+                    import torch
+                    obj = torch.load(path, map_location="cpu", weights_only=False)
+                except Exception:
+                    pass
+
+        if obj is None:
+            if hasattr(self, "process") and callable(self.process):
+                self.process()
+                try:
+                    with open(path, "rb") as f:
+                        obj = pickle.load(f)
+                except Exception:
+                    try:
+                        import torch
+                        obj = torch.load(path, map_location="cpu", weights_only=False)
+                    except Exception:
+                        pass
+            if obj is None:
+                raise RuntimeError(f"Cannot load dataset from {path}")
+
+        if isinstance(obj, tuple):
+            if len(obj) == 2:
+                data, self.slices = obj
+            elif len(obj) == 3:
+                data, self.slices, extra = obj
+                if isinstance(extra, dict):
+                    self.sizes = extra
+            elif len(obj) >= 4:
+                data, self.slices = obj[0], obj[1]
+                if isinstance(obj[2], dict):
+                    self.sizes = obj[2]
+            else:
+                data = obj[0]
+            if isinstance(data, dict):
+                data = _from_dict(data)
+            self._data = data
+        elif isinstance(obj, list):
+            self._data_list = obj
+        elif isinstance(obj, dict):
+            self._data = _from_dict(obj)
+        else:
+            self._data = obj
+
+        if self._data is not None and hasattr(self._data, "to_backend"):
+            try:
+                self._data.to_backend()
+            except Exception:
+                pass
+
+        if isinstance(self.slices, dict):
+            from k3_node.data.storage import is_tensor_like, to_numpy
+            self.slices = {k: to_numpy(v) if is_tensor_like(v) else v for k, v in self.slices.items()}
+
+
+
+def _from_dict(mapping):
+    """Rebuilds a saved graph: a ``HeteroData`` if the mapping holds per-type attribute
+    dictionaries (node types and ``(src, rel, dst)`` edge types), else a ``Data``."""
+    if any(isinstance(v, dict) or not isinstance(k, str) for k, v in mapping.items()):
+        from k3_node.data.hetero_data import HeteroData
+
+        data = HeteroData()
+        for key, value in mapping.items():
+            if isinstance(value, dict):
+                for attr, item in value.items():
+                    setattr(data[key], attr, item)
+            else:
+                setattr(data, key, value)
+        return data
+    return Data(**mapping)

k3_node/data/makedirs.py

@@ -0,0 +1,7 @@
+import os
+
+
+def makedirs(path: str):
+    """Recursively creates a directory."""
+    os.makedirs(path, exist_ok=True)
+

k3_node/data/on_disk_dataset.py

@@ -0,0 +1,77 @@
+import os
+import os.path as osp
+from typing import Any, Callable, List, Optional, Union
+
+from k3_node.data.data import BaseData
+from k3_node.data.database import Database, RocksDatabase, SQLiteDatabase, Schema
+from k3_node.data.dataset import Dataset
+
+
+class OnDiskDataset(Dataset):
+    """Dataset base class for out-of-core graph datasets using a Database backend."""
+
+    BACKENDS = {
+        "sqlite": SQLiteDatabase,
+        "rocksdb": RocksDatabase,
+    }
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        backend: str = "sqlite",
+        schema: Schema = object,
+        log: bool = True,
+    ):
+        if backend not in self.BACKENDS:
+            raise ValueError(f"Database backend must be one of {set(self.BACKENDS.keys())}, got '{backend}'")
+
+        self.backend = backend
+        self.schema = schema
+        self._db: Optional[Database] = None
+
+        super().__init__(root, transform, pre_filter=pre_filter, log=log)
+
+    @property
+    def processed_file_names(self) -> str:
+        return f"{self.backend}.db"
+
+    @property
+    def db(self) -> Database:
+        if self._db is not None:
+            return self._db
+
+        cls = self.BACKENDS[self.backend]
+        os.makedirs(self.processed_dir, exist_ok=True)
+        path = osp.join(self.processed_dir, self.processed_file_names)
+        self._db = cls(path=path, schema=self.schema)
+        return self._db
+
+    def close(self):
+        if self._db is not None:
+            self._db.close()
+            self._db = None
+
+    def serialize(self, data: BaseData) -> Any:
+        return data
+
+    def deserialize(self, data: Any) -> BaseData:
+        return data
+
+    def len(self) -> int:
+        return len(self.db)
+
+    def get(self, idx: int) -> BaseData:
+        return self.deserialize(self.db[idx])
+
+    def append(self, data: BaseData):
+        idx = len(self)
+        self.db[idx] = self.serialize(data)
+
+    def extend(self, data_list: List[BaseData]):
+        start = len(self)
+        indices = list(range(start, start + len(data_list)))
+        serialized = [self.serialize(d) for d in data_list]
+        self.db.multi_insert(indices, serialized)
+