real-ladybug 0.0.1.dev1__cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

real_ladybug/database.py

@@ -0,0 +1,307 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from . import _lbug
+from .types import Type
+
+if TYPE_CHECKING:
+    import sys
+    from types import TracebackType
+
+    from numpy.typing import NDArray
+    from torch_geometric.data.feature_store import IndexType
+
+    from .torch_geometric_feature_store import LbugFeatureStore
+    from .torch_geometric_graph_store import LbugGraphStore
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
+
+class Database:
+    """Lbug database instance."""
+
+    def __init__(
+        self,
+        database_path: str | Path | None = None,
+        *,
+        buffer_pool_size: int = 0,
+        max_num_threads: int = 0,
+        compression: bool = True,
+        lazy_init: bool = False,
+        read_only: bool = False,
+        max_db_size: int = (1 << 43),
+        auto_checkpoint: bool = True,
+        checkpoint_threshold: int = -1,
+        throw_on_wal_replay_failure: bool = True,
+        enable_checksums: bool = True,
+    ):
+        """
+        Parameters
+        ----------
+        database_path : str, Path
+            The path to database files. If the path is not specified, or empty, or equal to `:memory:`, the database
+            will be created in memory.
+
+        buffer_pool_size : int
+            The maximum size of buffer pool in bytes. Defaults to ~80% of system memory.
+
+        max_num_threads : int
+            The maximum number of threads to use for executing queries.
+
+        compression : bool
+            Enable database compression.
+
+        lazy_init : bool
+            If True, the database will not be initialized until the first query.
+            This is useful when the database is not used in the main thread or
+            when the main process is forked.
+            Default to False.
+
+        read_only : bool
+            If true, the database is opened read-only. No write transactions is
+            allowed on the `Database` object. Multiple read-only `Database`
+            objects can be created with the same database path. However, there
+            cannot be multiple `Database` objects created with the same
+            database path.
+            Default to False.
+
+        max_db_size : int
+            The maximum size of the database in bytes. Note that this is introduced
+            temporarily for now to get around with the default 8TB mmap address
+             space limit some environment. This will be removed once we implemente
+             a better solution later. The value is default to 1 << 43 (8TB) under 64-bit
+             environment and 1GB under 32-bit one.
+
+        auto_checkpoint: bool
+            If true, the database will automatically checkpoint when the size of
+            the WAL file exceeds the checkpoint threshold.
+
+        checkpoint_threshold: int
+            The threshold of the WAL file size in bytes. When the size of the
+            WAL file exceeds this threshold, the database will checkpoint if autoCheckpoint is true.
+
+        throw_on_wal_replay_failure: bool
+            If true, any WAL replaying failure when loading the database will throw an error.
+            Otherwise, Lbug will silently ignore the failure and replay up to where the error
+            occured.
+
+        enable_checksums: bool
+            If true, the database will use checksums to detect corruption in the
+            WAL file.
+
+        """
+        if database_path is None:
+            database_path = ":memory:"
+        if isinstance(database_path, Path):
+            database_path = str(database_path)
+
+        self.database_path = database_path
+        self.buffer_pool_size = buffer_pool_size
+        self.max_num_threads = max_num_threads
+        self.compression = compression
+        self.read_only = read_only
+        self.max_db_size = max_db_size
+        self.auto_checkpoint = auto_checkpoint
+        self.checkpoint_threshold = checkpoint_threshold
+        self.throw_on_wal_replay_failure = throw_on_wal_replay_failure
+        self.enable_checksums = enable_checksums
+        self.is_closed = False
+
+        self._database: Any = None  # (type: _lbug.Database from pybind11)
+        if not lazy_init:
+            self.init_database()
+
+    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()
+
+    @staticmethod
+    def get_version() -> str:
+        """
+        Get the version of the database.
+
+        Returns
+        -------
+        str
+            The version of the database.
+        """
+        return _lbug.Database.get_version()  # type: ignore[union-attr]
+
+    @staticmethod
+    def get_storage_version() -> int:
+        """
+        Get the storage version of the database.
+
+        Returns
+        -------
+        int
+            The storage version of the database.
+        """
+        return _lbug.Database.get_storage_version()  # type: ignore[union-attr]
+
+    def __getstate__(self) -> dict[str, Any]:
+        state = {
+            "database_path": self.database_path,
+            "buffer_pool_size": self.buffer_pool_size,
+            "compression": self.compression,
+            "read_only": self.read_only,
+            "_database": None,
+        }
+        return state
+
+    def init_database(self) -> None:
+        """Initialize the database."""
+        self.check_for_database_close()
+        if self._database is None:
+            self._database = _lbug.Database(  # type: ignore[union-attr]
+                self.database_path,
+                self.buffer_pool_size,
+                self.max_num_threads,
+                self.compression,
+                self.read_only,
+                self.max_db_size,
+                self.auto_checkpoint,
+                self.checkpoint_threshold,
+                self.throw_on_wal_replay_failure,
+                self.enable_checksums,
+            )
+
+    def get_torch_geometric_remote_backend(
+        self, num_threads: int | None = None
+    ) -> tuple[LbugFeatureStore, LbugGraphStore]:
+        """
+        Use the database as the remote backend for torch_geometric.
+
+        For the interface of the remote backend, please refer to
+        https://pytorch-geometric.readthedocs.io/en/latest/advanced/remote.html.
+        The current implementation is read-only and does not support edge
+        features. The IDs of the nodes are based on the internal IDs (i.e., node
+        offsets). For the remote node IDs to be consistent with the positions in
+        the output tensors, please ensure that no deletion has been performed
+        on the node tables.
+
+        The remote backend can also be plugged into the data loader of
+        torch_geometric, which is useful for mini-batch training. For example:
+
+        ```python
+            loader_lbug = NeighborLoader(
+                data=(feature_store, graph_store),
+                num_neighbors={('paper', 'cites', 'paper'): [12, 12, 12]},
+                batch_size=LOADER_BATCH_SIZE,
+                input_nodes=('paper', input_nodes),
+                num_workers=4,
+                filter_per_worker=False,
+            )
+        ```
+
+        Please note that the database instance is not fork-safe, so if more than
+        one worker is used, `filter_per_worker` must be set to False.
+
+        Parameters
+        ----------
+        num_threads : int
+            Number of threads to use for data loading. Default to None, which
+            means using the number of CPU cores.
+
+        Returns
+        -------
+        feature_store : LbugFeatureStore
+            Feature store compatible with torch_geometric.
+        graph_store : LbugGraphStore
+            Graph store compatible with torch_geometric.
+        """
+        self.check_for_database_close()
+        from .torch_geometric_feature_store import LbugFeatureStore
+        from .torch_geometric_graph_store import LbugGraphStore
+
+        return (
+            LbugFeatureStore(self, num_threads),
+            LbugGraphStore(self, num_threads),
+        )
+
+    def _scan_node_table(
+        self,
+        table_name: str,
+        prop_name: str,
+        prop_type: str,
+        dim: int,
+        indices: IndexType,
+        num_threads: int,
+    ) -> NDArray[Any]:
+        self.check_for_database_close()
+        import numpy as np
+
+        """
+        Scan a node table from storage directly, bypassing query engine.
+        Used internally by torch_geometric remote backend only.
+        """
+        self.init_database()
+        indices_cast = np.array(indices, dtype=np.uint64)
+        result = None
+
+        if prop_type == Type.INT64.value:
+            result = np.empty(len(indices) * dim, dtype=np.int64)
+            self._database.scan_node_table_as_int64(table_name, prop_name, indices_cast, result, num_threads)
+        elif prop_type == Type.INT32.value:
+            result = np.empty(len(indices) * dim, dtype=np.int32)
+            self._database.scan_node_table_as_int32(table_name, prop_name, indices_cast, result, num_threads)
+        elif prop_type == Type.INT16.value:
+            result = np.empty(len(indices) * dim, dtype=np.int16)
+            self._database.scan_node_table_as_int16(table_name, prop_name, indices_cast, result, num_threads)
+        elif prop_type == Type.DOUBLE.value:
+            result = np.empty(len(indices) * dim, dtype=np.float64)
+            self._database.scan_node_table_as_double(table_name, prop_name, indices_cast, result, num_threads)
+        elif prop_type == Type.FLOAT.value:
+            result = np.empty(len(indices) * dim, dtype=np.float32)
+            self._database.scan_node_table_as_float(table_name, prop_name, indices_cast, result, num_threads)
+
+        if result is not None:
+            return result
+
+        msg = f"Unsupported property type: {prop_type}"
+        raise ValueError(msg)
+
+    def close(self) -> None:
+        """
+        Close the database. Once the database is closed, the lock on the database
+        files is released and the database can be opened in another process.
+
+        Note: Call to this method is not required. The Python garbage collector
+        will automatically close the database when no references to the database
+        object exist. It is recommended not to call this method explicitly. If you
+        decide to manually close the database, make sure that all the QueryResult
+        and Connection objects are closed before calling this method.
+        """
+        if self.is_closed:
+            return
+        self.is_closed = True
+        if self._database is not None:
+            self._database.close()
+            self._database: Any = None  # (type: _lbug.Database from pybind11)
+
+    def check_for_database_close(self) -> None:
+        """
+        Check if the database is closed and raise an exception if it is.
+
+        Raises
+        ------
+        Exception
+            If the database is closed.
+
+        """
+        if not self.is_closed:
+            return
+        msg = "Database is closed"
+        raise RuntimeError(msg)

real_ladybug/prepared_statement.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from .connection import Connection
+
+
+class PreparedStatement:
+    """
+    A prepared statement is a parameterized query which can avoid planning the
+    same query for repeated execution.
+    """
+
+    def __init__(self, connection: Connection, query: str, parameters: dict[str, Any] | None = None):
+        """
+        Parameters
+        ----------
+        connection : Connection
+            Connection to a database.
+        query : str
+            Query to prepare.
+        parameters : dict[str, Any]
+            Parameters for the query.
+        """
+        if parameters is None:
+            parameters = {}
+        self._prepared_statement = connection._connection.prepare(query, parameters)
+        self._connection = connection
+
+    def is_success(self) -> bool:
+        """
+        Check if the prepared statement is successfully prepared.
+
+        Returns
+        -------
+        bool
+            True if the prepared statement is successfully prepared.
+        """
+        return self._prepared_statement.is_success()
+
+    def get_error_message(self) -> str:
+        """
+        Get the error message if the query is not prepared successfully.
+
+        Returns
+        -------
+        str
+            Error message.
+        """
+        return self._prepared_statement.get_error_message()