kumoai 2.12.0.dev202511031731__cp313-cp313-macosx_11_0_arm64.whl → 2.12.0.dev202511101731__cp313-cp313-macosx_11_0_arm64.whl

kumoai/_version.py

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

kumoai/client/endpoints.py

@@ -147,3 +147,4 @@ class RFMEndpoints:
     explain = Endpoint(f"{BASE}/explain", HTTPMethod.POST)
     evaluate = Endpoint(f"{BASE}/evaluate", HTTPMethod.POST)
     validate_query = Endpoint(f"{BASE}/validate_query", HTTPMethod.POST)
+    parse_query = Endpoint(f"{BASE}/parse_query", HTTPMethod.POST)

kumoai/client/rfm.py

@@ -1,7 +1,11 @@
+from typing import Any
+
 from kumoapi.json_serde import to_json_dict
 from kumoapi.rfm import (
     RFMEvaluateResponse,
     RFMExplanationResponse,
+    RFMParseQueryRequest,
+    RFMParseQueryResponse,
     RFMPredictResponse,
     RFMValidateQueryRequest,
     RFMValidateQueryResponse,
@@ -26,25 +30,32 @@ class RFMAPI:
         Returns:
             RFMPredictResponse containing the predictions
         """
-        # Send binary data to the predict endpoint
         response = self._client._request(
-            RFMEndpoints.predict, data=request,
-            headers={'Content-Type': 'application/x-protobuf'})
+            RFMEndpoints.predict,
+            data=request,
+            headers={'Content-Type': 'application/x-protobuf'},
+        )
         raise_on_error(response)
         return parse_response(RFMPredictResponse, response)
 
-    def explain(self, request: bytes) -> RFMExplanationResponse:
+    def explain(
+        self,
+        request: bytes,
+        skip_summary: bool = False,
+    ) -> RFMExplanationResponse:
         """Explain the RFM model on the given context.
 
         Args:
             request: The predict request as serialized protobuf.
+            skip_summary: Whether to skip generating a human-readable summary
+                of the explanation.
 
         Returns:
             RFMPredictResponse containing the explanations
         """
-        # Send binary data to the explain endpoint
+        params: dict[str, Any] = {'generate_summary': not skip_summary}
         response = self._client._request(
-            RFMEndpoints.explain, data=request,
+            RFMEndpoints.explain, data=request, params=params,
             headers={'Content-Type': 'application/x-protobuf'})
         raise_on_error(response)
         return parse_response(RFMExplanationResponse, response)
@@ -58,7 +69,6 @@ class RFMAPI:
         Returns:
             RFMEvaluateResponse containing the computed metrics
         """
-        # Send binary data to the evaluate endpoint
         response = self._client._request(
             RFMEndpoints.evaluate, data=request,
             headers={'Content-Type': 'application/x-protobuf'})
@@ -82,3 +92,21 @@ class RFMAPI:
                                          json=to_json_dict(request))
         raise_on_error(response)
         return parse_response(RFMValidateQueryResponse, response)
+
+    def parse_query(
+        self,
+        request: RFMParseQueryRequest,
+    ) -> RFMParseQueryResponse:
+        """Validate a predictive query against a graph.
+
+        Args:
+            request: The request object containing
+                the query and graph definition
+
+        Returns:
+            RFMParseQueryResponse containing the QueryDefinition
+        """
+        response = self._client._request(RFMEndpoints.parse_query,
+                                         json=to_json_dict(request))
+        raise_on_error(response)
+        return parse_response(RFMParseQueryResponse, response)

kumoai/experimental/rfm/__init__.py

@@ -36,7 +36,7 @@ import os
 import kumoai
 from .local_table import LocalTable
 from .local_graph import LocalGraph
-from .rfm import KumoRFM
+from .rfm import ExplainConfig, Explanation, KumoRFM
 from .authenticate import authenticate
 
 
@@ -60,6 +60,8 @@ __all__ = [
     'LocalTable',
     'LocalGraph',
     'KumoRFM',
+    'ExplainConfig',
+    'Explanation',
     'authenticate',
     'init',
 ]

kumoai/experimental/rfm/local_pquery_driver.py

@@ -1,23 +1,40 @@
 import warnings
-from typing import Dict, List, Literal, Optional, Tuple, Union
+from typing import Dict, List, Literal, NamedTuple, Optional, Set, Tuple, Union
 
 import numpy as np
 import pandas as pd
-from kumoapi.pquery import QueryType
-from kumoapi.rfm import PQueryDefinition
+from kumoapi.pquery import QueryType, ValidatedPredictiveQuery
+from kumoapi.pquery.AST import (
+    Aggregation,
+    ASTNode,
+    Column,
+    Condition,
+    Filter,
+    Join,
+    LogicalOperation,
+)
+from kumoapi.task import TaskType
+from kumoapi.typing import AggregationType, DateOffset, Stype
 
 import kumoai.kumolib as kumolib
 from kumoai.experimental.rfm.local_graph_store import LocalGraphStore
-from kumoai.experimental.rfm.pquery import PQueryPandasBackend
+from kumoai.experimental.rfm.pquery import PQueryPandasExecutor
 
 _coverage_warned = False
 
 
+class SamplingSpec(NamedTuple):
+    edge_type: Tuple[str, str, str]
+    hop: int
+    start_offset: Optional[DateOffset]
+    end_offset: Optional[DateOffset]
+
+
 class LocalPQueryDriver:
     def __init__(
         self,
         graph_store: LocalGraphStore,
-        query: PQueryDefinition,
+        query: ValidatedPredictiveQuery,
         random_seed: Optional[int] = None,
     ) -> None:
         self._graph_store = graph_store
@@ -33,7 +50,7 @@ class LocalPQueryDriver:
         if self._query.query_type == QueryType.TEMPORAL:
             assert exclude_node is None
 
-        table_name = self._query.entity.pkey.table_name
+        table_name = self._query.entity_table
         num_nodes = len(self._graph_store.df_dict[table_name])
         mask_dict = self._graph_store.mask_dict
 
@@ -66,7 +83,7 @@ class LocalPQueryDriver:
         anchor_time: pd.Timestamp,
     ) -> np.ndarray:
 
-        entity = self._query.entity.pkey.table_name
+        entity = self._query.entity_table
 
         # Filter out entities that do not exist yet in time:
         time_sec = self._graph_store.time_dict.get(entity)
@@ -124,8 +141,7 @@ class LocalPQueryDriver:
                 time = time.astype('datetime64[ns]').reset_index(drop=True)
             else:
                 assert anchor_time == 'entity'
-                time = self._graph_store.time_dict[
-                    self._query.entity.pkey.table_name]
+                time = self._graph_store.time_dict[self._query.entity_table]
                 time = pd.Series(time[node] * 1000**3, dtype='datetime64[ns]')
 
             y, mask = self(node, time)
@@ -223,8 +239,7 @@ class LocalPQueryDriver:
                 time = time.astype('datetime64[ns]').reset_index(drop=True)
             else:
                 assert anchor_time == 'entity'
-                time = self._graph_store.time_dict[
-                    self._query.entity.pkey.table_name]
+                time = self._graph_store.time_dict[self._query.entity_table]
                 time = pd.Series(time[node] * 1000**3, dtype='datetime64[ns]')
 
             y, mask = self(node, time)
@@ -249,7 +264,8 @@ class LocalPQueryDriver:
                     reached_end = True
                     break
                 candidate_offset = 0
-                anchor_time = anchor_time - (self._query.target.end_offset *
+                time_frame = self._query.target_timeframe.timeframe
+                anchor_time = anchor_time - (time_frame *
                                              self._query.num_forecasts)
                 if anchor_time < self._graph_store.min_time:
                     reached_end = True
@@ -304,26 +320,25 @@ class LocalPQueryDriver:
             time = time.astype('datetime64[ns]').reset_index(drop=True)
         else:
             assert anchor_time == 'entity'
-            time = self._graph_store.time_dict[
-                self._query.entity.pkey.table_name]
+            time = self._graph_store.time_dict[self._query.entity_table]
             time = pd.Series(time[node] * 1000**3, dtype='datetime64[ns]')
 
-        if self._query.entity.filter is not None:
+        if isinstance(self._query.entity_ast, Filter):
             # Mask out via (temporal) entity filter:
-            backend = PQueryPandasBackend()
+            executor = PQueryPandasExecutor()
             masks: List[np.ndarray] = []
             for start in range(0, len(node), batch_size):
                 feat_dict, time_dict, batch_dict = self._sample(
                     node[start:start + batch_size],
                     time.iloc[start:start + batch_size],
                 )
-                _mask = backend.eval_filter(
-                    filter=self._query.entity.filter,
+                _mask = executor.execute_filter(
+                    filter=self._query.entity_ast,
                     feat_dict=feat_dict,
                     time_dict=time_dict,
                     batch_dict=batch_dict,
                     anchor_time=time.iloc[start:start + batch_size],
-                )
+                )[1]
                 masks.append(_mask)
 
             _mask = np.concatenate(masks)
@@ -334,6 +349,96 @@ class LocalPQueryDriver:
 
         return mask
 
+    def _get_sampling_specs(
+        self,
+        node: ASTNode,
+        hop: int,
+        seed_table_name: str,
+        edge_types: List[Tuple[str, str, str]],
+        num_forecasts: int = 1,
+    ) -> List[SamplingSpec]:
+        if isinstance(node, (Aggregation, Column)):
+            if isinstance(node, Column):
+                table_name = node.fqn.split('.')[0]
+                if seed_table_name == table_name:
+                    return []
+            else:
+                table_name = node._get_target_column_name().split('.')[0]
+
+            target_edge_types = [
+                edge_type for edge_type in edge_types if
+                edge_type[2] == seed_table_name and edge_type[0] == table_name
+            ]
+            if len(target_edge_types) != 1:
+                raise ValueError(
+                    f"Could not find a unique foreign key from table "
+                    f"'{seed_table_name}' to '{table_name}'")
+
+            if isinstance(node, Column):
+                return [
+                    SamplingSpec(
+                        edge_type=target_edge_types[0],
+                        hop=hop + 1,
+                        start_offset=None,
+                        end_offset=None,
+                    )
+                ]
+            spec = SamplingSpec(
+                edge_type=target_edge_types[0],
+                hop=hop + 1,
+                start_offset=node.aggr_time_range.start_date_offset,
+                end_offset=node.aggr_time_range.end_date_offset *
+                num_forecasts,
+            )
+            return [spec] + self._get_sampling_specs(
+                node.target, hop=hop + 1, seed_table_name=table_name,
+                edge_types=edge_types, num_forecasts=num_forecasts)
+        specs = []
+        for child in node.children:
+            specs += self._get_sampling_specs(child, hop, seed_table_name,
+                                              edge_types, num_forecasts)
+        return specs
+
+    def get_sampling_specs(self) -> List[SamplingSpec]:
+        edge_types = self._graph_store.edge_types
+        specs = self._get_sampling_specs(
+            self._query.target_ast, hop=0,
+            seed_table_name=self._query.entity_table, edge_types=edge_types,
+            num_forecasts=self._query.num_forecasts)
+        specs += self._get_sampling_specs(
+            self._query.entity_ast, hop=0,
+            seed_table_name=self._query.entity_table, edge_types=edge_types)
+        if self._query.whatif_ast is not None:
+            specs += self._get_sampling_specs(
+                self._query.whatif_ast, hop=0,
+                seed_table_name=self._query.entity_table,
+                edge_types=edge_types)
+        # Group specs according to edge type and hop:
+        spec_dict: Dict[
+            Tuple[Tuple[str, str, str], int],
+            Tuple[Optional[DateOffset], Optional[DateOffset]],
+        ] = {}
+        for spec in specs:
+            if (spec.edge_type, spec.hop) not in spec_dict:
+                spec_dict[(spec.edge_type, spec.hop)] = (
+                    spec.start_offset,
+                    spec.end_offset,
+                )
+            else:
+                start_offset, end_offset = spec_dict[(
+                    spec.edge_type,
+                    spec.hop,
+                )]
+                spec_dict[(spec.edge_type, spec.hop)] = (
+                    min_date_offset(start_offset, spec.start_offset),
+                    max_date_offset(end_offset, spec.end_offset),
+                )
+
+        return [
+            SamplingSpec(edge, hop, start_offset, end_offset)
+            for (edge, hop), (start_offset, end_offset) in spec_dict.items()
+        ]
+
     def _sample(
         self,
         node: np.ndarray,
@@ -354,7 +459,7 @@ class LocalPQueryDriver:
             The feature dictionary, the time column dictionary and the batch
             dictionary.
         """
-        specs = self._query.get_sampling_specs(self._graph_store.edge_types)
+        specs = self.get_sampling_specs()
         num_hops = max([spec.hop for spec in specs] + [0])
         num_neighbors: Dict[Tuple[str, str, str], list[int]] = {}
         time_offsets: Dict[
@@ -380,7 +485,7 @@ class LocalPQueryDriver:
 
         edge_types = list(num_neighbors.keys()) + list(time_offsets.keys())
         node_types = list(
-            set([self._query.entity.pkey.table_name])
+            set([self._query.entity_table])
             | set(src for src, _, _ in edge_types)
             | set(dst for _, _, dst in edge_types))
 
@@ -412,21 +517,33 @@ class LocalPQueryDriver:
                 '__'.join(edge_type): np.array(values)
                 for edge_type, values in time_offsets.items()
             },
-            self._query.entity.pkey.table_name,
+            self._query.entity_table,
             node,
             anchor_time.astype(int).to_numpy() // 1000**3,
         )
 
         feat_dict: Dict[str, pd.DataFrame] = {}
         time_dict: Dict[str, pd.Series] = {}
-        column_dict = self._query.column_dict
-        time_tables = self._query.time_tables
+        column_dict: Dict[str, Set[str]] = {}
+        for col in self._query.all_query_columns:
+            table_name, col_name = col.split('.')
+            if table_name not in column_dict:
+                column_dict[table_name] = set()
+            if col_name != '*':
+                column_dict[table_name].add(col_name)
+        time_tables = self.find_time_tables()
         for table_name in set(list(column_dict.keys()) + time_tables):
             df = self._graph_store.df_dict[table_name]
             row_id = node_dict[table_name]
             df = df.iloc[row_id].reset_index(drop=True)
             if table_name in column_dict:
-                feat_dict[table_name] = df[list(column_dict[table_name])]
+                if len(column_dict[table_name]) == 0:
+                    # We are dealing with COUNT(table.*), insert a dummy col
+                    # to ensure we don't lose the information on node count
+                    feat_dict[table_name] = pd.DataFrame(
+                        {'ones': [1] * len(df)})
+                else:
+                    feat_dict[table_name] = df[list(column_dict[table_name])]
             if table_name in time_tables:
                 time_col = self._graph_store.time_column_dict[table_name]
                 time_dict[table_name] = df[time_col]
@@ -441,7 +558,7 @@ class LocalPQueryDriver:
 
         feat_dict, time_dict, batch_dict = self._sample(node, anchor_time)
 
-        y, mask = PQueryPandasBackend().eval_pquery(
+        y, mask = PQueryPandasExecutor().execute(
             query=self._query,
             feat_dict=feat_dict,
             time_dict=time_dict,
@@ -452,6 +569,62 @@ class LocalPQueryDriver:
 
         return y, mask
 
+    def find_time_tables(self) -> List[str]:
+        def _find_time_tables(node: ASTNode) -> List[str]:
+            time_tables = []
+            if isinstance(node, Aggregation):
+                time_tables.append(
+                    node._get_target_column_name().split('.')[0])
+            for child in node.children:
+                time_tables += _find_time_tables(child)
+            return time_tables
+
+        time_tables = _find_time_tables(
+            self._query.target_ast) + _find_time_tables(self._query.entity_ast)
+        if self._query.whatif_ast is not None:
+            time_tables += _find_time_tables(self._query.whatif_ast)
+        return list(set(time_tables))
+
+    @staticmethod
+    def get_task_type(
+        query: ValidatedPredictiveQuery,
+        edge_types: List[Tuple[str, str, str]],
+    ) -> TaskType:
+        if isinstance(query.target_ast, (Condition, LogicalOperation)):
+            return TaskType.BINARY_CLASSIFICATION
+
+        target = query.target_ast
+        if isinstance(target, Join):
+            target = target.rhs_target
+        if isinstance(target, Aggregation):
+            if target.aggr == AggregationType.LIST_DISTINCT:
+                table_name, col_name = target._get_target_column_name().split(
+                    '.')
+                target_edge_types = [
+                    edge_type for edge_type in edge_types
+                    if edge_type[0] == table_name and edge_type[1] == col_name
+                ]
+                if len(target_edge_types) != 1:
+                    raise NotImplementedError(
+                        f"Multilabel-classification queries based on "
+                        f"'LIST_DISTINCT' are not supported yet. If you "
+                        f"planned to write a link prediction query instead, "
+                        f"make sure to register '{col_name}' as a "
+                        f"foreign key.")
+                return TaskType.TEMPORAL_LINK_PREDICTION
+
+            return TaskType.REGRESSION
+
+        assert isinstance(target, Column)
+
+        if target.stype in {Stype.ID, Stype.categorical}:
+            return TaskType.MULTICLASS_CLASSIFICATION
+
+        if target.stype in {Stype.numerical}:
+            return TaskType.REGRESSION
+
+        raise NotImplementedError("Task type not yet supported")
+
 
 def date_offset_to_seconds(offset: pd.DateOffset) -> int:
     r"""Convert a :class:`pandas.DateOffset` into a maximum number of
@@ -492,3 +665,25 @@ def date_offset_to_seconds(offset: pd.DateOffset) -> int:
             total_ns += scaled_value
 
     return total_ns
+
+
+def min_date_offset(*args: Optional[DateOffset]) -> Optional[DateOffset]:
+    if any(arg is None for arg in args):
+        return None
+
+    anchor = pd.Timestamp('2000-01-01')
+    timestamps = [anchor + arg for arg in args]
+    assert len(timestamps) > 0
+    argmin = min(range(len(timestamps)), key=lambda i: timestamps[i])
+    return args[argmin]
+
+
+def max_date_offset(*args: DateOffset) -> DateOffset:
+    if any(arg is None for arg in args):
+        return None
+
+    anchor = pd.Timestamp('2000-01-01')
+    timestamps = [anchor + arg for arg in args]
+    assert len(timestamps) > 0
+    argmax = max(range(len(timestamps)), key=lambda i: timestamps[i])
+    return args[argmax]

kumoai/experimental/rfm/pquery/__init__.py

@@ -1,11 +1,7 @@
-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/pandas_executor.py

@@ -118,7 +118,7 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
             target_feat, target_mask = self.execute_column(
                 column=aggr.target,
                 feat_dict=feat_dict,
-                filter_na=False,
+                filter_na=True,
             )
         else:
             assert isinstance(aggr.target, Filter)
@@ -128,7 +128,7 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
                 time_dict=time_dict,
                 batch_dict=batch_dict,
                 anchor_time=anchor_time,
-                filter_na=False,
+                filter_na=True,
             )
 
         outs: List[pd.Series] = []
@@ -137,19 +137,20 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
             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)
+            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
-                curr_target_mask &= (target_time
+                time_filter_mask &= (target_time
                                      > anchor_target_time + start_offset)
             else:
                 assert num_forecasts == 1
+            curr_target_mask = target_mask & time_filter_mask
 
             out, mask = self.execute_aggregation_type(
                 aggr.aggr,
-                feat=target_feat[curr_target_mask],
+                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,
@@ -499,7 +500,32 @@ class PQueryPandasExecutor(PQueryExecutor[pd.DataFrame, pd.Series,
             )
         else:
             raise NotImplementedError(
-                f'{type(query.target)} compilation missing.')
+                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)