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

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

@@ -0,0 +1,511 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from .torch_geometric_result_converter import TorchGeometricResultConverter
+from .types import Type
+
+from .constants import ID, LABEL, SRC, DST, NODES, RELS
+
+if TYPE_CHECKING:
+    import sys
+    from collections.abc import Iterator
+    from types import TracebackType
+    from typing import Any
+
+    import networkx as nx
+    import pandas as pd
+    import polars as pl
+    import pyarrow as pa
+    import torch_geometric.data as geo
+
+    from . import _lbug
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
+
+class QueryResult:
+    """QueryResult stores the result of a query execution."""
+
+    def __init__(self, connection: _lbug.Connection, query_result: _lbug.QueryResult):  # type: ignore[name-defined]
+        """
+        Parameters
+        ----------
+        connection : _lbug.Connection
+            The underlying C++ connection object from pybind11.
+
+        query_result : _lbug.QueryResult
+            The underlying C++ query result object from pybind11.
+
+        """
+        self.connection = connection
+        self._query_result = query_result
+        self.is_closed = False
+        self.as_dict = False
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        exc_traceback: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    def __del__(self) -> None:
+        self.close()
+
+    def __iter__(self) -> Iterator[list[Any] | dict[str, Any]]:
+        return self
+
+    def __next__(self) -> list[Any] | dict[str, Any]:
+        if self.has_next():
+            return self.get_next()
+
+        raise StopIteration
+
+    def has_next(self) -> bool:
+        """
+        Check if there are more rows in the query result.
+
+        Returns
+        -------
+        bool
+            True if there are more rows in the query result, False otherwise.
+        """
+        self.check_for_query_result_close()
+        return self._query_result.hasNext()
+
+    def get_next(self) -> list[Any] | dict[str, Any]:
+        """
+        Get the next row in the query result.
+
+        Returns
+        -------
+        list
+            Next row in the query result.
+
+        Raises
+        ------
+        Exception
+            If there are no more rows.
+        """
+        self.check_for_query_result_close()
+        row = self._query_result.getNext()
+        return _row_to_dict(self.columns, row) if self.as_dict else row
+
+    def get_all(self) -> list[list[Any] | dict[str, Any]]:
+        """
+        Get the next row in the query result.
+
+        Returns
+        -------
+        list
+            All remaining rows in the query result.
+        """
+        return list(self)
+
+    def get_n(self, count: int) -> list[list[Any] | dict[str, Any]]:
+        """
+        Get many rows in the query result.
+
+        Returns
+        -------
+        list
+            Up to `count` rows in the query result.
+        """
+        results = []
+        while self.has_next() and count > 0:
+            results.append(self.get_next())
+            count -= 1
+        return results
+
+    def close(self) -> None:
+        """Close the query result."""
+        if not self.is_closed:
+            # Allows the connection to be garbage collected if the query result
+            # is closed manually by the user.
+            self._query_result.close()
+            self.connection = None
+            self.is_closed = True
+
+    def check_for_query_result_close(self) -> None:
+        """
+        Check if the query result is closed and raise an exception if it is.
+
+        Raises
+        ------
+        Exception
+            If the query result is closed.
+
+        """
+        if self.is_closed:
+            msg = "Query result is closed"
+            raise RuntimeError(msg)
+
+    def get_as_df(self) -> pd.DataFrame:
+        """
+        Get the query result as a Pandas DataFrame.
+
+        See Also
+        --------
+        get_as_pl : Get the query result as a Polars DataFrame.
+        get_as_arrow : Get the query result as a PyArrow Table.
+
+        Returns
+        -------
+        pandas.DataFrame
+            Query result as a Pandas DataFrame.
+
+        """
+        self.check_for_query_result_close()
+
+        return self._query_result.getAsDF()
+
+    def get_as_pl(self) -> pl.DataFrame:
+        """
+        Get the query result as a Polars DataFrame.
+
+        See Also
+        --------
+        get_as_df : Get the query result as a Pandas DataFrame.
+        get_as_arrow : Get the query result as a PyArrow Table.
+
+        Returns
+        -------
+        polars.DataFrame
+            Query result as a Polars DataFrame.
+        """
+        import polars as pl
+
+        self.check_for_query_result_close()
+
+        # note: polars should always export just a single chunk,
+        # (eg: "-1") otherwise it will just need to rechunk anyway
+        return pl.from_arrow(  # type: ignore[return-value]
+            data=self.get_as_arrow(chunk_size=-1, fallbackExtensionTypes=True),
+        )
+
+    def get_as_arrow(self, chunk_size: int | None = None, *, fallbackExtensionTypes: bool = False) -> pa.Table:
+        """
+        Get the query result as a PyArrow Table.
+
+        Parameters
+        ----------
+        chunk_size : Number of rows to include in each chunk.
+            None
+                The chunk size is adaptive and depends on the number of columns in the query result.
+            -1 or 0
+                The entire result is returned as a single chunk.
+            > 0
+                The chunk size is the number of rows specified.
+
+        fallbackExtensionTypes : bool
+            Avoid using Arrow extension types for compatibility with Polars
+
+        See Also
+        --------
+        get_as_pl : Get the query result as a Polars DataFrame.
+        get_as_df : Get the query result as a Pandas DataFrame.
+
+        Returns
+        -------
+        pyarrow.Table
+            Query result as a PyArrow Table.
+        """
+        self.check_for_query_result_close()
+
+        if chunk_size is None:
+            # Adaptive; target 10m total elements in each chunk.
+            # (eg: if we had 10 cols, this would result in a 1m row chunk_size).
+            target_n_elems = 10_000_000
+            chunk_size = max(target_n_elems // len(self.get_column_names()), 10)
+        elif chunk_size <= 0:
+            # No chunking: return the entire result as a single chunk
+            chunk_size = self.get_num_tuples()
+
+        return self._query_result.getAsArrow(chunk_size, fallbackExtensionTypes)
+
+    def get_column_data_types(self) -> list[str]:
+        """
+        Get the data types of the columns in the query result.
+
+        Returns
+        -------
+        list
+            Data types of the columns in the query result.
+
+        """
+        self.check_for_query_result_close()
+        return self._query_result.getColumnDataTypes()
+
+    def get_column_names(self) -> list[str]:
+        """
+        Get the names of the columns in the query result.
+
+        Returns
+        -------
+        list
+            Names of the columns in the query result.
+
+        """
+        self.check_for_query_result_close()
+        return self._query_result.getColumnNames()
+
+    def get_schema(self) -> dict[str, str]:
+        """
+        Get the column schema of the query result.
+
+        Returns
+        -------
+        dict
+            Schema of the query result.
+
+        """
+        self.check_for_query_result_close()
+        return dict(
+            zip(
+                self._query_result.getColumnNames(),
+                self._query_result.getColumnDataTypes(),
+            )
+        )
+
+    def reset_iterator(self) -> None:
+        """Reset the iterator of the query result."""
+        self.check_for_query_result_close()
+        self._query_result.resetIterator()
+
+    def get_as_networkx(
+        self,
+        directed: bool = True,  # noqa: FBT001
+    ) -> nx.MultiGraph | nx.MultiDiGraph:
+        """
+        Convert the nodes and rels in query result into a NetworkX directed or undirected graph
+        with the following rules:
+        Columns with data type other than node or rel will be ignored.
+        Duplicated nodes and rels will be converted only once.
+
+        Parameters
+        ----------
+        directed : bool
+            Whether the graph should be directed. Defaults to True.
+
+        Returns
+        -------
+        networkx.MultiDiGraph or networkx.MultiGraph
+            Query result as a NetworkX graph.
+
+        """
+        self.check_for_query_result_close()
+        import networkx as nx
+
+        nx_graph = nx.MultiDiGraph() if directed else nx.MultiGraph()
+        properties_to_extract = self._get_properties_to_extract()
+
+        self.reset_iterator()
+
+        nodes = {}
+        rels = {}
+        table_to_label_dict = {}
+        table_primary_key_dict = {}
+
+        def encode_node_id(node: dict[str, Any], table_primary_key_dict: dict[str, Any]) -> str:
+            node_label = node[LABEL]
+            return f"{node_label}_{node[table_primary_key_dict[node_label]]!s}"
+
+        def encode_rel_id(rel: dict[str, Any]) -> tuple[int, int]:
+            return rel[ID]["table"], rel[ID]["offset"]
+
+        # De-duplicate nodes and rels
+        while self.has_next():
+            row = self.get_next()
+            for i in properties_to_extract:
+                # Skip empty nodes and rels, which may be returned by
+                # OPTIONAL MATCH
+                if row[i] is None or row[i] == {}:
+                    continue
+                column_type, _ = properties_to_extract[i]
+                if column_type == Type.NODE.value:
+                    nid = row[i][ID]
+                    nodes[nid["table"], nid["offset"]] = row[i]
+                    table_to_label_dict[nid["table"]] = row[i][LABEL]
+
+                elif column_type == Type.REL.value:
+                    rels[encode_rel_id(row[i])] = row[i]
+
+                elif column_type == Type.RECURSIVE_REL.value:
+                    for node in row[i][NODES]:
+                        nid = node[ID]
+                        nodes[nid["table"], nid["offset"]] = node
+                        table_to_label_dict[nid["table"]] = node[LABEL]
+                    for rel in row[i][RELS]:
+                        for key in list(rel.keys()):
+                            if rel[key] is None:
+                                del rel[key]
+                        rels[encode_rel_id(rel)] = rel
+
+        # Add nodes
+        for node in nodes.values():
+            nid = node[ID]
+            node_id = node[LABEL] + "_" + str(nid["offset"])
+            if node[LABEL] not in table_primary_key_dict:
+                props = self.connection._get_node_property_names(node[LABEL])
+                for prop_name in props:
+                    if props[prop_name]["is_primary_key"]:
+                        table_primary_key_dict[node[LABEL]] = prop_name
+                        break
+            node_id = encode_node_id(node, table_primary_key_dict)
+            node[node[LABEL]] = True
+            nx_graph.add_node(node_id, **node)
+
+        # Add rels
+        for rel in rels.values():
+            src = rel[SRC]
+            dst = rel[DST]
+            src_node = nodes[src["table"], src["offset"]]
+            dst_node = nodes[dst["table"], dst["offset"]]
+            src_id = encode_node_id(src_node, table_primary_key_dict)
+            dst_id = encode_node_id(dst_node, table_primary_key_dict)
+            nx_graph.add_edge(src_id, dst_id, **rel)
+        return nx_graph
+
+    def _get_properties_to_extract(self) -> dict[int, tuple[str, str]]:
+        column_names = self.get_column_names()
+        column_types = self.get_column_data_types()
+        properties_to_extract = {}
+
+        # Iterate over columns and extract nodes and rels, ignoring other columns
+        for i in range(len(column_names)):
+            column_name = column_names[i]
+            column_type = column_types[i]
+            if column_type in [
+                Type.NODE.value,
+                Type.REL.value,
+                Type.RECURSIVE_REL.value,
+            ]:
+                properties_to_extract[i] = (column_type, column_name)
+        return properties_to_extract
+
+    def get_as_torch_geometric(self) -> tuple[geo.Data | geo.HeteroData, dict, dict, dict]:  # type: ignore[type-arg]
+        """
+        Convert the nodes and rels in query result into a PyTorch Geometric graph representation
+        torch_geometric.data.Data or torch_geometric.data.HeteroData.
+
+        For node conversion, numerical and boolean properties are directly converted into tensor and
+        stored in Data/HeteroData. For properties cannot be converted into tensor automatically
+        (please refer to the notes below for more detail), they are returned as unconverted_properties.
+
+        For rel conversion, rel is converted into edge_index tensor director. Edge properties are returned
+        as edge_properties.
+
+        Node properties that cannot be converted into tensor automatically:
+        - If the type of a node property is not one of INT64, DOUBLE, or BOOL, it cannot be converted
+          automatically.
+        - If a node property contains a null value, it cannot be converted automatically.
+        - If a node property contains a nested list of variable length (e.g. [[1,2],[3]]), it cannot be
+          converted automatically.
+        - If a node property is a list or nested list, but the shape is inconsistent (e.g. the list length
+          is 6 for one node but 5 for another node), it cannot be converted automatically.
+
+        Additional conversion rules:
+        - Columns with data type other than node or rel will be ignored.
+        - Duplicated nodes and rels will be converted only once.
+
+        Returns
+        -------
+        torch_geometric.data.Data or torch_geometric.data.HeteroData
+            Query result as a PyTorch Geometric graph. Containing numeric or boolean node properties
+            and edge_index tensor.
+
+        dict
+            A dictionary that maps the positional offset of each node in Data/HeteroData to its primary
+            key in the database.
+
+        dict
+            A dictionary contains node properties that cannot be converted into tensor automatically. The
+            order of values for each property is aligned with nodes in Data/HeteroData.
+
+        dict
+            A dictionary contains edge properties. The order of values for each property is aligned with
+            edge_index in Data/HeteroData.
+        """
+        self.check_for_query_result_close()
+        # Despite we are not using torch_geometric in this file, we need to
+        # import it here to throw an error early if the user does not have
+        # torch_geometric or torch installed.
+
+        converter = TorchGeometricResultConverter(self)
+        return converter.get_as_torch_geometric()
+
+    def get_execution_time(self) -> int:
+        """
+        Get the time in ms which was required for executing the query.
+
+        Returns
+        -------
+        double
+            Query execution time as double in ms.
+
+        """
+        self.check_for_query_result_close()
+        return self._query_result.getExecutionTime()
+
+    def get_compiling_time(self) -> int:
+        """
+        Get the time in ms which was required for compiling the query.
+
+        Returns
+        -------
+        double
+            Query compile time as double in ms.
+
+        """
+        self.check_for_query_result_close()
+        return self._query_result.getCompilingTime()
+
+    def get_num_tuples(self) -> int:
+        """
+        Get the number of tuples which the query returned.
+
+        Returns
+        -------
+        int
+            Number of tuples.
+
+        """
+        self.check_for_query_result_close()
+        return self._query_result.getNumTuples()
+
+    def rows_as_dict(self, state=True) -> Self:
+        """
+        Change the format of the results, such that each row is a dict with the
+        column name as a key.
+
+        Parameters
+        ----------
+        state
+            Whether to turn dict formatting on or off. Turns it on by default.
+
+        Returns
+        -------
+        self
+            The object itself.
+
+        """
+        self.as_dict = state
+        if state:
+            self.columns = self.get_column_names()
+        return self
+
+
+def _row_to_dict(columns: list[str], row: list[Any]) -> dict[str, Any]:
+    if len(columns) != len(row):
+        msg = "Number of columns in output row does not match number of columns"
+        raise RuntimeError(msg)
+
+    return dict(zip(columns, row))

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

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import multiprocessing
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+from torch import Tensor
+from torch_geometric.data.feature_store import FeatureStore, TensorAttr
+
+from .connection import Connection
+from .types import Type
+
+if TYPE_CHECKING:
+    from torch_geometric.typing import FeatureTensorType
+
+    from .database import Database
+
+
+class LbugFeatureStore(FeatureStore):  # type: ignore[misc]
+    """Feature store compatible with `torch_geometric`."""
+
+    def __init__(self, db: Database, num_threads: int | None = None):
+        super().__init__()
+        self.db = db
+        self.connection: Connection = None  # type: ignore[assignment]
+        self.node_properties_cache: dict[str, dict[str, Any]] = {}
+        if num_threads is None:
+            num_threads = multiprocessing.cpu_count()
+        self.num_threads = num_threads
+
+    def __get_connection(self) -> Connection:
+        if self.connection is None:
+            self.connection = Connection(self.db, self.num_threads)
+        return self.connection
+
+    def _put_tensor(self, tensor: FeatureTensorType, attr: TensorAttr) -> bool:
+        raise NotImplementedError
+
+    def _get_tensor(self, attr: TensorAttr) -> FeatureTensorType | None:
+        table_name = attr.group_name
+        attr_name = attr.attr_name
+        attr_info = self.__get_node_property(table_name, attr_name)
+        if (
+            attr_info["type"]
+            not in [
+                Type.INT64.value,
+                Type.INT32.value,
+                Type.INT16.value,
+                Type.DOUBLE.value,
+                Type.FLOAT.value,
+                Type.BOOL.value,
+            ]
+        ) or (attr_info["dimension"] > 0 and "shape" not in attr_info):
+            return self.__get_tensor_by_query(attr)
+        return self.__get_tensor_by_scan(attr)
+
+    def __get_tensor_by_scan(self, attr: TensorAttr) -> FeatureTensorType | None:
+        table_name = attr.group_name
+        attr_name = attr.attr_name
+        indices = attr.index
+
+        if indices is None:
+            shape = self.get_tensor_size(attr)
+            if indices is None:
+                indices = np.arange(shape[0], dtype=np.uint64)
+        elif isinstance(indices, slice):
+            if indices.step is None or indices.step == 1:
+                indices = np.arange(indices.start, indices.stop, dtype=np.uint64)
+            else:
+                indices = np.arange(indices.start, indices.stop, indices.step, dtype=np.uint64)
+        elif isinstance(indices, int):
+            indices = np.array([indices])
+
+        if table_name not in self.node_properties_cache:
+            self.node_properties_cache[table_name] = self.connection._get_node_property_names(table_name)
+        attr_info = self.node_properties_cache[table_name][attr_name]
+
+        flat_dim = 1
+        if attr_info["dimension"] > 0:
+            for i in range(attr_info["dimension"]):
+                flat_dim *= attr_info["shape"][i]
+        scan_result = self.connection.database._scan_node_table(
+            table_name, attr_name, attr_info["type"], flat_dim, indices, self.num_threads
+        )
+
+        if attr_info["dimension"] > 0 and "shape" in attr_info:
+            result_shape = (len(indices),) + attr_info["shape"]
+            scan_result = scan_result.reshape(result_shape)
+
+        result = torch.from_numpy(scan_result)
+        return result
+
+    def __get_tensor_by_query(self, attr: TensorAttr) -> FeatureTensorType | None:
+        table_name = attr.group_name
+        attr_name = attr.attr_name
+        indices = attr.index
+
+        self.__get_connection()
+
+        match_clause = f"MATCH (item:{table_name})"
+        return_clause = f"RETURN item.{attr_name}"
+
+        if indices is None:
+            where_clause = ""
+        elif isinstance(indices, int):
+            where_clause = f"WHERE offset(id(item)) = {indices}"
+        elif isinstance(indices, slice):
+            if indices.step is None or indices.step == 1:
+                where_clause = f"WHERE offset(id(item)) >= {indices.start} AND offset(id(item)) < {indices.stop}"
+            else:
+                where_clause = (
+                    f"WHERE offset(id(item)) >= {indices.start} AND offset(id(item)) < {indices.stop} "
+                    f"AND (offset(id(item)) - {indices.start}) % {indices.step} = 0"
+                )
+        elif isinstance(indices, (Tensor, list, np.ndarray, tuple)):
+            where_clause = "WHERE"
+            for i in indices:
+                where_clause += f" offset(id(item)) = {int(i)} OR"
+            where_clause = where_clause[:-3]
+        else:
+            msg = f"Invalid attr.index type: {type(indices)!s}"
+            raise ValueError(msg)
+
+        query = f"{match_clause} {where_clause} {return_clause}"
+        res = self.connection.execute(query)
+
+        result_list = []
+        while res.has_next():
+            value_array = res.get_next()
+            if len(value_array) == 1:
+                value_array = value_array[0]
+            result_list.append(value_array)
+        try:
+            return torch.tensor(result_list)
+        except Exception:
+            return result_list
+
+    def _remove_tensor(self, attr: TensorAttr) -> bool:
+        raise NotImplementedError
+
+    def _get_tensor_size(self, attr: TensorAttr) -> tuple[Any, ...]:
+        self.__get_connection()
+        length_query = f"MATCH (item:{attr.group_name}) RETURN count(item)"
+        res = self.connection.execute(length_query)
+        length = res.get_next()[0]
+        attr_info = self.__get_node_property(attr.group_name, attr.attr_name)
+        if attr_info["dimension"] == 0:
+            return (length,)
+        else:
+            return (length,) + attr_info["shape"]
+
+    def __get_node_property(self, table_name: str, attr_name: str) -> dict[str, Any]:
+        if table_name in self.node_properties_cache and attr_name in self.node_properties_cache[table_name]:
+            return self.node_properties_cache[table_name][attr_name]
+        self.__get_connection()
+        if table_name not in self.node_properties_cache:
+            self.node_properties_cache[table_name] = self.connection._get_node_property_names(table_name)
+        if attr_name not in self.node_properties_cache[table_name]:
+            msg = f"Attribute {attr_name} not found in group {table_name}"
+            raise ValueError(msg)
+        attr_info = self.node_properties_cache[table_name][attr_name]
+        return attr_info
+
+    def get_all_tensor_attrs(self) -> list[TensorAttr]:
+        """Return all TensorAttr from the table nodes."""
+        result_list = []
+        self.__get_connection()
+        for table_name in self.connection._get_node_table_names():
+            if table_name not in self.node_properties_cache:
+                self.node_properties_cache[table_name] = self.connection._get_node_property_names(table_name)
+            for attr_name in self.node_properties_cache[table_name]:
+                if self.node_properties_cache[table_name][attr_name]["type"] in [
+                    Type.INT64.value,
+                    Type.INT32.value,
+                    Type.INT16.value,
+                    Type.DOUBLE.value,
+                    Type.FLOAT.value,
+                    Type.BOOL.value,
+                ] and (
+                    self.node_properties_cache[table_name][attr_name]["dimension"] == 0
+                    or "shape" in self.node_properties_cache[table_name][attr_name]
+                ):
+                    result_list.append(TensorAttr(table_name, attr_name))
+        return result_list