kumoai 2.13.0.dev202512040651__cp312-cp312-win_amd64.whl → 2.14.0.dev202512111731__cp312-cp312-win_amd64.whl

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
@@ -502,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'):
@@ -548,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,
@@ -884,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):
@@ -980,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/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 +