kumoai 2.13.0.dev202512031731__cp312-cp312-macosx_11_0_arm64.whl → 2.14.0.dev202512301731__cp312-cp312-macosx_11_0_arm64.whl

kumoai/experimental/rfm/infer/__init__.py

@@ -1,17 +1,19 @@
 from .dtype import infer_dtype
-from .pkey import infer_primary_key
-from .time_col import infer_time_column
 from .id import contains_id
 from .timestamp import contains_timestamp
 from .categorical import contains_categorical
 from .multicategorical import contains_multicategorical
+from .stype import infer_stype
+from .pkey import infer_primary_key
+from .time_col import infer_time_column
 
 __all__ = [
     'infer_dtype',
-    'infer_primary_key',
-    'infer_time_column',
     'contains_id',
     'contains_timestamp',
     'contains_categorical',
     'contains_multicategorical',
+    'infer_stype',
+    'infer_primary_key',
+    'infer_time_column',
 ]

kumoai/experimental/rfm/infer/dtype.py

@@ -1,36 +1,27 @@
-from typing import Any, Dict
-
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 from kumoapi.typing import Dtype
 
-PANDAS_TO_DTYPE: Dict[Any, Dtype] = {
-    np.dtype('bool'): Dtype.bool,
-    pd.BooleanDtype(): Dtype.bool,
-    pa.bool_(): Dtype.bool,
-    np.dtype('byte'): Dtype.int,
-    pd.UInt8Dtype(): Dtype.int,
-    np.dtype('int16'): Dtype.int,
-    pd.Int16Dtype(): Dtype.int,
-    np.dtype('int32'): Dtype.int,
-    pd.Int32Dtype(): Dtype.int,
-    np.dtype('int64'): Dtype.int,
-    pd.Int64Dtype(): Dtype.int,
-    np.dtype('float32'): Dtype.float,
-    pd.Float32Dtype(): Dtype.float,
-    np.dtype('float64'): Dtype.float,
-    pd.Float64Dtype(): Dtype.float,
-    np.dtype('object'): Dtype.string,
-    pd.StringDtype(storage='python'): Dtype.string,
-    pd.StringDtype(storage='pyarrow'): Dtype.string,
-    pa.string(): Dtype.string,
-    pa.binary(): Dtype.binary,
-    np.dtype('datetime64[ns]'): Dtype.date,
-    np.dtype('timedelta64[ns]'): Dtype.timedelta,
-    pa.list_(pa.float32()): Dtype.floatlist,
-    pa.list_(pa.int64()): Dtype.intlist,
-    pa.list_(pa.string()): Dtype.stringlist,
+PANDAS_TO_DTYPE: dict[str, Dtype] = {
+    'bool': Dtype.bool,
+    'boolean': Dtype.bool,
+    'int8': Dtype.int,
+    'int16': Dtype.int,
+    'int32': Dtype.int,
+    'int64': Dtype.int,
+    'float': Dtype.float,
+    'double': Dtype.float,
+    'float16': Dtype.float,
+    'float32': Dtype.float,
+    'float64': Dtype.float,
+    'object': Dtype.string,
+    'string': Dtype.string,
+    'string[python]': Dtype.string,
+    'string[pyarrow]': Dtype.string,
+    'binary': Dtype.binary,
+    'binary[python]': Dtype.binary,
+    'binary[pyarrow]': Dtype.binary,
 }
 
 
@@ -45,46 +36,47 @@ def infer_dtype(ser: pd.Series) -> Dtype:
     """
     if pd.api.types.is_datetime64_any_dtype(ser.dtype):
         return Dtype.date
-
+    if pd.api.types.is_timedelta64_dtype(ser.dtype):
+        return Dtype.timedelta
     if isinstance(ser.dtype, pd.CategoricalDtype):
         return Dtype.string
 
-    if pd.api.types.is_object_dtype(ser.dtype):
+    if (pd.api.types.is_object_dtype(ser.dtype)
+            and not isinstance(ser.dtype, pd.ArrowDtype)):
         index = ser.iloc[:1000].first_valid_index()
         if index is not None and pd.api.types.is_list_like(ser[index]):
             pos = ser.index.get_loc(index)
             assert isinstance(pos, int)
             ser = ser.iloc[pos:pos + 1000].dropna()
-
-            if not ser.map(pd.api.types.is_list_like).all():
-                raise ValueError("Data contains a mix of list-like and "
-                                 "non-list-like values")
-
-            # Remove all empty Python lists without known data type:
-            ser = ser[ser.map(lambda x: not isinstance(x, list) or len(x) > 0)]
-
-            # Infer unique data types in this series:
-            dtypes = ser.apply(lambda x: PANDAS_TO_DTYPE.get(
-                np.array(x).dtype, Dtype.string)).unique().tolist()
-
-            invalid_dtypes = set(dtypes) - {
-                Dtype.string,
-                Dtype.int,
-                Dtype.float,
-            }
-            if len(invalid_dtypes) > 0:
-                raise ValueError(f"Data contains unsupported list data types: "
-                                 f"{list(invalid_dtypes)}")
-
-            if Dtype.string in dtypes:
-                return Dtype.stringlist
-
-            if dtypes == [Dtype.int]:
+            arr = pa.array(ser.tolist())
+            ser = pd.Series(arr, dtype=pd.ArrowDtype(arr.type))
+
+    if isinstance(ser.dtype, pd.ArrowDtype):
+        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
-
-            return Dtype.floatlist
-
-    if ser.dtype not in PANDAS_TO_DTYPE:
+            if pa.types.is_floating(elem_dtype):
+                return Dtype.floatlist
+            if pa.types.is_decimal(elem_dtype):
+                return Dtype.floatlist
+            if pa.types.is_string(elem_dtype):
+                return Dtype.stringlist
+            if pa.types.is_null(elem_dtype):
+                return Dtype.floatlist
+
+    if isinstance(ser.dtype, np.dtype):
+        dtype_str = str(ser.dtype).lower()
+    elif isinstance(ser.dtype, pd.api.extensions.ExtensionDtype):
+        dtype_str = ser.dtype.name.lower()
+        dtype_str = dtype_str.split('[')[0]  # Remove backend metadata
+    elif isinstance(ser.dtype, pa.DataType):
+        dtype_str = str(ser.dtype).lower()
+    else:
+        dtype_str = 'object'
+
+    if dtype_str not in PANDAS_TO_DTYPE:
         raise ValueError(f"Unsupported data type '{ser.dtype}'")
 
-    return PANDAS_TO_DTYPE[ser.dtype]
+    return PANDAS_TO_DTYPE[dtype_str]

kumoai/experimental/rfm/infer/multicategorical.py

@@ -40,7 +40,7 @@ def contains_multicategorical(
         sep = max(candidates, key=candidates.get)  # type: ignore
         ser = ser.str.split(sep)
 
-    num_unique_multi = ser.explode().nunique()
+    num_unique_multi = ser.astype('object').explode().nunique()
 
     if dtype.is_list():
         return num_unique_multi <= MAX_CAT

kumoai/experimental/rfm/infer/pkey.py

@@ -1,6 +1,5 @@
 import re
 import warnings
-from typing import Optional
 
 import pandas as pd
 
@@ -9,7 +8,7 @@ def infer_primary_key(
     table_name: str,
     df: pd.DataFrame,
     candidates: list[str],
-) -> Optional[str]:
+) -> str | None:
     r"""Auto-detect potential primary key column.
 
     Args:
@@ -20,6 +19,9 @@ def infer_primary_key(
     Returns:
         The name of the detected primary key, or ``None`` if not found.
     """
+    if len(candidates) == 0:
+        return None
+
     # A list of (potentially modified) table names that are eligible to match
     # with a primary key, i.e.:
     # - UserInfo -> User

kumoai/experimental/rfm/infer/stype.py

@@ -0,0 +1,35 @@
+import pandas as pd
+from kumoapi.typing import Dtype, Stype
+
+from kumoai.experimental.rfm.infer import (
+    contains_categorical,
+    contains_id,
+    contains_multicategorical,
+    contains_timestamp,
+)
+
+
+def infer_stype(ser: pd.Series, column_name: str, dtype: Dtype) -> Stype:
+    """Infers the :class:`Stype` from a :class:`pandas.Series`.
+
+    Args:
+        ser: A :class:`pandas.Series` to analyze.
+        column_name: The column name.
+        dtype: The data type.
+
+    Returns:
+        The semantic type.
+    """
+    if contains_id(ser, column_name, dtype):
+        return Stype.ID
+
+    if contains_timestamp(ser, column_name, dtype):
+        return Stype.timestamp
+
+    if contains_multicategorical(ser, column_name, dtype):
+        return Stype.multicategorical
+
+    if contains_categorical(ser, column_name, dtype):
+        return Stype.categorical
+
+    return dtype.default_stype

kumoai/experimental/rfm/infer/time_col.py

@@ -1,6 +1,5 @@
 import re
 import warnings
-from typing import Optional
 
 import pandas as pd
 
@@ -8,7 +7,7 @@ import pandas as pd
 def infer_time_column(
     df: pd.DataFrame,
     candidates: list[str],
-) -> Optional[str]:
+) -> str | None:
     r"""Auto-detect potential time column.
 
     Args:

kumoai/experimental/rfm/pquery/executor.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Generic, Tuple, TypeVar
+from typing import Generic, TypeVar
 
 from kumoapi.pquery import ValidatedPredictiveQuery
 from kumoapi.pquery.AST import (
@@ -21,82 +21,82 @@ class PQueryExecutor(Generic[TableData, ColumnData, IndexData], ABC):
     def execute_column(
         self,
         column: Column,
-        feat_dict: Dict[str, TableData],
+        feat_dict: dict[str, TableData],
         filter_na: bool = True,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_aggregation(
         self,
         aggr: Aggregation,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_condition(
         self,
         condition: Condition,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_logical_operation(
         self,
         logical_operation: LogicalOperation,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_join(
         self,
         join: Join,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute_filter(
         self,
         filter: Filter,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass
 
     @abstractmethod
     def execute(
         self,
         query: ValidatedPredictiveQuery,
-        feat_dict: Dict[str, TableData],
-        time_dict: Dict[str, ColumnData],
-        batch_dict: Dict[str, IndexData],
+        feat_dict: dict[str, TableData],
+        time_dict: dict[str, ColumnData],
+        batch_dict: dict[str, IndexData],
         anchor_time: ColumnData,
         num_forecasts: int = 1,
-    ) -> Tuple[ColumnData, IndexData]:
+    ) -> tuple[ColumnData, IndexData]:
         pass

kumoai/experimental/rfm/pquery/pandas_executor.py

@@ -1,5 +1,3 @@
-from typing import Dict, List, Tuple
-
 import numpy as np
 import pandas as pd
 from kumoapi.pquery import ValidatedPredictiveQuery
@@ -22,9 +20,9 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_column(
         self,
         column: Column,
-        feat_dict: Dict[str, pd.DataFrame],
+        feat_dict: dict[str, pd.DataFrame],
         filter_na: bool = True,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         table_name, column_name = column.fqn.split(".")
         if column_name == '*':
             out = pd.Series(np.ones(len(feat_dict[table_name]), dtype='int64'))
@@ -60,7 +58,7 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
         batch: np.ndarray,
         batch_size: int,
         filter_na: bool = True,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
 
         mask = feat.notna()
         feat, batch = feat[mask], batch[mask]
@@ -104,13 +102,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_aggregation(
         self,
         aggr: Aggregation,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         target_table = aggr._get_target_column_name().split('.')[0]
         target_batch = batch_dict[target_table]
         target_time = time_dict[target_table]
@@ -131,10 +129,10 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
                 filter_na=True,
             )
 
-        outs: List[pd.Series] = []
-        masks: List[np.ndarray] = []
+        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 +
@@ -226,13 +224,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_condition(
         self,
         condition: Condition,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if num_forecasts > 1:
             raise NotImplementedError("Forecasting not yet implemented for "
                                       "non-regression tasks")
@@ -306,13 +304,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_logical_operation(
         self,
         logical_operation: LogicalOperation,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if num_forecasts > 1:
             raise NotImplementedError("Forecasting not yet implemented for "
                                       "non-regression tasks")
@@ -370,13 +368,13 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_join(
         self,
         join: Join,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if isinstance(join.rhs_target, Aggregation):
             return self.execute_aggregation(
                 aggr=join.rhs_target,
@@ -393,12 +391,12 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute_filter(
         self,
         filter: Filter,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         filter_na: bool = True,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         out, mask = self.execute_column(
             column=filter.target,
             feat_dict=feat_dict,
@@ -431,12 +429,12 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
     def execute(
         self,
         query: ValidatedPredictiveQuery,
-        feat_dict: Dict[str, pd.DataFrame],
-        time_dict: Dict[str, pd.Series],
-        batch_dict: Dict[str, np.ndarray],
+        feat_dict: dict[str, pd.DataFrame],
+        time_dict: dict[str, pd.Series],
+        batch_dict: dict[str, np.ndarray],
         anchor_time: pd.Series,
         num_forecasts: int = 1,
-    ) -> Tuple[pd.Series, np.ndarray]:
+    ) -> tuple[pd.Series, np.ndarray]:
         if isinstance(query.entity_ast, Column):
             out, mask = self.execute_column(
                 column=query.entity_ast,

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