kumoai 2.13.0.dev202512061731__cp313-cp313-macosx_11_0_arm64.whl → 2.14.0.dev202512151351__cp313-cp313-macosx_11_0_arm64.whl

kumoai/experimental/rfm/base/source.py

@@ -9,6 +9,7 @@ class SourceColumn:
     dtype: Dtype
     is_primary_key: bool
     is_unique_key: bool
+    is_nullable: bool
 
 
 @dataclass

kumoai/experimental/rfm/base/sql_sampler.py

@@ -0,0 +1,56 @@
+from abc import abstractmethod
+from typing import Literal
+
+import numpy as np
+import pandas as pd
+
+from kumoai.experimental.rfm.base import Sampler, SamplerOutput
+
+
+class SQLSampler(Sampler):
+    def _sample_subgraph(
+        self,
+        entity_table_name: str,
+        entity_pkey: pd.Series,
+        anchor_time: pd.Series | Literal['entity'],
+        columns_dict: dict[str, set[str]],
+        num_neighbors: list[int],
+    ) -> SamplerOutput:
+
+        df, batch = self._by_pkey(
+            table_name=entity_table_name,
+            pkey=entity_pkey,
+            columns=columns_dict[entity_table_name],
+        )
+        if len(batch) != len(entity_pkey):
+            raise KeyError("Invalid primary keys")  # TODO
+
+        perm = batch.argsort()
+        batch = batch[perm]
+        df = df.iloc[perm].reset_index(drop=True)
+
+        if not isinstance(anchor_time, pd.Series):
+            time_column = self.time_column_dict[entity_table_name]
+            anchor_time = df[time_column]
+
+        return SamplerOutput(
+            anchor_time=anchor_time.astype(int).to_numpy(),
+            df_dict={entity_table_name: df},
+            inverse_dict={},
+            batch_dict={entity_table_name: batch},
+            num_sampled_nodes_dict={entity_table_name: [len(batch)]},
+            row_dict={},
+            col_dict={},
+            num_sampled_edges_dict={},
+        )
+
+    # Abstract Methods ########################################################
+
+    @abstractmethod
+    def _by_pkey(
+        self,
+        table_name: str,
+        pkey: pd.Series,
+        columns: set[str],
+    ) -> tuple[pd.DataFrame, np.ndarray]:
+        pass

kumoai/experimental/rfm/base/table.py

@@ -11,7 +11,12 @@ from kumoapi.typing import Stype
 from typing_extensions import Self
 
 from kumoai import in_notebook, in_snowflake_notebook
-from kumoai.experimental.rfm.base import Column, SourceColumn, SourceForeignKey
+from kumoai.experimental.rfm.base import (
+    Column,
+    DataBackend,
+    SourceColumn,
+    SourceForeignKey,
+)
 from kumoai.experimental.rfm.infer import (
     contains_categorical,
     contains_id,
@@ -503,6 +508,12 @@ class Table(ABC):
 
     # Abstract Methods ########################################################
 
+    @property
+    @abstractmethod
+    def backend(self) -> DataBackend:
+        r"""The data backend of this table."""
+        pass
+
     @cached_property
     def _source_column_dict(self) -> Dict[str, SourceColumn]:
         return {col.name: col for col in self._get_source_columns()}

kumoai/experimental/rfm/graph.py

@@ -13,7 +13,7 @@ from kumoapi.typing import Stype
 from typing_extensions import Self
 
 from kumoai import in_notebook, in_snowflake_notebook
-from kumoai.experimental.rfm import Table
+from kumoai.experimental.rfm.base import DataBackend, Table
 from kumoai.graph import Edge
 from kumoai.mixin import CastMixin
 
@@ -218,10 +218,12 @@ class Graph:
             connect,
         )
 
+        internal_connection = False
         if not isinstance(connection, Connection):
             connection = SqliteConnectionConfig._cast(connection)
             assert isinstance(connection, SqliteConnectionConfig)
             connection = connect(connection.uri, **connection.kwargs)
+            internal_connection = True
         assert isinstance(connection, Connection)
 
         if table_names is None:
@@ -234,6 +236,9 @@ class Graph:
 
         graph = cls(tables, edges=edges or [])
 
+        if internal_connection:
+            graph._connection = connection  # type: ignore
+
         if infer_metadata:
             graph.infer_metadata(False)
 
@@ -394,7 +399,14 @@ class Graph:
 
         return graph
 
-    # Tables ##############################################################
+    # Backend #################################################################
+
+    @property
+    def backend(self) -> DataBackend | None:
+        backends = [table.backend for table in self._tables.values()]
+        return backends[0] if len(backends) > 0 else None
+
+    # Tables ##################################################################
 
     def has_table(self, name: str) -> bool:
         r"""Returns ``True`` if the graph has a table with name ``name``;
@@ -433,13 +445,10 @@ class Graph:
             raise KeyError(f"Cannot add table with name '{table.name}' to "
                            f"this graph; table names must be globally unique.")
 
-        if len(self._tables) > 0:
-            cls = next(iter(self._tables.values())).__class__
-            if table.__class__ != cls:
-                raise ValueError(f"Cannot register a "
-                                 f"'{table.__class__.__name__}' to this "
-                                 f"graph since other tables are of type "
-                                 f"'{cls.__name__}'.")
+        if self.backend is not None and table.backend != self.backend:
+            raise ValueError(f"Cannot register a table with backend "
+                             f"'{table.backend}' to this graph since other "
+                             f"tables have backend '{self.backend}'.")
 
         self._tables[table.name] = table
 
@@ -826,6 +835,10 @@ class Graph:
             raise ValueError("At least one table needs to be added to the "
                              "graph")
 
+        backends = {table.backend for table in self._tables.values()}
+        if len(backends) != 1:
+            raise ValueError("Found multiple table backends in the graph")
+
         for edge in self.edges:
             src_table, fkey, dst_table = edge
 
@@ -1063,3 +1076,7 @@ class Graph:
                 f'  tables={tables},\n'
                 f'  edges={edges},\n'
                 f')')
+
+    def __del__(self) -> None:
+        if hasattr(self, '_connection'):
+            self._connection.close()

kumoai/experimental/rfm/pquery/pandas_executor.py

@@ -134,7 +134,7 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
         outs: List[pd.Series] = []
         masks: List[np.ndarray] = []
         for _ in range(num_forecasts):
-            anchor_target_time = anchor_time[target_batch]
+            anchor_target_time = anchor_time.iloc[target_batch]
             anchor_target_time = anchor_target_time.reset_index(drop=True)
 
             time_filter_mask = (target_time <= anchor_target_time +