kumoai 2.13.0.dev202512061731__cp313-cp313-macosx_11_0_arm64.whl → 2.14.0.dev202512151351__cp313-cp313-macosx_11_0_arm64.whl

kumoai/experimental/rfm/base/sampler.py

@@ -1,45 +1,58 @@
 import copy
 import re
+import warnings
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple
 
 import numpy as np
 import pandas as pd
+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.experimental.rfm.base import SourceColumn
+from kumoai.utils import ProgressLogger
+
 if TYPE_CHECKING:
     from kumoai.experimental.rfm import Graph
 
-
-@dataclass
-class EdgeSpec:
-    num_neighbors: int | None = None
-    time_offsets: tuple[
-        pd.DateOffset | None,
-        pd.DateOffset,
-    ] | None = None
-
-    def __post_init__(self) -> None:
-        if (self.num_neighbors is None) == (self.time_offsets is None):
-            raise ValueError("Only one of 'num_neighbors' and 'time_offsets' "
-                             "must be provided")
+_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]
     num_sampled_nodes_dict: dict[str, list[int]]
-    row_dict: dict[tuple[str, str, str], np.ndarray] | None = None
-    col_dict: dict[tuple[str, str, str], np.ndarray] | None = None
-    num_sampled_edges_dict: dict[tuple[str, str, str], list[int]] | None = None
+    row_dict: dict[tuple[str, str, str], np.ndarray]
+    col_dict: dict[tuple[str, str, str], np.ndarray]
+    num_sampled_edges_dict: dict[tuple[str, str, str], list[int]]
+
+
+class TargetOutput(NamedTuple):
+    entity_pkey: pd.Series
+    anchor_time: pd.Series
+    target: pd.Series
 
 
 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)
@@ -75,40 +88,106 @@ class Sampler(ABC):
                     continue
                 self._table_stype_dict[table.name][column.name] = column.stype
 
+        self._source_table_dict: dict[str, dict[str, SourceColumn]] = {}
+        for table in graph.tables.values():
+            self._source_table_dict[table.name] = table._source_column_dict
+
+        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 columns in all tables in
+        the graph.
+        """
         return self._table_stype_dict
 
+    @property
+    def source_table_dict(self) -> dict[str, dict[str, SourceColumn]]:
+        r"""Source column information for all tables in the graph."""
+        return self._source_table_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:
-
-        edge_spec_dict: dict[tuple[str, str, str], list[EdgeSpec]] = {
-            edge_type: [EdgeSpec(value) for value in num_neighbors]
-            for edge_type in self.edge_types
-        }
-
+        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
         if exclude_cols_dict is not None:
@@ -118,30 +197,29 @@ class Sampler(ABC):
                     del table_stype_dict[table_name][column_name]
 
         # Collect all columns being used as features:
-        column_spec_dict: dict[str, list[str]] = {
-            table_name: list(stype_dict.keys())
+        columns_dict: dict[str, set[str]] = {
+            table_name: set(stype_dict.keys())
             for table_name, stype_dict in table_stype_dict.items()
         }
         # Make sure to store primary key information for entity tables:
         for table_name in entity_table_names:
-            column_spec_dict[table_name] = (
-                [self.primary_key_dict[table_name]] +
-                column_spec_dict[table_name])
+            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(
+
+        out = self._sample_subgraph(
             entity_table_name=entity_table_names[0],
             entity_pkey=entity_pkey,
             anchor_time=anchor_time,
-            column_spec_dict=column_spec_dict,
-            edge_spec_dict=edge_spec_dict,
-            drop_duplicates=True,
-            return_edges=True,
+            columns_dict=columns_dict,
+            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={},
         )
@@ -150,18 +228,18 @@ class Sampler(ABC):
             if len(batch) == 0:
                 continue
 
-            primary_key = None
+            primary_key: str | None = None
             if table_name in entity_table_names:
-                primary_key = self.primary_key_dict.get(table_name, None)
+                primary_key = self.primary_key_dict[table_name]
 
             df = out.df_dict[table_name].reset_index(drop=True)
-            if table_name in self.end_time_column_dict:
+            if end_time_column := self.end_time_column_dict.get(table_name):
                 # Set end time to NaT for all values greater than anchor time:
-                end_time_column = self.end_time_column_dict[table_name]
+                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]
+                mask = ser.astype(int).to_numpy() > out.anchor_time[batch]
                 ser.iloc[mask] = pd.NaT
                 df[end_time_column] = ser
 
@@ -179,9 +257,6 @@ class Sampler(ABC):
                 primary_key=primary_key,
             )
 
-        assert out.row_dict is not None
-        assert out.col_dict is not None
-        assert out.num_sampled_edges_dict is not None
         for edge_type in out.row_dict.keys():
             row: np.ndarray | None = out.row_dict[edge_type]
             col: np.ndarray | None = out.col_dict[edge_type]
@@ -227,20 +302,408 @@ class Sampler(ABC):
 
         return subgraph
 
+    # Predictive Query ########################################################
+
+    def _get_query_columns_dict(
+        self,
+        query: ValidatedPredictiveQuery,
+    ) -> 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],
+        ] = {}
+
+        def _add_time_offset(node: ASTNode, num_forecasts: int = 1) -> None:
+            if isinstance(node, Aggregation):
+                table_name = node._get_target_column_name().split('.')[0]
+
+                edge_types = [
+                    edge_type for edge_type in self.edge_types
+                    if edge_type[0] == table_name
+                    and edge_type[2] == query.entity_table
+                ]
+                if len(edge_types) != 1:
+                    raise ValueError(f"Could not find a unique foreign key "
+                                     f"from table '{table_name}' to "
+                                     f"'{query.entity_table}'")
+                if edge_types[0] not in time_offset_dict:
+                    start = node.aggr_time_range.start_date_offset
+                    end = node.aggr_time_range.end_date_offset * num_forecasts
+                else:
+                    start, end = time_offset_dict[edge_types[0]]
+                    start = min_date_offset(
+                        start,
+                        node.aggr_time_range.start_date_offset,
+                    )
+                    end = max_date_offset(
+                        end,
+                        node.aggr_time_range.end_date_offset * num_forecasts,
+                    )
+                time_offset_dict[edge_types[0]] = (start, end)
+
+            for child in node.children:
+                _add_time_offset(child, num_forecasts)
+
+        _add_time_offset(query.target_ast, query.num_forecasts)
+        _add_time_offset(query.entity_ast)
+        if query.whatif_ast is not None:
+            _add_time_offset(query.whatif_ast)
+
+        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,
+            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,
+        )
+
+        # 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 sample(
+    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,
-        column_spec_dict: dict[str, list[str]],
-        edge_spec_dict: dict[tuple[str, str, str], list[EdgeSpec]],
-        drop_duplicates: bool = False,
-        return_edges: bool = False,
+        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,
+        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],
+        ],
+    ) -> 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 ############################################################
@@ -285,3 +748,26 @@ def _normalize_text(
     ser = ser.map(normalize_fn)
 
     return ser
+
+
+def min_date_offset(*args: pd.DateOffset | None) -> pd.DateOffset | None:
+    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: pd.DateOffset) -> pd.DateOffset:
+    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]
+
+
+def to_ser(value: Any, size: int) -> pd.Series:
+    return pd.Series([value]).repeat(size).reset_index(drop=True)