kumoai 2.13.0.dev202512040649__cp313-cp313-win_amd64.whl → 2.14.0.dev202601081732__cp313-cp313-win_amd64.whl

kumoai/experimental/rfm/base/sampler.py

@@ -0,0 +1,782 @@
+import copy
+import re
+import warnings
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from dataclasses import dataclass
+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.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]
+    num_sampled_nodes_dict: dict[str, list[int]]
+    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):
+    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)
+            self._edge_types.append(edge_type)
+            self._edge_types.append(Subgraph.rev_edge_type(edge_type))
+
+        # Source Table -> [(Foreign Key, Destination Table)]
+        self._foreign_key_dict: dict[str, list[tuple[str, str]]] = {}
+        # Destination Table -> [(Source Table, Foreign Key)]
+        self._rev_foreign_key_dict: dict[str, list[tuple[str, str]]] = {}
+        for table in graph.tables.values():
+            self._foreign_key_dict[table.name] = []
+            self._rev_foreign_key_dict[table.name] = []
+        for src_table, fkey, dst_table in graph.edges:
+            self._foreign_key_dict[src_table].append((fkey, dst_table))
+            self._rev_foreign_key_dict[dst_table].append((src_table, fkey))
+
+        self._primary_key_dict: dict[str, str] = {
+            table.name: table._primary_key
+            for table in graph.tables.values()
+            if table._primary_key is not None
+        }
+
+        self._time_column_dict: dict[str, str] = {
+            table.name: table._time_column
+            for table in graph.tables.values()
+            if table._time_column is not None
+        }
+
+        self._end_time_column_dict: dict[str, str] = {
+            table.name: table._end_time_column
+            for table in graph.tables.values()
+            if table._end_time_column is not None
+        }
+
+        foreign_keys = {(edge.src_table, edge.fkey) for edge in graph.edges}
+        self._table_stype_dict: dict[str, dict[str, Stype]] = {}
+        for table in graph.tables.values():
+            self._table_stype_dict[table.name] = {}
+            for column in table.columns:
+                if column == table.primary_key:
+                    continue
+                if (table.name, column.name) in foreign_keys:
+                    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 foreign_key_dict(self) -> dict[str, list[tuple[str, str]]]:
+        r"""The foreign keys for all tables in the graph."""
+        return self._foreign_key_dict
+
+    @property
+    def rev_foreign_key_dict(self) -> dict[str, list[tuple[str, str]]]:
+        r"""The foreign key back references for all tables in the graph."""
+        return self._rev_foreign_key_dict
+
+    @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 | 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
+        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():
+                for column_name in exclude_cols:
+                    del table_stype_dict[table_name][column_name]
+
+        # Collect all columns being used as features:
+        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:
+            columns_dict[table_name].add(self.primary_key_dict[table_name])
+
+        if (isinstance(anchor_time, pd.Series)
+                and anchor_time.dtype != 'datetime64[ns]'):
+            anchor_time = anchor_time.astype('datetime64[ns]')
+
+        out = self._sample_subgraph(
+            entity_table_name=entity_table_names[0],
+            entity_pkey=entity_pkey,
+            anchor_time=anchor_time,
+            columns_dict=columns_dict,
+            num_neighbors=num_neighbors,
+        )
+
+        # Parse `SubgraphOutput` into `Subgraph` structure:
+        subgraph = Subgraph(
+            anchor_time=out.anchor_time,
+            table_dict={},
+            link_dict={},
+        )
+
+        for table_name, batch in out.batch_dict.items():
+            if len(batch) == 0:
+                continue
+
+            primary_key: str | None = None
+            if table_name in entity_table_names:
+                primary_key = self.primary_key_dict[table_name]
+
+            df = out.df_dict[table_name].reset_index(drop=True)
+            if end_time_column := self.end_time_column_dict.get(table_name):
+                # 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]
+                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():
+                if stype == Stype.text:
+                    df[column_name] = _normalize_text(df[column_name])
+
+            subgraph.table_dict[table_name] = Table(
+                df=df,
+                row=out.inverse_dict.get(table_name),
+                batch=batch,
+                num_sampled_nodes=out.num_sampled_nodes_dict[table_name],
+                stype_dict=stype_dict,
+                primary_key=primary_key,
+            )
+
+        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]
+
+            if row is None or col is None or len(row) == 0:
+                continue
+
+            # Do not store reverse edge type if it is an exact replica:
+            rev_edge_type = Subgraph.rev_edge_type(edge_type)
+            if (rev_edge_type in subgraph.link_dict
+                    and np.array_equal(row, out.col_dict[rev_edge_type])
+                    and np.array_equal(col, out.row_dict[rev_edge_type])):
+                subgraph.link_dict[edge_type] = Link(
+                    layout=EdgeLayout.REV,
+                    row=None,
+                    col=None,
+                    num_sampled_edges=out.num_sampled_edges_dict[edge_type],
+                )
+                continue
+
+            # Do not store non-informative edges:
+            layout = EdgeLayout.COO
+            if np.array_equal(row, np.arange(len(row))):
+                row = None
+            if np.array_equal(col, np.arange(len(col))):
+                col = None
+
+            # Store in compressed representation if more efficient:
+            num_cols = subgraph.table_dict[edge_type[2]].num_rows
+            if col is not None and len(col) > num_cols + 1:
+                layout = EdgeLayout.CSC
+                colcount = np.bincount(col, minlength=num_cols)
+                col = np.empty(num_cols + 1, dtype=col.dtype)
+                col[0] = 0
+                np.cumsum(colcount, out=col[1:])
+
+            subgraph.link_dict[edge_type] = Link(
+                layout=layout,
+                row=row,
+                col=col,
+                num_sampled_edges=out.num_sampled_edges_dict[edge_type],
+            )
+
+        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 _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 | Literal['entity'],
+        columns_dict: dict[str, set[str]],
+        num_neighbors: list[int],
+    ) -> SamplerOutput:
+        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 ############################################################
+
+PUNCTUATION = re.compile(r"[\'\"\.,\(\)\!\?\;\:]")
+MULTISPACE = re.compile(r"\s+")
+
+
+def _normalize_text(
+    ser: pd.Series,
+    max_words: int | None = 50,
+) -> pd.Series:
+    r"""Normalizes text into a list of lower-case words.
+
+    Args:
+        ser: The :class:`pandas.Series` to normalize.
+        max_words: The maximum number of words to return.
+            This will auto-shrink any large text column to avoid blowing up
+            context size.
+    """
+    if len(ser) == 0 or pd.api.types.is_list_like(ser.iloc[0]):
+        return ser
+
+    def normalize_fn(line: str) -> list[str]:
+        line = PUNCTUATION.sub(" ", line)
+        line = re.sub(r"<br\s*/?>", " ", line)  # Handle <br /> or <br>
+        line = MULTISPACE.sub(" ", line)
+        words = line.split()
+        if max_words is not None:
+            words = words[:max_words]
+        return words
+
+    ser = ser.fillna('').astype(str)
+
+    if max_words is not None:
+        # We estimate the number of words as 5 characters + 1 space in an
+        # English text on average. We need this pre-filter here, as word
+        # splitting on a giant text can be very expensive:
+        ser = ser.str[:6 * max_words]
+
+    ser = ser.str.lower()
+    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)