kumoai 2.13.0.dev202512040252__cp310-cp310-win_amd64.whl → 2.13.0.dev202512041731__cp310-cp310-win_amd64.whl

kumoai/__init__.py

@@ -280,7 +280,19 @@ __all__ = [
 ]
 
 
+def in_snowflake_notebook() -> bool:
+    try:
+        from snowflake.snowpark.context import get_active_session
+        import streamlit  # noqa: F401
+        get_active_session()
+        return True
+    except Exception:
+        return False
+
+
 def in_notebook() -> bool:
+    if in_snowflake_notebook():
+        return True
     try:
         from IPython import get_ipython
         shell = get_ipython()

kumoai/_version.py

@@ -1 +1 @@
-__version__ = '2.13.0.dev202512040252'
+__version__ = '2.13.0.dev202512041731'

kumoai/experimental/rfm/backend/snow/table.py

@@ -4,7 +4,7 @@ from typing import List, Optional, Sequence
 import pandas as pd
 from kumoapi.typing import Dtype
 
-from kumoai.experimental.rfm.backend.sqlite import Connection
+from kumoai.experimental.rfm.backend.snow import Connection
 from kumoai.experimental.rfm.base import SourceColumn, SourceForeignKey, Table
 
 
@@ -14,6 +14,8 @@ class SnowTable(Table):
     Args:
         connection: The connection to a :class:`snowflake` database.
         name: The name of this table.
+        database: The database.
+        schema: The schema.
         columns: The selected columns of this table.
         primary_key: The name of the primary key of this table, if it exists.
         time_column: The name of the time column of this table, if it exists.

kumoai/experimental/rfm/base/__init__.py

@@ -1,10 +1,12 @@
 from .source import SourceColumn, SourceForeignKey
 from .column import Column
 from .table import Table
+from .sampler import Sampler
 
 __all__ = [
     'SourceColumn',
     'SourceForeignKey',
     'Column',
     'Table',
+    'Sampler',
 ]

kumoai/experimental/rfm/base/sampler.py

@@ -0,0 +1,134 @@
+import copy
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pandas as pd
+from kumoapi.rfm.context import Subgraph
+from kumoapi.typing import Stype
+
+if TYPE_CHECKING:
+    from kumoai.experimental.rfm import Graph
+
+
+@dataclass
+class EdgeSpec:
+    num_neighbors: int | None = None
+    time_offsets: tuple[
+        pd.DateOffset | None,
+        pd.DateOffset,
+    ] | None = None
+
+    def __post_init__(self) -> None:
+        if (self.num_neighbors is None) == (self.time_offsets is None):
+            raise ValueError("Only one of 'num_neighbors' and 'time_offsets' "
+                             "must be provided")
+
+
+@dataclass
+class SamplerOutput:
+    df_dict: dict[str, pd.DataFrame]
+    batch_dict: dict[str, pd.DataFrame]
+    num_sampled_nodes_dict: dict[str, list[int]]
+    edge_index_dict: dict[tuple[str, str, str], np.ndarray] | None = None
+    num_sampled_edges_dict: dict[tuple[str, str, str], list[int]] | None = None
+
+
+class Sampler(ABC):
+    def __init__(self, graph: 'Graph') -> None:
+        self._edge_types: list[tuple[str, str, str]] = []
+        for edge in graph.edges:
+            edge_type = (edge.src_table, edge.fkey, edge.dst_table)
+            self._edge_types.append(edge_type)
+            self._edge_types.append(Subgraph.rev_edge_type(edge_type))
+
+        self._primary_key_dict: dict[str, str] = {
+            table.name: table._primary_key
+            for table in graph.tables.values()
+            if table._primary_key is not None
+        }
+
+        self._time_column_dict: dict[str, str] = {
+            table.name: table._time_column
+            for table in graph.tables.values()
+            if table._time_column is not None
+        }
+
+        foreign_keys = {(edge.src_table, edge.fkey) for edge in graph.edges}
+        self._stype_dict: dict[str, dict[str, Stype]] = {}
+        for table in graph.tables.values():
+            self._stype_dict[table.name] = {}
+            for column in table.columns:
+                if column == table.primary_key:
+                    continue
+                if (table.name, column.name) in foreign_keys:
+                    continue
+                self._stype_dict[table.name][column.name] = column.stype
+
+    @property
+    def edge_types(self) -> list[tuple[str, str, str]]:
+        return self._edge_types
+
+    @property
+    def primary_key_dict(self) -> dict[str, str]:
+        return self._primary_key_dict
+
+    @property
+    def time_column_dict(self) -> dict[str, str]:
+        return self._time_column_dict
+
+    @property
+    def stype_dict(self) -> dict[str, dict[str, Stype]]:
+        return self._stype_dict
+
+    def sample_subgraph(
+        self,
+        entity_table_names: tuple[str, ...],
+        entity_pkey: pd.Series,
+        anchor_time: pd.Series,
+        num_neighbors: list[int],
+        exclude_cols_dict: dict[str, list[str]] | None = None,
+    ) -> Subgraph:
+
+        edge_spec: dict[tuple[str, str, str], list[EdgeSpec]] = {
+            edge_type: [EdgeSpec(value) for value in num_neighbors]
+            for edge_type in self.edge_types
+        }
+
+        stype_dict: dict[str, dict[str, Stype]] = self._stype_dict
+        if exclude_cols_dict is not None:
+            stype_dict = copy.deepcopy(stype_dict)
+            for table_name, exclude_cols in exclude_cols_dict.items():
+                for column_name in exclude_cols:
+                    del stype_dict[table_name][column_name]
+
+        column_spec: dict[str, list[str]] = {
+            table_name: list(stypes.keys())
+            for table_name, stypes in stype_dict.items()
+        }
+        for table_name in entity_table_names:
+            column_spec[table_name].append(self.primary_key_dict[table_name])
+
+        return self._sample(
+            entity_table_name=entity_table_names[0],
+            entity_pkey=entity_pkey,
+            anchor_time=anchor_time,
+            column_spec=column_spec,
+            edge_spec=edge_spec,
+            return_edges=True,
+        )
+
+    # Abstract Methods ########################################################
+
+    @abstractmethod
+    def _sample(
+        self,
+        entity_table_name: str,
+        entity_pkey: pd.Series,
+        anchor_time: pd.Series,
+        column_spec: dict[str, list[str]],
+        edge_spec: dict[tuple[str, str, str], list[EdgeSpec]],
+        return_edges: bool = False,
+    ) -> SamplerOutput:
+        pass

kumoai/experimental/rfm/base/table.py

@@ -10,7 +10,7 @@ from kumoapi.table import TableDefinition
 from kumoapi.typing import Stype
 from typing_extensions import Self
 
-from kumoai import in_notebook
+from kumoai import in_notebook, in_snowflake_notebook
 from kumoai.experimental.rfm.base import Column, SourceColumn, SourceForeignKey
 from kumoai.experimental.rfm.infer import (
     contains_categorical,
@@ -384,7 +384,12 @@ class Table(ABC):
         if self._num_rows is not None:
             num_rows_repr = ' ({self._num_rows:,} rows)'
 
-        if in_notebook():
+        if in_snowflake_notebook():
+            import streamlit as st
+            md_repr = f"### 🏷️ Metadata of Table `{self.name}`{num_rows_repr}"
+            st.markdown(md_repr)
+            st.dataframe(self.metadata, hide_index=True)
+        elif in_notebook():
             from IPython.display import Markdown, display
             md_repr = f"### 🏷️ Metadata of Table `{self.name}`{num_rows_repr}"
             display(Markdown(md_repr))
@@ -496,7 +501,7 @@ class Table(ABC):
                 f'  end_time_column={self._end_time_column},\n'
                 f')')
 
-    # Abstract method #########################################################
+    # Abstract Methods ########################################################
 
     @cached_property
     def _source_column_dict(self) -> Dict[str, SourceColumn]:

kumoai/experimental/rfm/graph.py

@@ -3,7 +3,6 @@ import io
 import warnings
 from collections import defaultdict
 from dataclasses import dataclass, field
-from importlib.util import find_spec
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union
 
@@ -13,7 +12,7 @@ from kumoapi.table import TableDefinition
 from kumoapi.typing import Stype
 from typing_extensions import Self
 
-from kumoai import in_notebook
+from kumoai import in_notebook, in_snowflake_notebook
 from kumoai.experimental.rfm import Table
 from kumoai.graph import Edge
 from kumoai.mixin import CastMixin
@@ -251,6 +250,8 @@ class Graph:
     def from_snowflake(
         cls,
         connection: Union['SnowflakeConnection', Dict[str, Any], None] = None,
+        database: Optional[str] = None,
+        schema: Optional[str] = None,
         table_names: Optional[Sequence[str]] = None,
         edges: Optional[Sequence[Edge]] = None,
         infer_metadata: bool = True,
@@ -267,7 +268,7 @@ class Graph:
             >>> import kumoai.experimental.rfm as rfm
 
             >>> # Create a graph directly in a Snowflake notebook:
-            >>> graph = rfm.Graph.from_snowflake()
+            >>> graph = rfm.Graph.from_snowflake(schema='my_schema')
 
         Args:
             connection: An open connection from
@@ -276,6 +277,8 @@ class Graph:
                 connection. If ``None``, will re-use an active session in case
                 it exists, or create a new connection from credentials stored
                 in environment variables.
+            database: The database.
+            schema: The schema.
             table_names: Set of table names to include. If ``None``, will add
                 all tables present in the database.
             edges: An optional list of :class:`~kumoai.graph.Edge` objects to
@@ -297,17 +300,27 @@ class Graph:
 
         if table_names is None:
             with connection.cursor() as cursor:
-                cursor.execute("SELECT CURRENT_DATABASE(), CURRENT_SCHEMA()")
-                database, schema = cursor.fetchone()
-                query = f"""
+                if database is None and schema is None:
+                    cursor.execute("SELECT CURRENT_DATABASE(), "
+                                   "CURRENT_SCHEMA()")
+                    result = cursor.fetchone()
+                    database = database or result[0]
+                    schema = schema or result[1]
+                cursor.execute(f"""
                 SELECT TABLE_NAME
                 FROM {database}.INFORMATION_SCHEMA.TABLES
                 WHERE TABLE_SCHEMA = '{schema}'
-                """
-                cursor.execute(query)
+                """)
                 table_names = [row[0] for row in cursor.fetchall()]
 
-        tables = [SnowTable(connection, name) for name in table_names]
+        tables = [
+            SnowTable(
+                connection,
+                name=table_name,
+                database=database,
+                schema=schema,
+            ) for table_name in table_names
+        ]
 
         graph = cls(tables, edges=edges or [])
 
@@ -488,9 +501,13 @@ class Graph:
 
     def print_metadata(self) -> None:
         r"""Prints the :meth:`~Graph.metadata` of the graph."""
-        if in_notebook():
+        if in_snowflake_notebook():
+            import streamlit as st
+            st.markdown("### 🗂️ Graph Metadata")
+            st.dataframe(self.metadata, hide_index=True)
+        elif in_notebook():
             from IPython.display import Markdown, display
-            display(Markdown('### 🗂️ Graph Metadata'))
+            display(Markdown("### 🗂️ Graph Metadata"))
             df = self.metadata
             try:
                 if hasattr(df.style, 'hide'):
@@ -534,26 +551,36 @@ class Graph:
                   edge.src_table, edge.fkey) for edge in self.edges]
         edges = sorted(edges)
 
-        if in_notebook():
+        if in_snowflake_notebook():
+            import streamlit as st
+            st.markdown("### 🕸️ Graph Links (FK ↔️ PK)")
+            if len(edges) > 0:
+                st.markdown('\n'.join([
+                    f"- `{edge[2]}.{edge[3]}` ↔️ `{edge[0]}.{edge[1]}`"
+                    for edge in edges
+                ]))
+            else:
+                st.markdown("*No links registered*")
+        elif in_notebook():
             from IPython.display import Markdown, display
-            display(Markdown('### 🕸️ Graph Links (FK ↔️ PK)'))
+            display(Markdown("### 🕸️ Graph Links (FK ↔️ PK)"))
             if len(edges) > 0:
                 display(
                     Markdown('\n'.join([
-                        f'- `{edge[2]}.{edge[3]}` ↔️ `{edge[0]}.{edge[1]}`'
+                        f"- `{edge[2]}.{edge[3]}` ↔️ `{edge[0]}.{edge[1]}`"
                         for edge in edges
                     ])))
             else:
-                display(Markdown('*No links registered*'))
+                display(Markdown("*No links registered*"))
         else:
             print("🕸️ Graph Links (FK ↔️ PK):")
             if len(edges) > 0:
                 print('\n'.join([
-                    f'• {edge[2]}.{edge[3]} ↔️ {edge[0]}.{edge[1]}'
+                    f"• {edge[2]}.{edge[3]} ↔️ {edge[0]}.{edge[1]}"
                     for edge in edges
                 ]))
             else:
-                print('No links registered')
+                print("No links registered")
 
     def link(
         self,
@@ -870,19 +897,19 @@ class Graph:
 
             return True
 
-        # Check basic dependency:
-        if not find_spec('graphviz'):
-            raise ModuleNotFoundError("The 'graphviz' package is required for "
-                                      "visualization")
-        elif not has_graphviz_executables():
+        try:  # Check basic dependency:
+            import graphviz
+        except ImportError as e:
+            raise ImportError("The 'graphviz' package is required for "
+                              "visualization") from e
+
+        if not in_snowflake_notebook() and not has_graphviz_executables():
             raise RuntimeError("Could not visualize graph as 'graphviz' "
                                "executables are not installed. These "
                                "dependencies are required in addition to the "
                                "'graphviz' Python package. Please install "
                                "them as described at "
                                "https://graphviz.org/download/.")
-        else:
-            import graphviz
 
         format: Optional[str] = None
         if isinstance(path, str):
@@ -966,6 +993,9 @@ class Graph:
             graph.render(path, cleanup=True)
         elif isinstance(path, io.BytesIO):
             path.write(graph.pipe())
+        elif in_snowflake_notebook():
+            import streamlit as st
+            st.graphviz_chart(graph)
         elif in_notebook():
             from IPython.display import display
             display(graph)

{kumoai-2.13.0.dev202512040252.dist-info → kumoai-2.13.0.dev202512041731.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kumoai
-Version: 2.13.0.dev202512040252
+Version: 2.13.0.dev202512041731
 Summary: AI on the Modern Data Stack
 Author-email: "Kumo.AI" <hello@kumo.ai>
 License-Expression: MIT

{kumoai-2.13.0.dev202512040252.dist-info → kumoai-2.13.0.dev202512041731.dist-info}/RECORD

@@ -1,13 +1,13 @@
-kumoai/__init__.py,sha256=qu-qohU2cQlManX1aZIlzA3ivKl52m-cSQBPSW8urUU,10837
+kumoai/__init__.py,sha256=aDhb7KGetDnOz54u1Fd45zfM2N8oAha6XT2CvJqOvgc,11146
 kumoai/_logging.py,sha256=qL4JbMQwKXri2f-SEJoFB8TY5ALG12S-nobGTNWxW-A,915
 kumoai/_singleton.py,sha256=i2BHWKpccNh5SJGDyU0IXsnYzJAYr8Xb0wz4c6LRbpo,861
-kumoai/_version.py,sha256=16u1rVm-N2IEE7QbyS9U5nn_hjp7P_wxBIQzzAKSnDA,39
+kumoai/_version.py,sha256=nAPj2gzW1byjXavAIRphLnIsySItnJOq_dJijgrT4kk,39
 kumoai/databricks.py,sha256=ahwJz6DWLXMkndT0XwEDBxF-hoqhidFR8wBUQ4TLZ68,490
 kumoai/exceptions.py,sha256=7TMs0SC8xrU009_Pgd4QXtSF9lxJq8MtRbeX9pcQUy4,859
 kumoai/formatting.py,sha256=o3uCnLwXPhe1KI5WV9sBgRrcU7ed4rgu_pf89GL9Nc0,983
 kumoai/futures.py,sha256=J8rtZMEYFzdn5xF_x-LAiKJz3KGL6PT02f6rq_2bOJk,3836
 kumoai/jobs.py,sha256=dCi7BAdfm2tCnonYlGU4WJokJWbh3RzFfaOX2EYCIHU,2576
-kumoai/kumolib.cp310-win_amd64.pyd,sha256=3iE0thfrVDx0Yhh0I0li-BwZcIpQfRpaYxYMsSpYofc,194048
+kumoai/kumolib.cp310-win_amd64.pyd,sha256=4fWOeR-nrmdLzThTLrD0wexcVo5F4KECAdbjPSHPydM,194048
 kumoai/mixin.py,sha256=IaiB8SAI0VqOoMVzzIaUlqMt53-QPUK6OB0HikG-V9E,840
 kumoai/spcs.py,sha256=KWfENrwSLruprlD-QPh63uU0N6npiNrwkeKfBk3EUyQ,4260
 kumoai/artifact_export/__init__.py,sha256=UXAQI5q92ChBzWAk8o3J6pElzYHudAzFZssQXd4o7i8,247
@@ -55,7 +55,7 @@ kumoai/encoder/__init__.py,sha256=8FeP6mUyCeXxr1b8kUIi5dxe5vEXQRft9tPoaV1CBqg,18
 kumoai/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kumoai/experimental/rfm/__init__.py,sha256=EFZz6IvvskmeO85Vig6p1m_6jdimS_BkeREOndHuRsc,6247
 kumoai/experimental/rfm/authenticate.py,sha256=G89_4TMeUpr5fG_0VTzMF5sdNhaciitA1oc2loTlTmo,19321
-kumoai/experimental/rfm/graph.py,sha256=kSWve-Fn_9qERFjEpCDO5zDnngtd9T4MOhR_o46PI7s,39602
+kumoai/experimental/rfm/graph.py,sha256=SL3-WinoLnkZC6VVjebYGLuQJJyEVFJdCm6h3FNE0e4,40816
 kumoai/experimental/rfm/local_graph_sampler.py,sha256=dQ3JnuozTNeZyUFRu2h8OTMNmV1RAoaCA0gvkpgOstg,8110
 kumoai/experimental/rfm/local_graph_store.py,sha256=6jY1ciVIlnBBhZCxWwBTl7SKX1fxRIDLszwrftD0Cdk,13485
 kumoai/experimental/rfm/local_pquery_driver.py,sha256=Yd_yHIrvuDj16IC1pvsqiQvZS41vvOOCRMiuDGtN6Fk,26851
@@ -65,13 +65,14 @@ kumoai/experimental/rfm/backend/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm
 kumoai/experimental/rfm/backend/local/__init__.py,sha256=usMh0fuDxKK-aOVT1sU30BQWFS0eSkfUrhUVILisQQI,934
 kumoai/experimental/rfm/backend/local/table.py,sha256=1PqNOROzlnK3SaZHNcU2hyzeifs0N4wssQAS3-Z0Myc,3674
 kumoai/experimental/rfm/backend/snow/__init__.py,sha256=viMeR9VWpB1kjRdSWCTNFMdM7a8Mj_Dtck1twJW8dV8,962
-kumoai/experimental/rfm/backend/snow/table.py,sha256=HVrPtCVvfsisFmq9jMovowsE5Wl5oti3O-kru7ruXlc,4312
+kumoai/experimental/rfm/backend/snow/table.py,sha256=Rf4hUPOUtsjpaIc9vBKWPZ3yz20OOg6DZqCGeih4KC8,4372
 kumoai/experimental/rfm/backend/sqlite/__init__.py,sha256=xw5NNLrWSvUvRkD49X_9hZYjas5EuP1XDANPy0EEjOg,874
 kumoai/experimental/rfm/backend/sqlite/table.py,sha256=mBiZC21gQwfR4demFrP37GmawMHfIm-G82mLQeBqIZo,3901
-kumoai/experimental/rfm/base/__init__.py,sha256=oXPkeBemtuDxRUK61-0sOT84GZB_oQ6HvaZNU1KFNaw,199
+kumoai/experimental/rfm/base/__init__.py,sha256=U4GMz-sHnagddIQEqhcfpY9JRhmOc0_Rs9c5_Zc9rqU,245
 kumoai/experimental/rfm/base/column.py,sha256=OE-PRQ8HO4uTq0e3_3eHJFfhp5nzw79zd-43g3iMh4g,2385
+kumoai/experimental/rfm/base/sampler.py,sha256=E0CC5Rfyjd8SjyjK1bCKKYDesPGYYbIiY96yW6xOXww,4563
 kumoai/experimental/rfm/base/source.py,sha256=H5yN9xAwK3i_69EdqOV_x58muPGKQiI8ev5BhHQDZEo,290
-kumoai/experimental/rfm/base/table.py,sha256=glyAg4LCQdddM3lIRClJSA7qMyfoHUVAGBf1rEs6B8Y,20113
+kumoai/experimental/rfm/base/table.py,sha256=6GlWUz3tAu2g1QqTA5idWGmfo2KEJoNApDFZRn8e0pg,20388
 kumoai/experimental/rfm/infer/__init__.py,sha256=qKg8or-SpgTApD6ePw1PJ4aUZPrOLTHLRCmBIJ92hrk,486
 kumoai/experimental/rfm/infer/categorical.py,sha256=bqmfrE5ZCBTcb35lA4SyAkCu3MgttAn29VBJYMBNhVg,893
 kumoai/experimental/rfm/infer/dtype.py,sha256=Hf_drluYNuN59lTSe-8GuXalg20Pv93kCktB6Hb9f74,2686
@@ -104,8 +105,8 @@ kumoai/utils/__init__.py,sha256=wAKgmwtMIGuiauW9D_GGKH95K-24Kgwmld27mm4nsro,278
 kumoai/utils/datasets.py,sha256=UyAII-oAn7x3ombuvpbSQ41aVF9SYKBjQthTD-vcT2A,3011
 kumoai/utils/forecasting.py,sha256=ZgKeUCbWLOot0giAkoigwU5du8LkrwAicFOi5hVn6wg,7624
 kumoai/utils/progress_logger.py,sha256=MZsWgHd4UZQKCXiJZgQeW-Emi_BmzlCKPLPXOL_HqBo,5239
-kumoai-2.13.0.dev202512040252.dist-info/licenses/LICENSE,sha256=ZUilBDp--4vbhsEr6f_Upw9rnIx09zQ3K9fXQ0rfd6w,1111
-kumoai-2.13.0.dev202512040252.dist-info/METADATA,sha256=T-O--qEm_2QPzB-dDkwR6Ei7r79H7v6qQBbR1e1J8gg,2580
-kumoai-2.13.0.dev202512040252.dist-info/WHEEL,sha256=KUuBC6lxAbHCKilKua8R9W_TM71_-9Sg5uEP3uDWcoU,101
-kumoai-2.13.0.dev202512040252.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
-kumoai-2.13.0.dev202512040252.dist-info/RECORD,,
+kumoai-2.13.0.dev202512041731.dist-info/licenses/LICENSE,sha256=ZUilBDp--4vbhsEr6f_Upw9rnIx09zQ3K9fXQ0rfd6w,1111
+kumoai-2.13.0.dev202512041731.dist-info/METADATA,sha256=H4PwggDYmXKPTArJoXK3Lxg4__QTOIUb0nZRR2heOpM,2580
+kumoai-2.13.0.dev202512041731.dist-info/WHEEL,sha256=KUuBC6lxAbHCKilKua8R9W_TM71_-9Sg5uEP3uDWcoU,101
+kumoai-2.13.0.dev202512041731.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
+kumoai-2.13.0.dev202512041731.dist-info/RECORD,,