kumoai 2.10.1__cp311-cp311-macosx_11_0_arm64.whl → 2.12.0.dev202511031731__cp311-cp311-macosx_11_0_arm64.whl

kumoai/__init__.py

@@ -200,9 +200,11 @@ def init(
 
     logger = logging.getLogger('kumoai')
     log_level = logging.getLevelName(logger.getEffectiveLevel())
+
     logger.info(
-        "Successfully initialized the Kumo SDK against deployment %s, with "
-        "log level %s.", url, log_level)
+        f"Successfully initialized the Kumo SDK (version {__version__}) "
+        f"against deployment {url}, with "
+        f"log level {log_level}.")
 
 
 def set_log_level(level: str) -> None:

kumoai/_version.py

@@ -1 +1 @@
-__version__ = '2.10.1'
+__version__ = '2.12.0.dev202511031731'

kumoai/experimental/rfm/local_graph.py

@@ -611,8 +611,8 @@ class LocalGraph:
                 raise ValueError(f"{edge} is invalid as foreign key "
                                  f"'{fkey}' and primary key '{dst_key.name}' "
                                  f"have incompatible data types (got "
-                                 f"fkey.dtype '{dst_key.dtype}' and "
-                                 f"pkey.dtype '{src_key.dtype}')")
+                                 f"fkey.dtype '{src_key.dtype}' and "
+                                 f"pkey.dtype '{dst_key.dtype}')")
 
         return self
 

kumoai/experimental/rfm/pquery/__init__.py

@@ -1,7 +1,11 @@
 from .backend import PQueryBackend
 from .pandas_backend import PQueryPandasBackend
+from .executor import PQueryExecutor
+from .pandas_executor import PQueryPandasExecutor
 
 __all__ = [
     'PQueryBackend',
     'PQueryPandasBackend',
+    'PQueryExecutor',
+    'PQueryPandasExecutor',
 ]

kumoai/experimental/rfm/pquery/executor.py

@@ -0,0 +1,102 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Generic, Tuple, TypeVar
+
+from kumoapi.pquery import ValidatedPredictiveQuery
+from kumoapi.pquery.AST import (
+    Aggregation,
+    Column,
+    Condition,
+    Filter,
+    Join,
+    LogicalOperation,
+)
+
+TableData = TypeVar('TableData')
+ColumnData = TypeVar('ColumnData')
+IndexData = TypeVar('IndexData')
+
+
+class PQueryExecutor(Generic[TableData, ColumnData, IndexData], ABC):
+    @abstractmethod
+    def execute_column(
+        self,
+        column: Column,
+        feat_dict: Dict[str, TableData],
+        filter_na: bool = True,
+    ) -> 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],
+        anchor_time: ColumnData,
+        filter_na: bool = True,
+        num_forecasts: int = 1,
+    ) -> 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],
+        anchor_time: ColumnData,
+        filter_na: bool = True,
+        num_forecasts: int = 1,
+    ) -> 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],
+        anchor_time: ColumnData,
+        filter_na: bool = True,
+        num_forecasts: int = 1,
+    ) -> 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],
+        anchor_time: ColumnData,
+        filter_na: bool = True,
+        num_forecasts: int = 1,
+    ) -> 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],
+        anchor_time: ColumnData,
+    ) -> 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],
+        anchor_time: ColumnData,
+        num_forecasts: int = 1,
+    ) -> Tuple[ColumnData, IndexData]:
+        pass

kumoai/experimental/rfm/pquery/pandas_executor.py

@@ -0,0 +1,506 @@
+from typing import Dict, List, Tuple
+
+import numpy as np
+import pandas as pd
+from kumoapi.pquery import ValidatedPredictiveQuery
+from kumoapi.pquery.AST import (
+    Aggregation,
+    Column,
+    Condition,
+    Constant,
+    Filter,
+    Join,
+    LogicalOperation,
+)
+from kumoapi.typing import AggregationType, BoolOp, MemberOp, RelOp
+
+from kumoai.experimental.rfm.pquery import PQueryExecutor
+
+
+class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
+                                          np.ndarray]):
+    def execute_column(
+        self,
+        column: Column,
+        feat_dict: Dict[str, pd.DataFrame],
+        filter_na: bool = True,
+    ) -> 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'))
+        else:
+            out = feat_dict[table_name][column_name]
+            out = out.reset_index(drop=True)
+
+        if pd.api.types.is_float_dtype(out):
+            out = out.astype('float32')
+
+        out.name = None
+        out.index.name = None
+
+        mask = out.notna().to_numpy()
+
+        if not filter_na:
+            return out, mask
+
+        out = out[mask].reset_index(drop=True)
+
+        # Cast to primitive dtype:
+        if pd.api.types.is_integer_dtype(out):
+            out = out.astype('int64')
+        elif pd.api.types.is_bool_dtype(out):
+            out = out.astype('bool')
+
+        return out, mask
+
+    def execute_aggregation_type(
+        self,
+        op: AggregationType,
+        feat: pd.Series,
+        batch: np.ndarray,
+        batch_size: int,
+        filter_na: bool = True,
+    ) -> Tuple[pd.Series, np.ndarray]:
+
+        mask = feat.notna()
+        feat, batch = feat[mask], batch[mask]
+
+        if op == AggregationType.LIST_DISTINCT:
+            df = pd.DataFrame(dict(feat=feat, batch=batch))
+            df = df.drop_duplicates()
+            out = df.groupby('batch')['feat'].agg(list)
+
+        else:
+            df = pd.DataFrame(dict(feat=feat, batch=batch))
+            if op == AggregationType.AVG:
+                agg = 'mean'
+            elif op == AggregationType.COUNT:
+                agg = 'size'
+            else:
+                agg = op.lower()
+            out = df.groupby('batch')['feat'].agg(agg)
+
+            if not pd.api.types.is_datetime64_any_dtype(out):
+                out = out.astype('float32')
+
+        out.name = None
+        out.index.name = None
+
+        if op in {AggregationType.SUM, AggregationType.COUNT}:
+            out = out.reindex(range(batch_size), fill_value=0)
+            mask = np.ones(batch_size, dtype=bool)
+            return out, mask
+
+        mask = np.zeros(batch_size, dtype=bool)
+        mask[batch] = True
+
+        if filter_na:
+            return out.reset_index(drop=True), mask
+
+        out = out.reindex(range(batch_size), fill_value=pd.NA)
+
+        return out, mask
+
+    def execute_aggregation(
+        self,
+        aggr: Aggregation,
+        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]:
+        target_table = aggr._get_target_column_name().split('.')[0]
+        target_batch = batch_dict[target_table]
+        target_time = time_dict[target_table]
+        if isinstance(aggr.target, Column):
+            target_feat, target_mask = self.execute_column(
+                column=aggr.target,
+                feat_dict=feat_dict,
+                filter_na=False,
+            )
+        else:
+            assert isinstance(aggr.target, Filter)
+            target_feat, target_mask = self.execute_filter(
+                filter=aggr.target,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=False,
+            )
+
+        outs: List[pd.Series] = []
+        masks: List[np.ndarray] = []
+        for _ in range(num_forecasts):
+            anchor_target_time = anchor_time[target_batch]
+            anchor_target_time = anchor_target_time.reset_index(drop=True)
+
+            curr_target_mask = target_mask & (
+                target_time
+                <= anchor_target_time + aggr.aggr_time_range.end_date_offset)
+            if aggr.aggr_time_range.start is not None:
+                start_offset = aggr.aggr_time_range.start_date_offset
+                curr_target_mask &= (target_time
+                                     > anchor_target_time + start_offset)
+            else:
+                assert num_forecasts == 1
+
+            out, mask = self.execute_aggregation_type(
+                aggr.aggr,
+                feat=target_feat[curr_target_mask],
+                batch=target_batch[curr_target_mask],
+                batch_size=len(anchor_time),
+                filter_na=False if num_forecasts > 1 else filter_na,
+            )
+            outs.append(out)
+            masks.append(mask)
+
+            if num_forecasts > 1:
+                anchor_time = (anchor_time +
+                               aggr.aggr_time_range.end_date_offset)
+        if len(outs) == 1:
+            assert len(masks) == 1
+            return outs[0], masks[0]
+
+        out = pd.Series([list(ser) for ser in zip(*outs)])
+        mask = np.stack(masks, axis=-1).any(axis=-1)  # type: ignore
+
+        if filter_na:
+            out = out[mask].reset_index(drop=True)
+
+        return out, mask
+
+    def execute_rel_op(
+        self,
+        left: pd.Series,
+        op: RelOp,
+        right: Constant,
+    ) -> pd.Series:
+
+        if right.typed_value() is None:
+            if op == RelOp.EQ:
+                return left.isna()
+            assert op == RelOp.NEQ
+            return left.notna()
+
+        # Promote left to float if right is a float to avoid lossy coercion.
+        right_value = right.typed_value()
+        if pd.api.types.is_integer_dtype(left) and isinstance(
+                right_value, float):
+            left = left.astype('float64')
+        value = pd.Series([right_value], dtype=left.dtype).iloc[0]
+
+        if op == RelOp.EQ:
+            return (left == value).fillna(False).astype(bool)
+        if op == RelOp.NEQ:
+            out = (left != value).fillna(False).astype(bool)
+            out[left.isna()] = False  # N/A != right should always be `False`.
+            return out
+        if op == RelOp.LEQ:
+            return (left <= value).fillna(False).astype(bool)
+        if op == RelOp.GEQ:
+            return (left >= value).fillna(False).astype(bool)
+        if op == RelOp.LT:
+            return (left < value).fillna(False).astype(bool)
+        if op == RelOp.GT:
+            return (left > value).fillna(False).astype(bool)
+
+        raise NotImplementedError(f"Operator '{op}' not implemented")
+
+    def execute_member_op(
+        self,
+        left: pd.Series,
+        op: MemberOp,
+        right: Constant,
+    ) -> pd.Series:
+
+        if op == MemberOp.IN:
+            ser = pd.Series(right.typed_value(), dtype=left.dtype)
+            return left.isin(ser).astype(bool)
+
+        raise NotImplementedError(f"Operator '{op}' not implemented")
+
+    def execute_condition(
+        self,
+        condition: Condition,
+        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]:
+        if num_forecasts > 1:
+            raise NotImplementedError("Forecasting not yet implemented for "
+                                      "non-regression tasks")
+
+        assert isinstance(condition.value, Constant)
+        value_is_na = condition.value.typed_value() is None
+        if isinstance(condition.target, Column):
+            left, mask = self.execute_column(
+                column=condition.target,
+                feat_dict=feat_dict,
+                filter_na=filter_na if not value_is_na else False,
+            )
+        elif isinstance(condition.target, Join):
+            left, mask = self.execute_join(
+                join=condition.target,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=filter_na if not value_is_na else False,
+            )
+        else:
+            assert isinstance(condition.target, Aggregation)
+            left, mask = self.execute_aggregation(
+                aggr=condition.target,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=filter_na if not value_is_na else False,
+            )
+
+        if filter_na and value_is_na:
+            mask = np.ones(len(left), dtype=bool)
+
+        if isinstance(condition.op, RelOp):
+            out = self.execute_rel_op(
+                left=left,
+                op=condition.op,
+                right=condition.value,
+            )
+        else:
+            assert isinstance(condition.op, MemberOp)
+            out = self.execute_member_op(
+                left=left,
+                op=condition.op,
+                right=condition.value,
+            )
+
+        return out, mask
+
+    def execute_bool_op(
+        self,
+        left: pd.Series,
+        op: BoolOp,
+        right: pd.Series | None,
+    ) -> pd.Series:
+
+        # TODO Implement Kleene-Priest three-value logic.
+        if op == BoolOp.AND:
+            assert right is not None
+            return left & right
+        if op == BoolOp.OR:
+            assert right is not None
+            return left | right
+        if op == BoolOp.NOT:
+            return ~left
+
+        raise NotImplementedError(f"Operator '{op}' not implemented")
+
+    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],
+        anchor_time: pd.Series,
+        filter_na: bool = True,
+        num_forecasts: int = 1,
+    ) -> Tuple[pd.Series, np.ndarray]:
+        if num_forecasts > 1:
+            raise NotImplementedError("Forecasting not yet implemented for "
+                                      "non-regression tasks")
+
+        if isinstance(logical_operation.left, Condition):
+            left, mask = self.execute_condition(
+                condition=logical_operation.left,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=False,
+            )
+        else:
+            assert isinstance(logical_operation.left, LogicalOperation)
+            left, mask = self.execute_logical_operation(
+                logical_operation=logical_operation.left,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=False,
+            )
+
+        right = right_mask = None
+        if isinstance(logical_operation.right, Condition):
+            right, right_mask = self.execute_condition(
+                condition=logical_operation.right,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=False,
+            )
+        elif isinstance(logical_operation.right, LogicalOperation):
+            right, right_mask = self.execute_logical_operation(
+                logical_operation=logical_operation.right,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=False,
+            )
+
+        out = self.execute_bool_op(left, logical_operation.bool_op, right)
+
+        if right_mask is not None:
+            mask &= right_mask
+
+        if filter_na:
+            out = out[mask].reset_index(drop=True)
+
+        return out, mask
+
+    def execute_join(
+        self,
+        join: Join,
+        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]:
+        if isinstance(join.rhs_target, Aggregation):
+            return self.execute_aggregation(
+                aggr=join.rhs_target,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=True,
+                num_forecasts=num_forecasts,
+            )
+        raise NotImplementedError(
+            f'Unexpected {type(join.rhs_target)} nested in Join')
+
+    def execute_filter(
+        self,
+        filter: Filter,
+        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]:
+        out, mask = self.execute_column(
+            column=filter.target,
+            feat_dict=feat_dict,
+            filter_na=False,
+        )
+        if isinstance(filter.condition, Condition):
+            _mask = self.execute_condition(
+                condition=filter.condition,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=False,
+            )[0].to_numpy()
+        else:
+            assert isinstance(filter.condition, LogicalOperation)
+            _mask = self.execute_logical_operation(
+                logical_operation=filter.condition,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=False,
+            )[0].to_numpy()
+        if filter_na:
+            return out[_mask & mask].reset_index(drop=True), _mask & mask
+        else:
+            return out[_mask].reset_index(drop=True), mask & _mask
+
+    def execute(
+        self,
+        query: ValidatedPredictiveQuery,
+        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]:
+        if isinstance(query.entity_ast, Column):
+            out, mask = self.execute_column(
+                column=query.entity_ast,
+                feat_dict=feat_dict,
+                filter_na=True,
+            )
+        else:
+            assert isinstance(query.entity_ast, Filter)
+            out, mask = self.execute_filter(
+                filter=query.entity_ast,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+            )
+        if isinstance(query.target_ast, Column):
+            out, _mask = self.execute_column(
+                column=query.target_ast,
+                feat_dict=feat_dict,
+                filter_na=True,
+            )
+        elif isinstance(query.target_ast, Condition):
+            out, _mask = self.execute_condition(
+                condition=query.target_ast,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=True,
+                num_forecasts=num_forecasts,
+            )
+        elif isinstance(query.target_ast, Aggregation):
+            out, _mask = self.execute_aggregation(
+                aggr=query.target_ast,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=True,
+                num_forecasts=num_forecasts,
+            )
+        elif isinstance(query.target_ast, Join):
+            out, _mask = self.execute_join(
+                join=query.target_ast,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=True,
+                num_forecasts=num_forecasts,
+            )
+        elif isinstance(query.target_ast, LogicalOperation):
+            out, _mask = self.execute_logical_operation(
+                logical_operation=query.target_ast,
+                feat_dict=feat_dict,
+                time_dict=time_dict,
+                batch_dict=batch_dict,
+                anchor_time=anchor_time,
+                filter_na=True,
+                num_forecasts=num_forecasts,
+            )
+        else:
+            raise NotImplementedError(
+                f'{type(query.target)} compilation missing.')
+        out = out[mask[_mask]]
+        mask &= _mask
+        out = out.reset_index(drop=True)
+        return out, mask

kumoai/experimental/rfm/rfm.py

@@ -199,6 +199,7 @@ class KumoRFM:
         max_pq_iterations: int = 20,
         random_seed: Optional[int] = _RANDOM_SEED,
         verbose: Union[bool, ProgressLogger] = True,
+        use_prediction_time: bool = False,
     ) -> pd.DataFrame:
         pass
 
@@ -217,6 +218,7 @@ class KumoRFM:
         max_pq_iterations: int = 20,
         random_seed: Optional[int] = _RANDOM_SEED,
         verbose: Union[bool, ProgressLogger] = True,
+        use_prediction_time: bool = False,
     ) -> Explanation:
         pass
 
@@ -234,6 +236,7 @@ class KumoRFM:
         max_pq_iterations: int = 20,
         random_seed: Optional[int] = _RANDOM_SEED,
         verbose: Union[bool, ProgressLogger] = True,
+        use_prediction_time: bool = False,
     ) -> Union[pd.DataFrame, Explanation]:
         """Returns predictions for a predictive query.
 
@@ -264,6 +267,9 @@ class KumoRFM:
                 entities to find valid labels.
             random_seed: A manual seed for generating pseudo-random numbers.
             verbose: Whether to print verbose output.
+            use_prediction_time: Whether to use the anchor timestamp as an
+                additional feature during prediction. This is typically
+                beneficial for time series forecasting tasks.
 
         Returns:
             The predictions as a :class:`pandas.DataFrame`.
@@ -353,6 +359,7 @@ class KumoRFM:
                 request = RFMPredictRequest(
                     context=context,
                     run_mode=RunMode(run_mode),
+                    use_prediction_time=use_prediction_time,
                 )
                 with warnings.catch_warnings():
                     warnings.filterwarnings('ignore', message='gencode')
@@ -503,6 +510,7 @@ class KumoRFM:
         max_pq_iterations: int = 20,
         random_seed: Optional[int] = _RANDOM_SEED,
         verbose: Union[bool, ProgressLogger] = True,
+        use_prediction_time: bool = False,
     ) -> pd.DataFrame:
         """Evaluates a predictive query.
 
@@ -526,6 +534,9 @@ class KumoRFM:
                 entities to find valid labels.
             random_seed: A manual seed for generating pseudo-random numbers.
             verbose: Whether to print verbose output.
+            use_prediction_time: Whether to use the anchor timestamp as an
+                additional feature during prediction. This is typically
+                beneficial for time series forecasting tasks.
 
         Returns:
             The metrics as a :class:`pandas.DataFrame`
@@ -569,6 +580,7 @@ class KumoRFM:
                 context=context,
                 run_mode=RunMode(run_mode),
                 metrics=metrics,
+                use_prediction_time=use_prediction_time,
             )
             with warnings.catch_warnings():
                 warnings.filterwarnings('ignore', message='Protobuf gencode')
@@ -1006,7 +1018,7 @@ class KumoRFM:
         elif task_type == TaskType.MULTICLASS_CLASSIFICATION:
             supported_metrics = ['acc', 'precision', 'recall', 'f1', 'mrr']
         elif task_type == TaskType.REGRESSION:
-            supported_metrics = ['mae', 'mape', 'mse', 'rmse', 'smape']
+            supported_metrics = ['mae', 'mape', 'mse', 'rmse', 'smape', 'r2']
         elif task_type == TaskType.TEMPORAL_LINK_PREDICTION:
             supported_metrics = [
                 'map@', 'ndcg@', 'mrr@', 'precision@', 'recall@', 'f1@',

kumoai/trainer/trainer.py

@@ -20,7 +20,6 @@ from kumoapi.jobs import (
     TrainingJobResource,
 )
 from kumoapi.model_plan import ModelPlan
-from kumoapi.task import TaskType
 
 from kumoai import global_state
 from kumoai.artifact_export.config import OutputConfig
@@ -405,15 +404,15 @@ class Trainer:
                 pred_table_data_path = prediction_table.table_data_uri
 
         api = global_state.client.batch_prediction_job_api
-
-        from kumoai.pquery.predictive_query import PredictiveQuery
-        pquery = PredictiveQuery.load_from_training_job(training_job_id)
-        if pquery.get_task_type() == TaskType.BINARY_CLASSIFICATION:
-            if binary_classification_threshold is None:
-                logger.warning("No binary classification threshold provided. "
-                               "Using default threshold of 0.5.")
-                binary_classification_threshold = 0.5
-
+        # Remove to resolve https://github.com/kumo-ai/kumo/issues/24250
+        # from kumoai.pquery.predictive_query import PredictiveQuery
+        # pquery = PredictiveQuery.load_from_training_job(training_job_id)
+        # if pquery.get_task_type() == TaskType.BINARY_CLASSIFICATION:
+        #     if binary_classification_threshold is None:
+        #         logger.warning(
+        # "No binary classification threshold provided. "
+        # "Using default threshold of 0.5.")
+        #         binary_classification_threshold = 0.5
         job_id, response = api.maybe_create(
             BatchPredictionRequest(
                 dict(custom_tags),

{kumoai-2.10.1.dist-info → kumoai-2.12.0.dev202511031731.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kumoai
-Version: 2.10.1
+Version: 2.12.0.dev202511031731
 Summary: AI on the Modern Data Stack
 Author-email: "Kumo.AI" <hello@kumo.ai>
 License-Expression: MIT
@@ -23,7 +23,7 @@ Requires-Dist: requests>=2.28.2
 Requires-Dist: urllib3
 Requires-Dist: plotly
 Requires-Dist: typing_extensions>=4.5.0
-Requires-Dist: kumo-api==0.38.0
+Requires-Dist: kumo-api==0.40.0
 Requires-Dist: tqdm>=4.66.0
 Requires-Dist: aiohttp>=3.10.0
 Requires-Dist: pydantic>=1.10.21

{kumoai-2.10.1.dist-info → kumoai-2.12.0.dev202511031731.dist-info}/RECORD

@@ -1,7 +1,7 @@
 kumoai/_logging.py,sha256=U2_5ROdyk92P4xO4H2WJV8EC7dr6YxmmnM-b7QX9M7I,886
 kumoai/mixin.py,sha256=MP413xzuCqWhxAPUHmloLA3j4ZyF1tEtfi516b_hOXQ,812
-kumoai/_version.py,sha256=fy1qvJHPd7FWOAuuVNKh9cFE7RrqWon8x59x3EjYTCc,23
-kumoai/__init__.py,sha256=x3DjDsWBgWSNwo7mDwb3XAoRm2NuSO09yvhQTL9tBT8,10673
+kumoai/_version.py,sha256=X5C9cHVsjznMq0N29k8V18IjmrXq8NyKWG7IEMkjaBc,39
+kumoai/__init__.py,sha256=LU1zmKYc0KV5hy2VGKUuXgSvbJwj2rSRQ_R_bpHyl1o,10708
 kumoai/formatting.py,sha256=jA_rLDCGKZI8WWCha-vtuLenVKTZvli99Tqpurz1H84,953
 kumoai/futures.py,sha256=oJFIfdCM_3nWIqQteBKYMY4fPhoYlYWE_JA2o6tx-ng,3737
 kumoai/kumolib.cpython-311-darwin.so,sha256=AmB_Fysmud1y7Gm5CuBQ5lWDuSzpxVDV_iTA2cjH1s8,232544
@@ -12,17 +12,19 @@ kumoai/spcs.py,sha256=N4ddeoHAc4I3bKrDitsb91lUx5VKvCyPyMT3zWiuCcY,4275
 kumoai/_singleton.py,sha256=UTwrbDkoZSGB8ZelorvprPDDv9uZkUi1q_SrmsyngpQ,836
 kumoai/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kumoai/experimental/rfm/local_graph_sampler.py,sha256=o60_sdMa_fr60DrdmCIaE6lKQAD2msp1t-GGubFNt-o,6738
-kumoai/experimental/rfm/local_graph.py,sha256=2LTllKKnjkThM7lr6jg_miypWN_oLC3YmIcZHLkAa4U,30076
+kumoai/experimental/rfm/local_graph.py,sha256=2iJDlsGVzqCe1bD_puXWlhwGkn7YnQyJ4p4C-fwCZNE,30076
 kumoai/experimental/rfm/local_pquery_driver.py,sha256=xqAQ9fJfkqM1axknFpg0NLQbIYmExh-s7vGdUyDEkwA,18600
 kumoai/experimental/rfm/__init__.py,sha256=F1aUOCLDN2yrIRDAiOlogDfXKUkUQgp8Mt0pVX9rLX8,1641
 kumoai/experimental/rfm/utils.py,sha256=3IiBvT_aLBkkcJh3H11_50yt_XlEzHR0cm9Kprrtl8k,11123
 kumoai/experimental/rfm/local_table.py,sha256=r8xZ33Mjs6JD8ud6h23tZ99Dag2DvZ4h6tWjmGrKQg4,19605
-kumoai/experimental/rfm/rfm.py,sha256=cKHwIinGmnRg1_QiDpzhC_ZnNTWhvX3tA4ZhXgjlNEU,45513
+kumoai/experimental/rfm/rfm.py,sha256=BcC0EqXfz2OhMT-g8gBGv7M6yUTboj-PyGWIQZPUf70,46227
 kumoai/experimental/rfm/local_graph_store.py,sha256=8BqonuaMftAAsjgZpB369i5AeNd1PkisMbbEqc0cKBo,13847
 kumoai/experimental/rfm/authenticate.py,sha256=FiuHMvP7V3zBZUlHMDMbNLhc-UgDZgz4hjVSTuQ7DRw,18888
 kumoai/experimental/rfm/pquery/backend.py,sha256=6wtB0yFpxQUraBSA2TbKMVSIMD0dcLwYV5P4SQx2g_k,3287
-kumoai/experimental/rfm/pquery/__init__.py,sha256=bsNcdn7DnPw9kpSQ_bQVmQX1RmXzPQhzfA1y6G-n7I8,146
+kumoai/experimental/rfm/pquery/__init__.py,sha256=9uLXixjp78y0IzO2F__lFqKNm37OGhN3iDh56akWLNU,283
 kumoai/experimental/rfm/pquery/pandas_backend.py,sha256=pgHCErSo6U-KJMhgIYijYt96uubtFB2WtsrTdLU7NYc,15396
+kumoai/experimental/rfm/pquery/pandas_executor.py,sha256=BgF3saosisgLHx1RyLj-HSEbMp4xLatNuARdKWwiiLY,17326
+kumoai/experimental/rfm/pquery/executor.py,sha256=f7-pJhL0BgFU9E4o4gQpQyArOvyrZtwxFmks34-QOAE,2741
 kumoai/experimental/rfm/infer/multicategorical.py,sha256=0-cLpDnGryhr76QhZNO-klKokJ6MUSfxXcGdQ61oykY,1102
 kumoai/experimental/rfm/infer/categorical.py,sha256=VwNaKwKbRYkTxEJ1R6gziffC8dGsEThcDEfbi-KqW5c,853
 kumoai/experimental/rfm/infer/id.py,sha256=ZIO0DWIoiEoS_8MVc5lkqBfkTWWQ0yGCgjkwLdaYa_Q,908
@@ -90,9 +92,9 @@ kumoai/trainer/job.py,sha256=Wk69nzFhbvuA3nEvtCstI04z5CxkgvQ6tHnGchE0Lkg,44938
 kumoai/trainer/baseline_trainer.py,sha256=LlfViNOmswNv4c6zJJLsyv0pC2mM2WKMGYx06ogtEVc,4024
 kumoai/trainer/__init__.py,sha256=zUdFl-f-sBWmm2x8R-rdVzPBeU2FaMzUY5mkcgoTa1k,939
 kumoai/trainer/online_serving.py,sha256=9cddb5paeZaCgbUeceQdAOxysCtV5XP-KcsgFz_XR5w,9566
-kumoai/trainer/trainer.py,sha256=nPeZMMp17TtRFd4lKbF-TlMPnhYR4_VyPDPI0T9W9PU,20094
-kumoai-2.10.1.dist-info/RECORD,,
-kumoai-2.10.1.dist-info/WHEEL,sha256=sunMa2yiYbrNLGeMVDqEA0ayyJbHlex7SCn1TZrEq60,136
-kumoai-2.10.1.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
-kumoai-2.10.1.dist-info/METADATA,sha256=6XcvFFVYccShSlMdWchiyehRoiG93v2gU18n7D6pwD4,2036
-kumoai-2.10.1.dist-info/licenses/LICENSE,sha256=TbWlyqRmhq9PEzCaTI0H0nWLQCCOywQM8wYH8MbjfLo,1102
+kumoai/trainer/trainer.py,sha256=hBXO7gwpo3t59zKFTeIkK65B8QRmWCwO33sbDuEAPlY,20133
+kumoai-2.12.0.dev202511031731.dist-info/RECORD,,
+kumoai-2.12.0.dev202511031731.dist-info/WHEEL,sha256=sunMa2yiYbrNLGeMVDqEA0ayyJbHlex7SCn1TZrEq60,136
+kumoai-2.12.0.dev202511031731.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
+kumoai-2.12.0.dev202511031731.dist-info/METADATA,sha256=yf8LuBryiRverUZLTN389Y_94sZrrzNnITX6sDAlfy0,2052
+kumoai-2.12.0.dev202511031731.dist-info/licenses/LICENSE,sha256=TbWlyqRmhq9PEzCaTI0H0nWLQCCOywQM8wYH8MbjfLo,1102