kumoai 2.10.0.dev202509231831__cp313-cp313-macosx_11_0_arm64.whl → 2.14.0.dev202512161731__cp313-cp313-macosx_11_0_arm64.whl

kumoai/experimental/rfm/pquery/{pandas_backend.py → pandas_executor.py}

@@ -1,45 +1,69 @@
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Tuple
 
 import numpy as np
 import pandas as pd
-from kumoapi.rfm import PQueryDefinition
-from kumoapi.rfm.pquery import (
+from kumoapi.pquery import ValidatedPredictiveQuery
+from kumoapi.pquery.AST import (
     Aggregation,
-    AggregationType,
-    BoolOp,
     Column,
     Condition,
+    Constant,
     Filter,
-    Float,
-    FloatList,
-    Int,
-    IntList,
+    Join,
     LogicalOperation,
-    MemberOp,
-    RelOp,
-    Str,
-    StrList,
 )
+from kumoapi.typing import AggregationType, BoolOp, MemberOp, RelOp
 
-from kumoai.experimental.rfm.pquery import PQueryBackend
+from kumoai.experimental.rfm.pquery import PQueryExecutor
 
 
-class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
-    def eval_aggregation_type(
+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: Optional[pd.Series],
+        feat: pd.Series,
         batch: np.ndarray,
         batch_size: int,
         filter_na: bool = True,
     ) -> Tuple[pd.Series, np.ndarray]:
 
-        if op != AggregationType.COUNT:
-            assert feat is not None
-
-        if feat is not None:
-            mask = feat.notna()
-            feat, batch = feat[mask], batch[mask]
+        mask = feat.notna()
+        feat, batch = feat[mask], batch[mask]
 
         if op == AggregationType.LIST_DISTINCT:
             df = pd.DataFrame(dict(feat=feat, batch=batch))
@@ -77,102 +101,7 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
 
         return out, mask
 
-    def eval_rel_op(
-        self,
-        left: pd.Series,
-        op: RelOp,
-        right: Union[Int, Float, Str, None],
-    ) -> pd.Series:
-
-        if right is None:
-            if op == RelOp.EQ:
-                return left.isna()
-            assert op == RelOp.NEQ
-            return left.notna()
-
-        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 eval_member_op(
-        self,
-        left: pd.Series,
-        op: MemberOp,
-        right: Union[IntList, FloatList, StrList],
-    ) -> pd.Series:
-
-        if op == MemberOp.IN:
-            ser = pd.Series(right.value, dtype=left.dtype)
-            return left.isin(ser).astype(bool)
-
-        raise NotImplementedError(f"Operator '{op}' not implemented")
-
-    def eval_bool_op(
-        self,
-        left: pd.Series,
-        op: BoolOp,
-        right: Optional[pd.Series],
-    ) -> 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 eval_column(
-        self,
-        column: Column,
-        feat_dict: Dict[str, pd.DataFrame],
-        filter_na: bool = True,
-    ) -> Tuple[pd.Series, np.ndarray]:
-
-        out = feat_dict[column.table_name][column.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 eval_aggregation(
+    def execute_aggregation(
         self,
         aggr: Aggregation,
         feat_dict: Dict[str, pd.DataFrame],
@@ -182,49 +111,47 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
         filter_na: bool = True,
         num_forecasts: int = 1,
     ) -> Tuple[pd.Series, np.ndarray]:
-
-        target_table = aggr.column.table_name
+        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=True,
+            )
+        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=True,
+            )
 
         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)
 
-            target_mask = target_time <= anchor_target_time + aggr.end_offset
-
-            if aggr.start is not None:
-                start_offset = aggr.start * aggr.time_unit.to_offset()
-                target_mask &= target_time > anchor_target_time + start_offset
+            time_filter_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
+                time_filter_mask &= (target_time
+                                     > anchor_target_time + start_offset)
             else:
                 assert num_forecasts == 1
+            curr_target_mask = target_mask & time_filter_mask
 
-            if aggr.filter is not None:
-                target_mask &= self.eval_filter(
-                    filter=aggr.filter,
-                    feat_dict=feat_dict,
-                    time_dict=time_dict,
-                    batch_dict=batch_dict,
-                    anchor_time=anchor_time,
-                )
-
-            if (aggr.type == AggregationType.COUNT
-                    and aggr.column.column_name == '*'):
-                target_feat = None
-            else:
-                target_feat, _ = self.eval_column(
-                    aggr.column,
-                    feat_dict,
-                    filter_na=False,
-                )
-                target_feat = target_feat[target_mask]
-
-            out, mask = self.eval_aggregation_type(
-                aggr.type,
-                feat=target_feat,
-                batch=target_batch[target_mask],
+            out, mask = self.execute_aggregation_type(
+                aggr.aggr,
+                feat=target_feat[time_filter_mask[target_mask].reset_index(
+                    drop=True)],
+                batch=target_batch[curr_target_mask],
                 batch_size=len(anchor_time),
                 filter_na=False if num_forecasts > 1 else filter_na,
             )
@@ -232,8 +159,8 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
             masks.append(mask)
 
             if num_forecasts > 1:
-                anchor_time = anchor_time + aggr.end_offset
-
+                anchor_time = (anchor_time +
+                               aggr.aggr_time_range.end_date_offset)
         if len(outs) == 1:
             assert len(masks) == 1
             return outs[0], masks[0]
@@ -246,7 +173,57 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
 
         return out, mask
 
-    def eval_condition(
+    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],
@@ -256,48 +233,77 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
         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(condition.left, Column):
-            left, mask = self.eval_column(
-                column=condition.left,
+        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 condition.right is not None else False,
+                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.left, Aggregation)
-            left, mask = self.eval_aggregation(
-                aggr=condition.left,
+            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 condition.right is not None else False,
+                filter_na=filter_na if not value_is_na else False,
             )
 
-        if filter_na and condition.right is None:
+        if filter_na and value_is_na:
             mask = np.ones(len(left), dtype=bool)
 
         if isinstance(condition.op, RelOp):
-            out = self.eval_rel_op(
+            out = self.execute_rel_op(
                 left=left,
                 op=condition.op,
-                right=condition.right,
+                right=condition.value,
             )
         else:
             assert isinstance(condition.op, MemberOp)
-            out = self.eval_member_op(
+            out = self.execute_member_op(
                 left=left,
                 op=condition.op,
-                right=condition.right,
+                right=condition.value,
             )
 
         return out, mask
 
-    def eval_logical_operation(
+    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],
@@ -307,13 +313,12 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
         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.eval_condition(
+            left, mask = self.execute_condition(
                 condition=logical_operation.left,
                 feat_dict=feat_dict,
                 time_dict=time_dict,
@@ -323,7 +328,7 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
             )
         else:
             assert isinstance(logical_operation.left, LogicalOperation)
-            left, mask = self.eval_logical_operation(
+            left, mask = self.execute_logical_operation(
                 logical_operation=logical_operation.left,
                 feat_dict=feat_dict,
                 time_dict=time_dict,
@@ -334,7 +339,7 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
 
         right = right_mask = None
         if isinstance(logical_operation.right, Condition):
-            right, right_mask = self.eval_condition(
+            right, right_mask = self.execute_condition(
                 condition=logical_operation.right,
                 feat_dict=feat_dict,
                 time_dict=time_dict,
@@ -343,7 +348,7 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
                 filter_na=False,
             )
         elif isinstance(logical_operation.right, LogicalOperation):
-            right, right_mask = self.eval_logical_operation(
+            right, right_mask = self.execute_logical_operation(
                 logical_operation=logical_operation.right,
                 feat_dict=feat_dict,
                 time_dict=time_dict,
@@ -352,7 +357,7 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
                 filter_na=False,
             )
 
-        out = self.eval_bool_op(left, logical_operation.op, right)
+        out = self.execute_bool_op(left, logical_operation.bool_op, right)
 
         if right_mask is not None:
             mask &= right_mask
@@ -362,16 +367,45 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
 
         return out, mask
 
-    def eval_filter(
+    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,
-    ) -> np.ndarray:
+        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):
-            return self.eval_condition(
+            _mask = self.execute_condition(
                 condition=filter.condition,
                 feat_dict=feat_dict,
                 time_dict=time_dict,
@@ -381,7 +415,7 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
             )[0].to_numpy()
         else:
             assert isinstance(filter.condition, LogicalOperation)
-            return self.eval_logical_operation(
+            _mask = self.execute_logical_operation(
                 logical_operation=filter.condition,
                 feat_dict=feat_dict,
                 time_dict=time_dict,
@@ -389,58 +423,44 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
                 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 eval_pquery(
+    def execute(
         self,
-        query: PQueryDefinition,
+        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]:
-
-        mask = np.ones(len(anchor_time), dtype=bool)
-
-        if query.entity.filter is not None:
-            mask &= self.eval_filter(
-                filter=query.entity.filter,
+        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 getattr(query, 'assuming', None) is not None:
-            if isinstance(query.assuming, Condition):
-                mask &= self.eval_condition(
-                    condition=query.assuming,
-                    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(query.assuming, LogicalOperation)
-                mask &= self.eval_logical_operation(
-                    logical_operation=query.assuming,
-                    feat_dict=feat_dict,
-                    time_dict=time_dict,
-                    batch_dict=batch_dict,
-                    anchor_time=anchor_time,
-                    filter_na=False,
-                )[0].to_numpy()
-
-        if isinstance(query.target, Column):
-            out, _mask = self.eval_column(
-                column=query.target,
+        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, Aggregation):
-            out, _mask = self.eval_aggregation(
-                aggr=query.target,
+        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,
@@ -448,9 +468,9 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
                 filter_na=True,
                 num_forecasts=num_forecasts,
             )
-        elif isinstance(query.target, Condition):
-            out, _mask = self.eval_condition(
-                condition=query.target,
+        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,
@@ -458,10 +478,9 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
                 filter_na=True,
                 num_forecasts=num_forecasts,
             )
-        else:
-            assert isinstance(query.target, LogicalOperation)
-            out, _mask = self.eval_logical_operation(
-                logical_operation=query.target,
+        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,
@@ -469,10 +488,45 @@ class PQueryPandasBackend(PQueryBackend[pd.DataFrame, pd.Series, np.ndarray]):
                 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_ast)} compilation missing.')
+        if query.whatif_ast is not None:
+            if isinstance(query.whatif_ast, Condition):
+                mask &= self.execute_condition(
+                    condition=query.whatif_ast,
+                    feat_dict=feat_dict,
+                    time_dict=time_dict,
+                    batch_dict=batch_dict,
+                    anchor_time=anchor_time,
+                    filter_na=True,
+                    num_forecasts=num_forecasts,
+                )[0]
+            elif isinstance(query.whatif_ast, LogicalOperation):
+                mask &= self.execute_logical_operation(
+                    logical_operation=query.whatif_ast,
+                    feat_dict=feat_dict,
+                    time_dict=time_dict,
+                    batch_dict=batch_dict,
+                    anchor_time=anchor_time,
+                    filter_na=True,
+                    num_forecasts=num_forecasts,
+                )[0]
+            else:
+                raise ValueError(
+                    f'Unsupported ASSUMING condition {type(query.whatif_ast)}')
 
         out = out[mask[_mask]]
         mask &= _mask
-
         out = out.reset_index(drop=True)
-
         return out, mask