kumoai 2.14.0.dev202512271732__cp310-cp310-macosx_11_0_arm64.whl → 2.14.0rc2__cp310-cp310-macosx_11_0_arm64.whl

kumoai/experimental/rfm/graph.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 import copy
 import io
@@ -19,6 +21,7 @@ from kumoai import in_notebook, in_snowflake_notebook
 from kumoai.experimental.rfm.base import ColumnSpec, DataBackend, Table
 from kumoai.graph import Edge
 from kumoai.mixin import CastMixin
+from kumoai.utils import display
 
 if TYPE_CHECKING:
     import graphviz
@@ -536,6 +539,35 @@ class Graph:
 
         return graph
 
+    @classmethod
+    def from_relbench(
+        cls,
+        dataset: str,
+        verbose: bool = True,
+    ) -> Graph:
+        r"""Loads a `RelBench <https://relbench.stanford.edu>`_ dataset into a
+        :class:`Graph` instance.
+
+        .. code-block:: python
+
+            >>> # doctest: +SKIP
+            >>> import kumoai.experimental.rfm as rfm
+
+            >>> graph = rfm.Graph.from_relbench("f1")
+
+        Args:
+            dataset: The RelBench dataset name.
+            verbose: Whether to print verbose output.
+        """
+        from kumoai.experimental.rfm.relbench import from_relbench
+        graph = from_relbench(dataset, verbose=verbose)
+
+        if verbose:
+            graph.print_metadata()
+            graph.print_links()
+
+        return graph
+
     # Backend #################################################################
 
     @property
@@ -617,28 +649,28 @@ class Graph:
         r"""Returns a :class:`pandas.DataFrame` object containing metadata
         information about the tables in this graph.
 
-        The returned dataframe has columns ``name``, ``primary_key``,
-        ``time_column``, and ``end_time_column``, which provide an aggregate
-        view of the properties of the tables of this graph.
+        The returned dataframe has columns ``"Name"``, ``"Primary Key"``,
+        ``"Time Column"``, and ``"End Time Column"``, which provide an
+        aggregated view of the properties of the tables of this graph.
 
         Example:
             >>> # doctest: +SKIP
             >>> import kumoai.experimental.rfm as rfm
             >>> graph = rfm.Graph(tables=...).infer_metadata()
             >>> graph.metadata  # doctest: +SKIP
-                name   primary_key  time_column  end_time_column
-            0   users      user_id            -                -
+                Name   Primary Key  Time Column  End Time Column
+            0   users  user_id      -            -
         """
         tables = list(self.tables.values())
 
         return pd.DataFrame({
-            'name':
+            'Name':
             pd.Series(dtype=str, data=[t.name for t in tables]),
-            'primary_key':
+            'Primary Key':
             pd.Series(dtype=str, data=[t._primary_key or '-' for t in tables]),
-            'time_column':
+            'Time Column':
             pd.Series(dtype=str, data=[t._time_column or '-' for t in tables]),
-            'end_time_column':
+            'End Time Column':
             pd.Series(
                 dtype=str,
                 data=[t._end_time_column or '-' for t in tables],
@@ -647,24 +679,8 @@ class Graph:
 
     def print_metadata(self) -> None:
         r"""Prints the :meth:`~Graph.metadata` of the graph."""
-        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"))
-            df = self.metadata
-            try:
-                if hasattr(df.style, 'hide'):
-                    display(df.style.hide(axis='index'))  # pandas=2
-                else:
-                    display(df.style.hide_index())  # pandas<1.3
-            except ImportError:
-                print(df.to_string(index=False))  # missing jinja2
-        else:
-            print("🗂️ Graph Metadata:")
-            print(self.metadata.to_string(index=False))
+        display.title("🗂️ Graph Metadata")
+        display.dataframe(self.metadata)
 
     def infer_metadata(self, verbose: bool = True) -> Self:
         r"""Infers metadata for all tables in the graph.
@@ -693,40 +709,21 @@ class Graph:
 
     def print_links(self) -> None:
         r"""Prints the :meth:`~Graph.edges` of the graph."""
-        edges = [(edge.dst_table, self[edge.dst_table]._primary_key,
-                  edge.src_table, edge.fkey) for edge in self.edges]
-        edges = sorted(edges)
-
-        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)"))
-            if len(edges) > 0:
-                display(
-                    Markdown('\n'.join([
-                        f"- `{edge[2]}.{edge[3]}` ↔️ `{edge[0]}.{edge[1]}`"
-                        for edge in edges
-                    ])))
-            else:
-                display(Markdown("*No links registered*"))
+        edges = sorted([(
+            edge.dst_table,
+            self[edge.dst_table]._primary_key,
+            edge.src_table,
+            edge.fkey,
+        ) for edge in self.edges])
+
+        display.title("🕸️ Graph Links (FK ↔️ PK)")
+        if len(edges) > 0:
+            display.unordered_list(items=[
+                f"`{edge[2]}.{edge[3]}` ↔️ `{edge[0]}.{edge[1]}`"
+                for edge in edges
+            ])
         else:
-            print("🕸️ Graph Links (FK ↔️ PK):")
-            if len(edges) > 0:
-                print('\n'.join([
-                    f"• {edge[2]}.{edge[3]} ↔️ {edge[0]}.{edge[1]}"
-                    for edge in edges
-                ]))
-            else:
-                print("No links registered")
+            display.italic("No links registered")
 
     def link(
         self,

kumoai/experimental/rfm/infer/dtype.py

@@ -52,7 +52,8 @@ def infer_dtype(ser: pd.Series) -> Dtype:
             ser = pd.Series(arr, dtype=pd.ArrowDtype(arr.type))
 
     if isinstance(ser.dtype, pd.ArrowDtype):
-        if pa.types.is_list(ser.dtype.pyarrow_dtype):
+        if (pa.types.is_list(ser.dtype.pyarrow_dtype)
+                or pa.types.is_fixed_size_list(ser.dtype.pyarrow_dtype)):
             elem_dtype = ser.dtype.pyarrow_dtype.value_type
             if pa.types.is_integer(elem_dtype):
                 return Dtype.intlist

kumoai/experimental/rfm/relbench.py

@@ -0,0 +1,76 @@
+import difflib
+import json
+from functools import lru_cache
+from urllib.request import urlopen
+
+import pooch
+import pyarrow as pa
+
+from kumoai.experimental.rfm import Graph
+from kumoai.experimental.rfm.backend.local import LocalTable
+
+PREFIX = 'rel-'
+CACHE_DIR = pooch.os_cache('relbench')
+HASH_URL = ('https://raw.githubusercontent.com/snap-stanford/relbench/main/'
+            'relbench/datasets/hashes.json')
+
+
+@lru_cache
+def get_registry() -> pooch.Pooch:
+    with urlopen(HASH_URL) as r:
+        hashes = json.load(r)
+
+    return pooch.create(
+        path=CACHE_DIR,
+        base_url='https://relbench.stanford.edu/download/',
+        registry=hashes,
+    )
+
+
+def from_relbench(dataset: str, verbose: bool = True) -> Graph:
+    dataset = dataset.lower()
+    if dataset.startswith(PREFIX):
+        dataset = dataset[len(PREFIX):]
+
+    registry = get_registry()
+
+    datasets = [key.split('/')[0][len(PREFIX):] for key in registry.registry]
+    if dataset not in datasets:
+        matches = difflib.get_close_matches(dataset, datasets, n=1)
+        hint = f" Did you mean '{matches[0]}'?" if len(matches) > 0 else ''
+        raise ValueError(f"Unknown RelBench dataset '{dataset}'.{hint} Valid "
+                         f"datasets are {str(datasets)[1:-1]}.")
+
+    registry.fetch(
+        f'{PREFIX}{dataset}/db.zip',
+        processor=pooch.Unzip(extract_dir='.'),
+        progressbar=verbose,
+    )
+
+    graph = Graph(tables=[])
+    edges: list[tuple[str, str, str]] = []
+    for path in (CACHE_DIR / f'{PREFIX}{dataset}' / 'db').glob('*.parquet'):
+        data = pa.parquet.read_table(path)
+        metadata = {
+            key.decode('utf-8'): json.loads(value.decode('utf-8'))
+            for key, value in data.schema.metadata.items()
+            if key in [b"fkey_col_to_pkey_table", b"pkey_col", b"time_col"]
+        }
+
+        table = LocalTable(
+            df=data.to_pandas(),
+            name=path.stem,
+            primary_key=metadata['pkey_col'],
+            time_column=metadata['time_col'],
+        )
+        graph.add_table(table)
+
+        edges.extend([
+            (path.stem, fkey, dst_table)
+            for fkey, dst_table in metadata['fkey_col_to_pkey_table'].items()
+        ])
+
+    for edge in edges:
+        graph.link(*edge)
+
+    return graph