kumoai 2.13.0.dev202512091732__cp311-cp311-macosx_11_0_arm64.whl → 2.14.0.dev202601051732__cp311-cp311-macosx_11_0_arm64.whl

kumoai/experimental/rfm/base/sampler.py

@@ -1,23 +1,29 @@
 import copy
 import re
+import warnings
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple
 
 import numpy as np
 import pandas as pd
-from kumoapi.pquery import ValidatedPredictiveQuery
+from kumoapi.pquery import QueryType, ValidatedPredictiveQuery
 from kumoapi.pquery.AST import Aggregation, ASTNode
 from kumoapi.rfm.context import EdgeLayout, Link, Subgraph, Table
 from kumoapi.typing import Stype
 
+from kumoai.utils import ProgressLogger
+
 if TYPE_CHECKING:
     from kumoai.experimental.rfm import Graph
 
+_coverage_warned = False
+
 
 @dataclass
 class SamplerOutput:
+    anchor_time: np.ndarray
     df_dict: dict[str, pd.DataFrame]
     inverse_dict: dict[str, np.ndarray]
     batch_dict: dict[str, np.ndarray]
@@ -27,16 +33,26 @@ class SamplerOutput:
     num_sampled_edges_dict: dict[tuple[str, str, str], list[int]]
 
 
-@dataclass
-class TargetOutput:
+class TargetOutput(NamedTuple):
     entity_pkey: pd.Series
     anchor_time: pd.Series
     target: pd.Series
-    num_trials: int
 
 
 class Sampler(ABC):
-    def __init__(self, graph: 'Graph') -> None:
+    r"""A base class to sample relational data (*i.e.*, subgraphs and
+    ground-truth targets) from a custom backend.
+
+    Args:
+        graph: The graph.
+        verbose: Whether to print verbose output.
+    """
+    def __init__(
+        self,
+        graph: 'Graph',
+        verbose: bool | ProgressLogger = True,
+    ) -> None:
+
         self._edge_types: list[tuple[str, str, str]] = []
         for edge in graph.edges:
             edge_type = (edge.src_table, edge.fkey, edge.dst_table)
@@ -72,37 +88,99 @@ class Sampler(ABC):
                     continue
                 self._table_stype_dict[table.name][column.name] = column.stype
 
+        self._min_time_dict: dict[str, pd.Timestamp] = {}
+        self._max_time_dict: dict[str, pd.Timestamp] = {}
+
+    # Properties ##############################################################
+
     @property
     def edge_types(self) -> list[tuple[str, str, str]]:
+        r"""All available edge types in the graph."""
         return self._edge_types
 
     @property
     def primary_key_dict(self) -> dict[str, str]:
+        r"""All available primary keys in the graph."""
         return self._primary_key_dict
 
     @property
     def time_column_dict(self) -> dict[str, str]:
+        r"""All available time columns in the graph."""
         return self._time_column_dict
 
     @property
     def end_time_column_dict(self) -> dict[str, str]:
+        r"""All available end time columns in the graph."""
         return self._end_time_column_dict
 
     @property
     def table_stype_dict(self) -> dict[str, dict[str, Stype]]:
+        r"""The registered semantic types for all feature columns in all tables
+        in the graph.
+        """
         return self._table_stype_dict
 
+    def get_min_time(
+        self,
+        table_names: list[str] | None = None,
+    ) -> pd.Timestamp:
+        r"""Returns the minimal timestamp in the union of a set of tables.
+
+        Args:
+            table_names: The set of tables.
+        """
+        if table_names is None or len(table_names) == 0:
+            table_names = list(self.time_column_dict.keys())
+        unknown = list(set(table_names) - set(self._min_time_dict.keys()))
+        if len(unknown) > 0:
+            min_max_time_dict = self._get_min_max_time_dict(unknown)
+            for table_name, (min_time, max_time) in min_max_time_dict.items():
+                self._min_time_dict[table_name] = min_time
+                self._max_time_dict[table_name] = max_time
+        return min([self._min_time_dict[table]
+                    for table in table_names] + [pd.Timestamp.max])
+
+    def get_max_time(
+        self,
+        table_names: list[str] | None = None,
+    ) -> pd.Timestamp:
+        r"""Returns the maximum timestamp in the union of a set of tables.
+
+        Args:
+            table_names: The set of tables.
+        """
+        if table_names is None or len(table_names) == 0:
+            table_names = list(self.time_column_dict.keys())
+        unknown = list(set(table_names) - set(self._max_time_dict.keys()))
+        if len(unknown) > 0:
+            min_max_time_dict = self._get_min_max_time_dict(unknown)
+            for table_name, (min_time, max_time) in min_max_time_dict.items():
+                self._min_time_dict[table_name] = min_time
+                self._max_time_dict[table_name] = max_time
+        return max([self._max_time_dict[table]
+                    for table in table_names] + [pd.Timestamp.min])
+
+    # Subgraph Sampling #######################################################
+
     def sample_subgraph(
         self,
         entity_table_names: tuple[str, ...],
         entity_pkey: pd.Series,
-        anchor_time: pd.Series,
+        anchor_time: pd.Series | Literal['entity'],
         num_neighbors: list[int],
         exclude_cols_dict: dict[str, list[str]] | None = None,
     ) -> Subgraph:
-
+        r"""Samples distinct subgraphs for each entity primary key.
+
+        Args:
+            entity_table_names: The entity table names.
+            entity_pkey: The primary keys to use as seed nodes.
+            anchor_time: The anchor time of the subgraphs.
+            num_neighbors: The number of neighbors to sample for each hop.
+            exclude_cols_dict: The columns to exclude from the subgraph.
+        """
         # Exclude all columns that leak target information:
-        table_stype_dict: dict[str, dict[str, Stype]] = self._table_stype_dict
+        table_stype_dict: dict[str, dict[str, Stype]] = self.table_stype_dict
         if exclude_cols_dict is not None:
             table_stype_dict = copy.deepcopy(table_stype_dict)
             for table_name, exclude_cols in exclude_cols_dict.items():
@@ -118,7 +196,8 @@ class Sampler(ABC):
         for table_name in entity_table_names:
             columns_dict[table_name].add(self.primary_key_dict[table_name])
 
-        if anchor_time.dtype != 'datetime64[ns]':
+        if (isinstance(anchor_time, pd.Series)
+                and anchor_time.dtype != 'datetime64[ns]'):
             anchor_time = anchor_time.astype('datetime64[ns]')
 
         out = self._sample_subgraph(
@@ -129,8 +208,9 @@ class Sampler(ABC):
             num_neighbors=num_neighbors,
         )
 
+        # Parse `SubgraphOutput` into `Subgraph` structure:
         subgraph = Subgraph(
-            anchor_time=anchor_time.astype(int).to_numpy(),
+            anchor_time=out.anchor_time,
             table_dict={},
             link_dict={},
         )
@@ -148,11 +228,8 @@ class Sampler(ABC):
                 # Set end time to NaT for all values greater than anchor time:
                 assert table_name not in out.inverse_dict
                 ser = df[end_time_column]
-                if ser.dtype != 'datetime64[ns]':
-                    ser = ser.astype('datetime64[ns]')
-                mask = ser > anchor_time.iloc[batch]
-                ser.iloc[mask] = pd.NaT
-                df[end_time_column] = ser
+                mask = ser.astype(int).to_numpy() > out.anchor_time[batch]
+                df.loc[mask, end_time_column] = pd.NaT
 
             stype_dict = table_stype_dict[table_name]
             for column_name, stype in stype_dict.items():
@@ -213,24 +290,31 @@ class Sampler(ABC):
 
         return subgraph
 
-    def sample_target(
+    # Predictive Query ########################################################
+
+    def _get_query_columns_dict(
         self,
         query: ValidatedPredictiveQuery,
-        num_examples: int,
-        anchor_time: pd.Timestamp | Literal['entity'],
-        random_seed: int | None = None,
-    ) -> TargetOutput:
-
+    ) -> dict[str, set[str]]:
         columns_dict: dict[str, set[str]] = defaultdict(set)
         for fqn in query.all_query_columns + [query.entity_column]:
             table_name, column_name = fqn.split('.')
+            if column_name == '*':
+                continue
             columns_dict[table_name].add(column_name)
-
         if column_name := self.time_column_dict.get(query.entity_table):
             columns_dict[table_name].add(column_name)
         if column_name := self.end_time_column_dict.get(query.entity_table):
             columns_dict[table_name].add(column_name)
+        return columns_dict
 
+    def _get_query_time_offset_dict(
+        self,
+        query: ValidatedPredictiveQuery,
+    ) -> dict[
+            tuple[str, str, str],
+            tuple[pd.DateOffset | None, pd.DateOffset],
+    ]:
         time_offset_dict: dict[
             tuple[str, str, str],
             tuple[pd.DateOffset | None, pd.DateOffset],
@@ -239,7 +323,6 @@ class Sampler(ABC):
         def _add_time_offset(node: ASTNode, num_forecasts: int = 1) -> None:
             if isinstance(node, Aggregation):
                 table_name = node._get_target_column_name().split('.')[0]
-                columns_dict[table_name].add(self.time_column_dict[table_name])
 
                 edge_types = [
                     edge_type for edge_type in self.edge_types
@@ -273,42 +356,342 @@ class Sampler(ABC):
         if query.whatif_ast is not None:
             _add_time_offset(query.whatif_ast)
 
-        return self._sample_target(
+        return time_offset_dict
+
+    def sample_target(
+        self,
+        query: ValidatedPredictiveQuery,
+        num_train_examples: int,
+        train_anchor_time: pd.Timestamp | Literal['entity'],
+        num_train_trials: int,
+        num_test_examples: int,
+        test_anchor_time: pd.Timestamp | Literal['entity'],
+        num_test_trials: int,
+        random_seed: int | None = None,
+    ) -> tuple[TargetOutput, TargetOutput]:
+        r"""Samples ground-truth targets given a predictive query, split into
+        training and test set.
+
+        Args:
+            query: The predictive query.
+            num_train_examples: How many training examples to produce.
+            train_anchor_time: The anchor timestamp for the training set.
+                If set to ``"entity"``, will use the timestamp of the entity.
+            num_train_trials: The number of training examples to try before
+                aborting.
+            num_test_examples: How many test examples to produce.
+            test_anchor_time: The anchor timestamp for the test set.
+                If set to ``"entity"``, will use the timestamp of the entity.
+            num_test_trials: The number of test examples to try before
+                aborting.
+            random_seed: A manual seed for generating pseudo-random numbers.
+        """
+        rng = np.random.default_rng(random_seed)
+
+        if num_train_examples == 0 or num_train_trials == 0:
+            num_train_examples = num_train_trials = 0
+        if num_test_examples == 0 or num_test_trials == 0:
+            num_test_examples = num_test_trials = 0
+
+        # 1. Collect information on what to query #############################
+        columns_dict = self._get_query_columns_dict(query)
+        time_offset_dict = self._get_query_time_offset_dict(query)
+        for table_name, _, _ in time_offset_dict.keys():
+            columns_dict[table_name].add(self.time_column_dict[table_name])
+
+        # 2. Sample random rows from entity table #############################
+        shared_train_test = query.query_type == QueryType.STATIC
+        shared_train_test &= train_anchor_time == test_anchor_time
+        if shared_train_test:
+            num_entity_rows = num_train_trials + num_test_trials
+        else:
+            num_entity_rows = max(num_train_trials, num_test_trials)
+        assert num_entity_rows > 0
+
+        entity_df = self._sample_entity_table(
+            table_name=query.entity_table,
+            columns=columns_dict[query.entity_table],
+            num_rows=num_entity_rows,
+            random_seed=random_seed,
+        )
+
+        if len(entity_df) == 0:
+            raise ValueError("Failed to find any rows in the entity table "
+                             "'{query.entity_table}'.")
+
+        entity_pkey = entity_df[self.primary_key_dict[query.entity_table]]
+        entity_time: pd.Series | None = None
+        if column_name := self.time_column_dict.get(query.entity_table):
+            entity_time = entity_df[column_name]
+        entity_end_time: pd.Series | None = None
+        if column_name := self.end_time_column_dict.get(query.entity_table):
+            entity_end_time = entity_df[column_name]
+
+        def get_valid_entity_index(
+            time: pd.Timestamp | Literal['entity'],
+            max_size: int | None = None,
+        ) -> np.ndarray:
+
+            if time == 'entity':
+                index: np.ndarray = np.arange(len(entity_pkey))
+            elif entity_time is None and entity_end_time is None:
+                index = np.arange(len(entity_pkey))
+            else:
+                mask: np.ndarray | None = None
+                if entity_time is not None:
+                    mask = (entity_time <= time).to_numpy()
+                if entity_end_time is not None:
+                    _mask = (entity_end_time > time).to_numpy()
+                    _mask |= entity_end_time.isna().to_numpy()
+                    mask = _mask if mask is None else mask & _mask
+                assert mask is not None
+                index = mask.nonzero()[0]
+
+            rng.shuffle(index)
+
+            if max_size is not None:
+                index = index[:max_size]
+
+            return index
+
+        # 3. Build training and test candidates ###############################
+        train_index = test_index = np.array([], dtype=np.int64)
+        train_time = test_time = pd.Series([], dtype='datetime64[ns]')
+
+        if shared_train_test:
+            train_index = get_valid_entity_index(train_anchor_time)
+            if train_anchor_time == 'entity':  # Sort by timestamp:
+                assert entity_time is not None
+                train_time = entity_time.iloc[train_index]
+                train_time = train_time.reset_index(drop=True)
+                train_time = train_time.sort_values(ascending=False)
+                perm = train_time.index.to_numpy()
+                train_index = train_index[perm]
+                train_time = train_time.reset_index(drop=True)
+            else:
+                train_time = to_ser(train_anchor_time, size=len(train_index))
+        else:
+            if num_test_examples > 0:
+                test_index = get_valid_entity_index(  #
+                    test_anchor_time, max_size=num_test_trials)
+                assert test_anchor_time != 'entity'
+                test_time = to_ser(test_anchor_time, len(test_index))
+
+            if query.query_type == QueryType.STATIC and num_train_examples > 0:
+                train_index = get_valid_entity_index(  #
+                    train_anchor_time, max_size=num_train_trials)
+                assert train_anchor_time != 'entity'
+                train_time = to_ser(train_anchor_time, len(train_index))
+            elif query.query_type == QueryType.TEMPORAL and num_train_examples:
+                aggr_table_names = [
+                    aggr._get_target_column_name().split('.')[0]
+                    for aggr in query.get_all_target_aggregations()
+                ]
+                offset = query.target_timeframe.timeframe * query.num_forecasts
+
+                train_indices: list[np.ndarray] = []
+                train_times: list[pd.Series] = []
+                while True:
+                    train_index = get_valid_entity_index(  #
+                        train_anchor_time, max_size=num_train_trials)
+                    assert train_anchor_time != 'entity'
+                    train_time = to_ser(train_anchor_time, len(train_index))
+                    train_indices.append(train_index)
+                    train_times.append(train_time)
+                    if sum(len(x) for x in train_indices) >= num_train_trials:
+                        break
+                    train_anchor_time -= offset
+                    if train_anchor_time < self.get_min_time(aggr_table_names):
+                        break
+                train_index = np.concatenate(train_indices, axis=0)
+                train_index = train_index[:num_train_trials]
+                train_time = pd.concat(train_times, axis=0, ignore_index=True)
+                train_time = train_time.iloc[:num_train_trials]
+
+        # 4. Sample training and test labels ##################################
+        train_y, train_mask, test_y, test_mask = self._sample_target(
             query=query,
-            num_examples=num_examples,
-            anchor_time=anchor_time,
+            entity_df=entity_df,
+            train_index=train_index,
+            train_time=train_time,
+            num_train_examples=(num_train_examples + num_test_examples
+                                if shared_train_test else num_train_examples),
+            test_index=test_index,
+            test_time=test_time,
+            num_test_examples=0 if shared_train_test else num_test_examples,
             columns_dict=columns_dict,
             time_offset_dict=time_offset_dict,
-            random_seed=random_seed,
+        )
+
+        # 5. Post-processing ##################################################
+        if shared_train_test:
+            num_examples = num_train_examples + num_test_examples
+            train_index = train_index[train_mask][:num_examples]
+            train_time = train_time.iloc[train_mask].iloc[:num_examples]
+            train_y = train_y.iloc[:num_examples]
+
+            _num_test = num_test_examples
+            _num_train = min(num_train_examples, 1000)
+            if (num_test_examples > 0 and num_train_examples > 0
+                    and len(train_y) < num_examples
+                    and len(train_y) < _num_test + _num_train):
+                # Not enough labels to satisfy requested split without losing
+                # large number of training examples:
+                _num_test = len(train_y) - _num_train
+                if _num_test < _num_train:  # Fallback to 50/50 split:
+                    _num_test = len(train_y) // 2
+
+            test_index = train_index[:_num_test]
+            test_pkey = entity_pkey.iloc[test_index]
+            test_time = train_time.iloc[:_num_test]
+            test_y = train_y.iloc[:_num_test]
+
+            train_index = train_index[_num_test:]
+            train_pkey = entity_pkey.iloc[train_index]
+            train_time = train_time.iloc[_num_test:]
+            train_y = train_y.iloc[_num_test:]
+        else:
+            train_index = train_index[train_mask][:num_train_examples]
+            train_pkey = entity_pkey.iloc[train_index]
+            train_time = train_time.iloc[train_mask].iloc[:num_train_examples]
+            train_y = train_y.iloc[:num_train_examples]
+
+            test_index = test_index[test_mask][:num_test_examples]
+            test_pkey = entity_pkey.iloc[test_index]
+            test_time = test_time.iloc[test_mask].iloc[:num_test_examples]
+            test_y = test_y.iloc[:num_test_examples]
+
+        train_pkey = train_pkey.reset_index(drop=True)
+        train_time = train_time.reset_index(drop=True)
+        train_y = train_y.reset_index(drop=True)
+        test_pkey = test_pkey.reset_index(drop=True)
+        test_time = test_time.reset_index(drop=True)
+        test_y = test_y.reset_index(drop=True)
+
+        if num_train_examples > 0 and len(train_y) == 0:
+            raise RuntimeError("Failed to collect any context examples. Is "
+                               "your predictive query too restrictive?")
+
+        if num_test_examples > 0 and len(test_y) == 0:
+            raise RuntimeError("Failed to collect any test examples for "
+                               "evaluation. Is your predictive query too "
+                               "restrictive?")
+
+        global _coverage_warned
+        if (not num_train_examples > 0  #
+                and not _coverage_warned  #
+                and len(entity_df) >= num_entity_rows
+                and len(train_y) < num_train_examples // 2):
+            _coverage_warned = True
+            warnings.warn(f"Failed to collect {num_train_examples:,} context "
+                          f"examples within {num_train_trials:,} candidates. "
+                          f"To improve coverage, consider increasing the "
+                          f"number of PQ iterations using the "
+                          f"'max_pq_iterations' option. This warning will not "
+                          f"be shown again in this run.")
+
+        if (not num_test_examples > 0  #
+                and not _coverage_warned  #
+                and len(entity_df) >= num_entity_rows
+                and len(test_y) < num_test_examples // 2):
+            _coverage_warned = True
+            warnings.warn(f"Failed to collect {num_test_examples:,} test "
+                          f"examples within {num_test_trials:,} candidates. "
+                          f"To improve coverage, consider increasing the "
+                          f"number of PQ iterations using the "
+                          f"'max_pq_iterations' option. This warning will not "
+                          f"be shown again in this run.")
+
+        return (
+            TargetOutput(train_pkey, train_time, train_y),
+            TargetOutput(test_pkey, test_time, test_y),
         )
 
     # Abstract Methods ########################################################
 
+    @abstractmethod
+    def _get_min_max_time_dict(
+        self,
+        table_names: list[str],
+    ) -> dict[str, tuple[pd.Timestamp, pd.Timestamp]]:
+        r"""Returns the minimum and maximum timestamps for a set of tables.
+
+        Args:
+            table_names: The tables.
+        """
+
     @abstractmethod
     def _sample_subgraph(
         self,
         entity_table_name: str,
         entity_pkey: pd.Series,
-        anchor_time: pd.Series,
+        anchor_time: pd.Series | Literal['entity'],
         columns_dict: dict[str, set[str]],
         num_neighbors: list[int],
     ) -> SamplerOutput:
-        pass
+        r"""Samples distinct subgraphs for each entity primary key.
+
+        Args:
+            entity_table_name: The entity table name.
+            entity_pkey: The primary keys to use as seed nodes.
+            anchor_time: The anchor time of the subgraphs.
+            columns_dict: The columns to return for each table.
+            num_neighbors: The number of neighbors to sample for each hop.
+        """
+
+    @abstractmethod
+    def _sample_entity_table(
+        self,
+        table_name: str,
+        columns: set[str],
+        num_rows: int,
+        random_seed: int | None = None,
+    ) -> pd.DataFrame:
+        r"""Returns a random sample of rows from the entity table.
+
+        Args:
+            table_name: The table.
+            columns: The columns to return.
+            num_rows: Maximum number of rows to return. Can be smaller in case
+                the entity table contains less rows.
+            random_seed: A manual seed for generating pseudo-random numbers.
+        """
 
     @abstractmethod
     def _sample_target(
         self,
         query: ValidatedPredictiveQuery,
-        num_examples: int,
-        anchor_time: pd.Timestamp | Literal['entity'],
+        entity_df: pd.DataFrame,
+        train_index: np.ndarray,
+        train_time: pd.Series,
+        num_train_examples: int,
+        test_index: np.ndarray,
+        test_time: pd.Series,
+        num_test_examples: int,
         columns_dict: dict[str, set[str]],
         time_offset_dict: dict[
             tuple[str, str, str],
             tuple[pd.DateOffset | None, pd.DateOffset],
         ],
-        random_seed: int | None = None,
-    ) -> TargetOutput:
-        pass
+    ) -> tuple[pd.Series, np.ndarray, pd.Series, np.ndarray]:
+        r"""Samples ground-truth targets given a predictive query from a set of
+        training and test candidates.
+
+        Args:
+            query: The predictive query.
+            entity_df: The entity data frame, containing the union of all train
+                and test candidates.
+            train_index: The indices of training candidates.
+            train_time: The anchor time of training candidates.
+            num_train_examples: How many training examples to produce.
+            test_index: The indices of test candidates.
+            test_time: The anchor time of test candidates.
+            num_test_examples: How many test examples to produce.
+            columns_dict: The columns that are being used to compute
+                ground-truth targets.
+            time_offset_dict: The date offsets to query for each edge type,
+                relative to the anchor time.
+        """
 
 
 # Helper Functions ############################################################
@@ -372,3 +755,7 @@ def max_date_offset(*args: pd.DateOffset) -> pd.DateOffset:
     assert len(timestamps) > 0
     argmax = max(range(len(timestamps)), key=lambda i: timestamps[i])
     return args[argmax]
+
+
+def to_ser(value: Any, size: int) -> pd.Series:
+    return pd.Series([value]).repeat(size).reset_index(drop=True)

kumoai/experimental/rfm/base/source.py

@@ -6,9 +6,10 @@ from kumoapi.typing import Dtype
 @dataclass
 class SourceColumn:
     name: str
-    dtype: Dtype
+    dtype: Dtype | None
     is_primary_key: bool
     is_unique_key: bool
+    is_nullable: bool
 
 
 @dataclass