kumoai 2.10.0.dev202510061830__cp313-cp313-macosx_11_0_arm64.whl → 2.13.0.dev202511261731__cp313-cp313-macosx_11_0_arm64.whl

kumoai/experimental/rfm/local_graph_sampler.py

@@ -2,7 +2,6 @@ from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import pandas as pd
-from kumoapi.model_plan import RunMode
 from kumoapi.rfm.context import EdgeLayout, Link, Subgraph, Table
 from kumoapi.typing import Stype
 
@@ -33,7 +32,6 @@ class LocalGraphSampler:
         entity_table_names: Tuple[str, ...],
         node: np.ndarray,
         time: np.ndarray,
-        run_mode: RunMode,
         num_neighbors: List[int],
         exclude_cols_dict: Dict[str, List[str]],
     ) -> Subgraph:
@@ -92,15 +90,23 @@ class LocalGraphSampler:
                 )
                 continue
 
-            # Only store unique rows in `df` above a certain threshold:
-            unique_node, inverse_node = np.unique(node, return_inverse=True)
-            if len(node) > 1.05 * len(unique_node):
-                df = df.iloc[unique_node]
-                row = inverse_node
+            row: Optional[np.ndarray] = None
+            if table_name in self._graph_store.end_time_column_dict:
+                # Set end time to NaT for all values greater than anchor time:
+                df = df.iloc[node].reset_index(drop=True)
+                col_name = self._graph_store.end_time_column_dict[table_name]
+                ser = df[col_name]
+                value = ser.astype('datetime64[ns]').astype(int).to_numpy()
+                mask = value > time[batch]
+                df.loc[mask, col_name] = pd.NaT
             else:
-                df = df.iloc[node]
-                row = None
-            df = df.reset_index(drop=True)
+                # Only store unique rows in `df` above a certain threshold:
+                unique_node, inverse = np.unique(node, return_inverse=True)
+                if len(node) > 1.05 * len(unique_node):
+                    df = df.iloc[unique_node].reset_index(drop=True)
+                    row = inverse
+                else:
+                    df = df.iloc[node].reset_index(drop=True)
 
             # Filter data frame to minimal set of columns:
             df = df[columns]

kumoai/experimental/rfm/local_graph_store.py

@@ -45,6 +45,7 @@ class LocalGraphStore:
 
             (
                 self.time_column_dict,
+                self.end_time_column_dict,
                 self.time_dict,
                 self.min_time,
                 self.max_time,
@@ -219,16 +220,21 @@ class LocalGraphStore:
         self,
         graph: LocalGraph,
     ) -> Tuple[
+            Dict[str, str],
             Dict[str, str],
             Dict[str, np.ndarray],
             pd.Timestamp,
             pd.Timestamp,
     ]:
         time_column_dict: Dict[str, str] = {}
+        end_time_column_dict: Dict[str, str] = {}
         time_dict: Dict[str, np.ndarray] = {}
         min_time = pd.Timestamp.max
         max_time = pd.Timestamp.min
         for table in graph.tables.values():
+            if table._end_time_column is not None:
+                end_time_column_dict[table.name] = table._end_time_column
+
             if table._time_column is None:
                 continue
 
@@ -243,7 +249,13 @@ class LocalGraphStore:
                 min_time = min(min_time, time.min())
                 max_time = max(max_time, time.max())
 
-        return time_column_dict, time_dict, min_time, max_time
+        return (
+            time_column_dict,
+            end_time_column_dict,
+            time_dict,
+            min_time,
+            max_time,
+        )
 
     def get_csc(
         self,

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
@@ -27,14 +44,13 @@ class LocalPQueryDriver:
 
     def _get_candidates(
         self,
-        anchor_time: Union[pd.Timestamp, Literal['entity']],
         exclude_node: Optional[np.ndarray] = None,
     ) -> np.ndarray:
 
         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
 
@@ -61,6 +77,30 @@ class LocalPQueryDriver:
 
         return candidate
 
+    def _filter_candidates_by_time(
+        self,
+        candidate: np.ndarray,
+        anchor_time: pd.Timestamp,
+    ) -> np.ndarray:
+
+        entity = self._query.entity_table
+
+        # Filter out entities that do not exist yet in time:
+        time_sec = self._graph_store.time_dict.get(entity)
+        if time_sec is not None:
+            mask = time_sec[candidate] <= (anchor_time.value // (1000**3))
+            candidate = candidate[mask]
+
+        # Filter out entities that no longer exist in time:
+        end_time_col = self._graph_store.end_time_column_dict.get(entity)
+        if end_time_col is not None:
+            ser = self._graph_store.df_dict[entity][end_time_col]
+            ser = ser.iloc[candidate]
+            mask = (anchor_time < ser) | ser.isna().to_numpy()
+            candidate = candidate[mask]
+
+        return candidate
+
     def collect_test(
         self,
         size: int,
@@ -84,7 +124,7 @@ class LocalPQueryDriver:
         """
         batch_size = size if batch_size is None else batch_size
 
-        candidate = self._get_candidates(anchor_time)
+        candidate = self._get_candidates()
 
         nodes: List[np.ndarray] = []
         times: List[pd.Series] = []
@@ -96,19 +136,12 @@ class LocalPQueryDriver:
             node = candidate[candidate_offset:candidate_offset + batch_size]
 
             if isinstance(anchor_time, pd.Timestamp):
-                # Filter out non-existent entities:
-                time = self._graph_store.time_dict.get(
-                    self._query.entity.pkey.table_name)
-                if time is not None:
-                    node = node[time[node] <= (anchor_time.value // (1000**3))]
-
-            if isinstance(anchor_time, pd.Timestamp):
+                node = self._filter_candidates_by_time(node, anchor_time)
                 time = pd.Series(anchor_time).repeat(len(node))
                 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)
@@ -185,7 +218,7 @@ class LocalPQueryDriver:
         """
         batch_size = size if batch_size is None else batch_size
 
-        candidate = self._get_candidates(anchor_time, exclude_node)
+        candidate = self._get_candidates(exclude_node)
 
         if len(candidate) == 0:
             raise RuntimeError("Failed to generate any context examples "
@@ -201,19 +234,12 @@ class LocalPQueryDriver:
             node = candidate[candidate_offset:candidate_offset + batch_size]
 
             if isinstance(anchor_time, pd.Timestamp):
-                # Filter out non-existent entities:
-                time = self._graph_store.time_dict.get(
-                    self._query.entity.pkey.table_name)
-                if time is not None:
-                    node = node[time[node] <= (anchor_time.value // (1000**3))]
-
-            if isinstance(anchor_time, pd.Timestamp):
+                node = self._filter_candidates_by_time(node, anchor_time)
                 time = pd.Series(anchor_time).repeat(len(node))
                 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)
@@ -238,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
@@ -288,37 +315,30 @@ class LocalPQueryDriver:
         mask: Optional[np.ndarray] = None
 
         if isinstance(anchor_time, pd.Timestamp):
-            # Mask out non-existent entities:
-            time = self._graph_store.time_dict.get(
-                self._query.entity.pkey.table_name)
-            if time is not None:
-                mask = time[node] <= (anchor_time.value // (1000**3))
-
-        if isinstance(anchor_time, pd.Timestamp):
+            node = self._filter_candidates_by_time(node, anchor_time)
             time = pd.Series(anchor_time).repeat(len(node))
             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)
@@ -329,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,
@@ -349,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[
@@ -375,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))
 
@@ -407,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]
@@ -436,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,
@@ -447,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
@@ -487,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]