real-ladybug 0.0.1.dev1__cp312-cp312-win_amd64.whl

real_ladybug-source/tools/python_api/build/real_ladybug/torch_geometric_graph_store.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import multiprocessing
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+from torch_geometric.data.graph_store import EdgeAttr, EdgeLayout, GraphStore
+
+from .connection import Connection
+
+if TYPE_CHECKING:
+    import sys
+
+    from torch_geometric.typing import EdgeTensorType
+
+    from .database import Database
+
+    if sys.version_info >= (3, 10):
+        from typing import TypeAlias
+    else:
+        from typing_extensions import TypeAlias
+
+StoreKeyType: TypeAlias = tuple[tuple[str], Any, bool]
+
+REL_BATCH_SIZE = 1000000
+
+
+@dataclass
+class Rel:  # noqa: D101
+    edge_type: tuple[str, ...]
+    layout: str
+    is_sorted: bool
+    size: tuple[int, ...]
+    materialized: bool = False
+    edge_index: EdgeTensorType | None = None
+
+
+class LbugGraphStore(GraphStore):  # type: ignore[misc]
+    """Graph store compatible with `torch_geometric`."""
+
+    def __init__(self, db: Database, num_threads: int | None = None):
+        super().__init__()
+        self.db = db
+        self.connection: Connection | None = None
+        self.store: dict[StoreKeyType, Rel] = {}
+        if num_threads is None:
+            num_threads = multiprocessing.cpu_count()
+        self.num_threads = num_threads
+        self.__populate_edge_attrs()
+
+    @staticmethod
+    def key(attr: EdgeAttr) -> tuple[tuple[str], Any, bool]:  # noqa: D102
+        return attr.edge_type, attr.layout.value, attr.is_sorted
+
+    def _put_edge_index(self, edge_index: EdgeTensorType, edge_attr: EdgeAttr) -> None:
+        key = self.key(edge_attr)
+        if key in self.store:
+            self.store[key].edge_index = edge_index
+            self.store[key].materialized = True
+            self.store[key].size = edge_attr.size
+        else:
+            self.store[key] = Rel(key[0], key[1], key[2], edge_attr.size, True, edge_index)
+
+    def _get_edge_index(self, edge_attr: EdgeAttr) -> EdgeTensorType | None:
+        if edge_attr.layout.value == EdgeLayout.COO.value:  # noqa: SIM102
+            # We always return a sorted COO edge index, if the request is
+            # for an unsorted COO edge index, we change the is_sorted flag
+            # to True and return the sorted COO edge index.
+            if edge_attr.is_sorted is False:
+                edge_attr.is_sorted = True
+
+        key = self.key(edge_attr)
+        if key in self.store:
+            rel = self.store[self.key(edge_attr)]
+            if not rel.materialized and rel.layout != EdgeLayout.COO.value:
+                msg = "Only COO layout is supported"
+                raise ValueError(msg)
+
+            if rel.layout == EdgeLayout.COO.value:
+                self.__get_edge_coo_from_database(self.key(edge_attr))
+            return rel.edge_index
+        else:
+            return None
+
+    def _remove_edge_index(self, edge_attr: EdgeAttr) -> None:
+        key = self.key(edge_attr)
+        if key in self.store:
+            del self.store[key]
+
+    def get_all_edge_attrs(self) -> list[EdgeAttr]:
+        """Return all EdgeAttr from the store values."""
+        return [EdgeAttr(rel.edge_type, rel.layout, rel.is_sorted, rel.size) for rel in self.store.values()]
+
+    def __get_edge_coo_from_database(self, key: StoreKeyType) -> None:
+        if not self.connection:
+            self.connection = Connection(self.db, self.num_threads)
+
+        rel = self.store[key]
+        if rel.layout != EdgeLayout.COO.value:
+            msg = "Only COO layout is supported"
+            raise ValueError(msg)
+        if rel.materialized:
+            return
+
+        edge_type = rel.edge_type
+        num_edges = self.connection._connection.get_num_rels(edge_type[1])
+        result = np.empty(2 * num_edges, dtype=np.int64)
+        self.connection._connection.get_all_edges_for_torch_geometric(
+            result, edge_type[0], edge_type[1], edge_type[2], REL_BATCH_SIZE
+        )
+        edge_list = torch.from_numpy(result)
+        edge_list = edge_list.reshape((2, edge_list.shape[0] // 2))
+        rel.edge_index = edge_list
+        rel.materialized = True
+
+    def __populate_edge_attrs(self) -> None:
+        if not self.connection:
+            self.connection = Connection(self.db, self.num_threads)
+        rel_tables = self.connection._get_rel_table_names()
+        for rel_table in rel_tables:
+            edge_type = (rel_table["src"], rel_table["name"], rel_table["dst"])
+            size = self.__get_size(edge_type)
+            rel = Rel(edge_type, EdgeLayout.COO.value, True, size, False, None)
+            self.store[self.key(EdgeAttr(edge_type, EdgeLayout.COO, True))] = rel
+
+    def __get_size(self, edge_type: tuple[str, ...]) -> tuple[int, int]:
+        num_nodes = self.connection._connection.get_num_nodes  # type: ignore[union-attr]
+        src_count, dst_count = num_nodes(edge_type[0]), num_nodes(edge_type[-1])
+        return (src_count, dst_count)

real_ladybug-source/tools/python_api/build/real_ladybug/torch_geometric_result_converter.py

@@ -0,0 +1,282 @@
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING, Any
+
+from .types import Type
+
+if TYPE_CHECKING:
+    import torch_geometric.data as geo
+
+    from .query_result import QueryResult
+
+from .constants import ID, LABEL, SRC, DST
+
+
+class TorchGeometricResultConverter:
+    """Convert graph results to `torch_geometric`."""
+
+    def __init__(self, query_result: QueryResult):
+        self.query_result = query_result
+        self.nodes_dict: dict[str, Any] = {}
+        self.edges_dict: dict[str, Any] = {}
+        self.edges_properties: dict[str | tuple[str, str], dict[str, Any]] = {}
+        self.rels: dict[tuple[Any, ...], dict[str, Any]] = {}
+        self.nodes_property_names_dict: dict[str, Any] = {}
+        self.table_to_label_dict: dict[int, str] = {}
+        self.internal_id_to_pos_dict: dict[tuple[int, int], int | None] = {}
+        self.pos_to_primary_key_dict: dict[str, Any] = {}
+        self.warning_messages: set[str] = set()
+        self.unconverted_properties: dict[str, Any] = {}
+        self.properties_to_extract = self.query_result._get_properties_to_extract()
+
+    def __get_node_property_names(self, table_name: str) -> dict[str, Any]:
+        if table_name in self.nodes_property_names_dict:
+            return self.nodes_property_names_dict[table_name]
+        results = self.query_result.connection._get_node_property_names(table_name)
+        self.nodes_property_names_dict[table_name] = results
+        return results
+
+    def __populate_nodes_dict_and_deduplicte_edges(self) -> None:
+        self.query_result.reset_iterator()
+        while self.query_result.has_next():
+            row = self.query_result.get_next()
+            for i in self.properties_to_extract:
+                column_type, _ = self.properties_to_extract[i]
+                if column_type == Type.NODE.value:
+                    node = row[i]
+                    label = node[LABEL]
+                    nid = node[ID]
+                    self.table_to_label_dict[nid["table"]] = label
+
+                    if (nid["table"], nid["offset"]) in self.internal_id_to_pos_dict:
+                        continue
+
+                    node_property_names = self.__get_node_property_names(label)
+
+                    pos, primary_key = self.__extract_properties_from_node(node, label, node_property_names)
+
+                    self.internal_id_to_pos_dict[nid["table"], nid["offset"]] = pos
+                    if label not in self.pos_to_primary_key_dict:
+                        self.pos_to_primary_key_dict[label] = {}
+                    self.pos_to_primary_key_dict[label][pos] = primary_key
+
+                elif column_type == Type.REL.value:
+                    src = row[i][SRC]
+                    dst = row[i][DST]
+                    self.rels[src["table"], src["offset"], dst["table"], dst["offset"]] = row[i]
+
+    def __extract_properties_from_node(
+        self,
+        node: dict[str, Any],
+        label: str,
+        node_property_names: dict[str, Any],
+    ) -> tuple[int | None, Any]:
+        pos = None
+        import torch
+
+        for prop_name in node_property_names:
+            # Read primary key
+            if node_property_names[prop_name]["is_primary_key"]:
+                primary_key = node[prop_name]
+
+            # If property is already marked as unconverted, then add it directly without further checks
+            if label in self.unconverted_properties and prop_name in self.unconverted_properties[label]:
+                pos = self.__add_unconverted_property(node, label, prop_name)
+                continue
+
+            # Mark properties that are not supported by torch_geometric as unconverted
+            if node_property_names[prop_name]["type"] not in [Type.INT64.value, Type.DOUBLE.value, Type.BOOL.value]:
+                self.warning_messages.add(
+                    "Property {}.{} of type {} is not supported by torch_geometric. The property is marked as unconverted.".format(
+                        label, prop_name, node_property_names[prop_name]["type"]
+                    )
+                )
+                self.__mark_property_unconverted(label, prop_name)
+                pos = self.__add_unconverted_property(node, label, prop_name)
+                continue
+            if node[prop_name] is None:
+                self.warning_messages.add(
+                    f"Property {label}.{prop_name} has a null value. torch_geometric does not support null values. The property is marked as unconverted."
+                )
+                self.__mark_property_unconverted(label, prop_name)
+                pos = self.__add_unconverted_property(node, label, prop_name)
+                continue
+
+            if node_property_names[prop_name]["dimension"] == 0:
+                curr_value = node[prop_name]
+            else:
+                try:
+                    if node_property_names[prop_name]["type"] == Type.INT64.value:
+                        curr_value = torch.LongTensor(node[prop_name])
+                    elif node_property_names[prop_name]["type"] == Type.DOUBLE.value:
+                        curr_value = torch.FloatTensor(node[prop_name])
+                    elif node_property_names[prop_name]["type"] == Type.BOOL.value:
+                        curr_value = torch.BoolTensor(node[prop_name])
+                except ValueError:
+                    self.warning_messages.add(
+                        f"Property {label}.{prop_name} cannot be converted to Tensor (likely due to nested list of variable length). The property is marked as unconverted."
+                    )
+                    self.__mark_property_unconverted(label, prop_name)
+                    pos = self.__add_unconverted_property(node, label, prop_name)
+                    continue
+
+                # Check if the shape of the property is consistent
+                if label in self.nodes_dict and prop_name in self.nodes_dict[label]:  # noqa: SIM102
+                    # If the shape is inconsistent, then mark the property as unconverted
+                    if curr_value.shape != self.nodes_dict[label][prop_name][0].shape:
+                        self.warning_messages.add(
+                            f"Property {label}.{prop_name} has an inconsistent shape. The property is marked as unconverted."
+                        )
+                        self.__mark_property_unconverted(label, prop_name)
+                        pos = self.__add_unconverted_property(node, label, prop_name)
+                        continue
+
+            # Create the dictionary for the label if it does not exist
+            if label not in self.nodes_dict:
+                self.nodes_dict[label] = {}
+            if prop_name not in self.nodes_dict[label]:
+                self.nodes_dict[label][prop_name] = []
+
+            # Add the property to the dictionary
+            self.nodes_dict[label][prop_name].append(curr_value)
+
+            # The pos will be overwritten for each property, but
+            # it should be the same for all properties
+            pos = len(self.nodes_dict[label][prop_name]) - 1
+        return pos, primary_key
+
+    def __add_unconverted_property(self, node: dict[str, Any], label: str, prop_name: str) -> int:
+        self.unconverted_properties[label][prop_name].append(node[prop_name])
+        return len(self.unconverted_properties[label][prop_name]) - 1
+
+    def __mark_property_unconverted(self, label: str, prop_name: str) -> None:
+        import torch
+
+        if label not in self.unconverted_properties:
+            self.unconverted_properties[label] = {}
+        if prop_name not in self.unconverted_properties[label]:
+            if label in self.nodes_dict and prop_name in self.nodes_dict[label]:
+                self.unconverted_properties[label][prop_name] = self.nodes_dict[label][prop_name]
+                del self.nodes_dict[label][prop_name]
+                if len(self.nodes_dict[label]) == 0:
+                    del self.nodes_dict[label]
+                for i in range(len(self.unconverted_properties[label][prop_name])):
+                    # If the property is a tensor, convert it back to list (consistent with the original type)
+                    if torch.is_tensor(self.unconverted_properties[label][prop_name][i]):  # type: ignore[no-untyped-call]
+                        self.unconverted_properties[label][prop_name][i] = self.unconverted_properties[label][
+                            prop_name
+                        ][i].tolist()
+            else:
+                self.unconverted_properties[label][prop_name] = []
+
+    def __populate_edges_dict(self) -> None:
+        # Post-process edges, map internal ids to positions
+        for r in self.rels:
+            src_pos = self.internal_id_to_pos_dict[r[0], r[1]]
+            dst_pos = self.internal_id_to_pos_dict[r[2], r[3]]
+            src_label = self.table_to_label_dict[r[0]]
+            dst_label = self.table_to_label_dict[r[2]]
+            if src_label not in self.edges_dict:
+                self.edges_dict[src_label] = {}
+            if dst_label not in self.edges_dict[src_label]:
+                self.edges_dict[src_label][dst_label] = []
+            self.edges_dict[src_label][dst_label].append((src_pos, dst_pos))
+            curr_edge_properties = self.rels[r]
+            if (src_label, dst_label) not in self.edges_properties:
+                self.edges_properties[src_label, dst_label] = {}
+            for prop_name in curr_edge_properties:
+                if prop_name in [SRC, DST, ID]:
+                    continue
+                if prop_name not in self.edges_properties[src_label, dst_label]:
+                    self.edges_properties[src_label, dst_label][prop_name] = []
+                self.edges_properties[src_label, dst_label][prop_name].append(curr_edge_properties[prop_name])
+
+    def __print_warnings(self) -> None:
+        for message in self.warning_messages:
+            warnings.warn(message, stacklevel=2)
+
+    def __convert_to_torch_geometric(
+        self,
+    ) -> tuple[
+        geo.Data | geo.HeteroData,
+        dict[str, Any],
+        dict[str, Any],
+        dict[str | tuple[str, str], dict[str, Any]],
+    ]:
+        import torch
+        import torch_geometric
+
+        if len(self.nodes_dict) == 0:
+            self.warning_messages.add("No nodes found or all node properties are not converted.")
+
+        # If there is only one node type, then convert to torch_geometric.data.Data
+        # Otherwise, convert to torch_geometric.data.HeteroData
+        if len(self.nodes_dict) == 1:
+            data = torch_geometric.data.Data()
+            is_hetero = False
+        else:
+            data = torch_geometric.data.HeteroData()
+            is_hetero = True
+
+        # Convert nodes to tensors
+        converted: torch.Tensor
+        for label in self.nodes_dict:
+            for prop_name in self.nodes_dict[label]:
+                prop_type = self.nodes_property_names_dict[label][prop_name]["type"]
+                prop_dimension = self.nodes_property_names_dict[label][prop_name]["dimension"]
+                if prop_dimension == 0:
+                    if prop_type == Type.INT64.value:
+                        converted = torch.LongTensor(self.nodes_dict[label][prop_name])
+                    elif prop_type == Type.BOOL.value:
+                        converted = torch.BoolTensor(self.nodes_dict[label][prop_name])
+                    elif prop_type == Type.DOUBLE.value:
+                        converted = torch.FloatTensor(self.nodes_dict[label][prop_name])
+                else:
+                    converted = torch.stack(self.nodes_dict[label][prop_name], dim=0)
+                if is_hetero:
+                    data[label][prop_name] = converted
+                else:
+                    data[prop_name] = converted
+
+        # Convert edges to tensors
+        for src_label in self.edges_dict:
+            for dst_label in self.edges_dict[src_label]:
+                edge_idx = torch.tensor(self.edges_dict[src_label][dst_label], dtype=torch.long).t().contiguous()
+                if is_hetero:
+                    data[src_label, dst_label].edge_index = edge_idx
+                else:
+                    data.edge_index = edge_idx
+
+        pos_to_primary_key_dict: dict[str, Any] = (
+            self.pos_to_primary_key_dict[label] if not is_hetero else self.pos_to_primary_key_dict
+        )
+
+        if is_hetero:
+            unconverted_properties = self.unconverted_properties
+            edge_properties = self.edges_properties
+        else:
+            if len(self.unconverted_properties) == 0:
+                unconverted_properties = {}
+            else:
+                unconverted_properties = self.unconverted_properties[next(iter(self.unconverted_properties))]
+            if len(self.edges_properties) == 0:
+                edge_properties = {}
+            else:
+                edge_properties = self.edges_properties[next(iter(self.edges_properties))]  # type: ignore[assignment]
+        return data, pos_to_primary_key_dict, unconverted_properties, edge_properties
+
+    def get_as_torch_geometric(
+        self,
+    ) -> tuple[
+        geo.Data | geo.HeteroData,
+        dict[str, Any],
+        dict[str, Any],
+        dict[str | tuple[str, str], dict[str, Any]],
+    ]:
+        """Convert graph data to `torch_geometric`."""
+        self.__populate_nodes_dict_and_deduplicte_edges()
+        self.__populate_edges_dict()
+        result = self.__convert_to_torch_geometric()
+        self.__print_warnings()
+        return result

real_ladybug-source/tools/python_api/build/real_ladybug/types.py

@@ -0,0 +1,39 @@
+from enum import Enum
+
+
+class Type(Enum):
+    """The type of a value in the database."""
+
+    ANY = "ANY"
+    NODE = "NODE"
+    REL = "REL"
+    RECURSIVE_REL = "RECURSIVE_REL"
+    SERIAL = "SERIAL"
+    BOOL = "BOOL"
+    INT64 = "INT64"
+    INT32 = "INT32"
+    INT16 = "INT16"
+    INT8 = "INT8"
+    UINT64 = "UINT64"
+    UINT32 = "UINT32"
+    UINT16 = "UINT16"
+    UINT8 = "UINT8"
+    INT128 = "INT128"
+    DOUBLE = "DOUBLE"
+    FLOAT = "FLOAT"
+    DATE = "DATE"
+    TIMESTAMP = "TIMESTAMP"
+    TIMSTAMP_TZ = "TIMESTAMP_TZ"
+    TIMESTAMP_NS = "TIMESTAMP_NS"
+    TIMESTAMP_MS = "TIMESTAMP_MS"
+    TIMESTAMP_SEC = "TIMESTAMP_SEC"
+    INTERVAL = "INTERVAL"
+    INTERNAL_ID = "INTERNAL_ID"
+    STRING = "STRING"
+    BLOB = "BLOB"
+    UUID = "UUID"
+    LIST = "LIST"
+    ARRAY = "ARRAY"
+    STRUCT = "STRUCT"
+    MAP = "MAP"
+    UNION = "UNION"

real_ladybug-source/tools/python_api/src_py/__init__.py

@@ -0,0 +1,83 @@
+"""
+# Lbug Python API bindings.
+
+This package provides a Python API for Lbug graph database management system.
+
+To install the package, run:
+```
+python3 -m pip install real_ladybug
+```
+
+Example usage:
+```python
+import real_ladybug as lb
+
+db = lb.Database("./test")
+conn = lb.Connection(db)
+
+# Define the schema
+conn.execute("CREATE NODE TABLE User(name STRING, age INT64, PRIMARY KEY (name))")
+conn.execute("CREATE NODE TABLE City(name STRING, population INT64, PRIMARY KEY (name))")
+conn.execute("CREATE REL TABLE Follows(FROM User TO User, since INT64)")
+conn.execute("CREATE REL TABLE LivesIn(FROM User TO City)")
+
+# Load some data
+conn.execute('COPY User FROM "user.csv"')
+conn.execute('COPY City FROM "city.csv"')
+conn.execute('COPY Follows FROM "follows.csv"')
+conn.execute('COPY LivesIn FROM "lives-in.csv"')
+
+# Query the data
+results = conn.execute("MATCH (u:User) RETURN u.name, u.age;")
+while results.has_next():
+    print(results.get_next())
+```
+
+The dataset used in this example can be found [here](https://github.com/LadybugDB/ladybug/tree/master/dataset/demo-db/csv).
+
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+# Set RTLD_GLOBAL and RTLD_LAZY flags on Linux to fix the issue with loading
+# extensions
+if sys.platform == "linux":
+    original_dlopen_flags = sys.getdlopenflags()
+    sys.setdlopenflags(os.RTLD_GLOBAL | os.RTLD_LAZY)
+
+from .async_connection import AsyncConnection
+from .connection import Connection
+from .database import Database
+from .prepared_statement import PreparedStatement
+from .query_result import QueryResult
+from .types import Type
+
+
+def __getattr__(name: str) -> str | int:
+    if name in ("version", "__version__"):
+        return Database.get_version()
+    elif name == "storage_version":
+        return Database.get_storage_version()
+    else:
+        msg = f"module {__name__!r} has no attribute {name!r}"
+        raise AttributeError(msg)
+
+
+# Restore the original dlopen flags
+if sys.platform == "linux":
+    sys.setdlopenflags(original_dlopen_flags)
+
+__all__ = [
+    "AsyncConnection",
+    "Connection",
+    "Database",
+    "PreparedStatement",
+    "QueryResult",
+    "Type",
+    "__version__",  # noqa: F822
+    "storage_version",  # noqa: F822
+    "version",  # noqa: F822
+]