kumoai 2.13.0.dev202512040649__cp313-cp313-win_amd64.whl → 2.14.0.dev202512211732__cp313-cp313-win_amd64.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,17 +1,17 @@
-from typing import Dict
-
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 from kumoapi.typing import Dtype
 
-PANDAS_TO_DTYPE: Dict[str, Dtype] = {
+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,

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,