kumoai 2.14.0.dev202601011731__cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

kumoai/experimental/rfm/rfm.py

@@ -0,0 +1,1184 @@
+import json
+import time
+import warnings
+from collections import defaultdict
+from collections.abc import Generator, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass, replace
+from typing import Any, Literal, overload
+
+import numpy as np
+import pandas as pd
+from kumoapi.model_plan import RunMode
+from kumoapi.pquery import QueryType, ValidatedPredictiveQuery
+from kumoapi.pquery.AST import (
+    Aggregation,
+    Column,
+    Condition,
+    Join,
+    LogicalOperation,
+)
+from kumoapi.rfm import Context
+from kumoapi.rfm import Explanation as ExplanationConfig
+from kumoapi.rfm import (
+    RFMEvaluateRequest,
+    RFMParseQueryRequest,
+    RFMPredictRequest,
+)
+from kumoapi.task import TaskType
+from kumoapi.typing import AggregationType, Stype
+
+from kumoai.client.rfm import RFMAPI
+from kumoai.exceptions import HTTPException
+from kumoai.experimental.rfm import Graph
+from kumoai.experimental.rfm.base import DataBackend, Sampler
+from kumoai.mixin import CastMixin
+from kumoai.utils import ProgressLogger, display
+
+_RANDOM_SEED = 42
+
+_MAX_PRED_SIZE: dict[TaskType, int] = defaultdict(lambda: 1_000)
+_MAX_PRED_SIZE[TaskType.TEMPORAL_LINK_PREDICTION] = 200
+
+_MAX_CONTEXT_SIZE = {
+    RunMode.DEBUG: 100,
+    RunMode.FAST: 1_000,
+    RunMode.NORMAL: 5_000,
+    RunMode.BEST: 10_000,
+}
+_MAX_TEST_SIZE = {  # Share test set size across run modes for fair comparison:
+    RunMode.DEBUG: 100,
+    RunMode.FAST: 2_000,
+    RunMode.NORMAL: 2_000,
+    RunMode.BEST: 2_000,
+}
+
+_MAX_SIZE = 30 * 1024 * 1024
+_SIZE_LIMIT_MSG = ("Context size exceeds the 30MB limit. {stats}\nPlease "
+                   "reduce either the number of tables in the graph, their "
+                   "number of columns (e.g., large text columns), "
+                   "neighborhood configuration, or the run mode. If none of "
+                   "this is possible, please create a feature request at "
+                   "'https://github.com/kumo-ai/kumo-rfm' if you must go "
+                   "beyond this for your use-case.")
+
+
+@dataclass(repr=False)
+class ExplainConfig(CastMixin):
+    """Configuration for explainability.
+
+    Args:
+        skip_summary: Whether to skip generating a human-readable summary of
+            the explanation.
+    """
+    skip_summary: bool = False
+
+
+@dataclass(repr=False)
+class Explanation:
+    prediction: pd.DataFrame
+    summary: str
+    details: ExplanationConfig
+
+    @overload
+    def __getitem__(self, index: Literal[0]) -> pd.DataFrame:
+        pass
+
+    @overload
+    def __getitem__(self, index: Literal[1]) -> str:
+        pass
+
+    def __getitem__(self, index: int) -> pd.DataFrame | str:
+        if index == 0:
+            return self.prediction
+        if index == 1:
+            return self.summary
+        raise IndexError("Index out of range")
+
+    def __iter__(self) -> Iterator[pd.DataFrame | str]:
+        return iter((self.prediction, self.summary))
+
+    def __repr__(self) -> str:
+        return str((self.prediction, self.summary))
+
+    def print(self) -> None:
+        r"""Prints the explanation."""
+        display.dataframe(self.prediction)
+        display.message(self.summary)
+
+    def _ipython_display_(self) -> None:
+        self.print()
+
+
+class KumoRFM:
+    r"""The Kumo Relational Foundation model (RFM) from the `KumoRFM: A
+    Foundation Model for In-Context Learning on Relational Data
+    <https://kumo.ai/research/kumo_relational_foundation_model.pdf>`_ paper.
+
+    :class:`KumoRFM` is a foundation model to generate predictions for any
+    relational dataset without training.
+    The model is pre-trained and the class provides an interface to query the
+    model from a :class:`Graph` object.
+
+    .. code-block:: python
+
+        from kumoai.experimental.rfm import Graph, KumoRFM
+
+        df_users = pd.DataFrame(...)
+        df_items = pd.DataFrame(...)
+        df_orders = pd.DataFrame(...)
+
+        graph = Graph.from_data({
+            'users': df_users,
+            'items': df_items,
+            'orders': df_orders,
+        })
+
+        rfm = KumoRFM(graph)
+
+        query = ("PREDICT COUNT(orders.*, 0, 30, days)>0 "
+                 "FOR users.user_id=1")
+        result = rfm.predict(query)
+
+        print(result)  # user_id  COUNT(transactions.*, 0, 30, days) > 0
+                       # 1        0.85
+
+    Args:
+        graph: The graph.
+        verbose: Whether to print verbose output.
+        optimize: If set to ``True``, will optimize the underlying data backend
+            for optimal querying. For example, for transactional database
+            backends, will create any missing indices. Requires write-access to
+            the data backend.
+    """
+    def __init__(
+        self,
+        graph: Graph,
+        verbose: bool | ProgressLogger = True,
+        optimize: bool = False,
+    ) -> None:
+        graph = graph.validate()
+        self._graph_def = graph._to_api_graph_definition()
+
+        if graph.backend == DataBackend.LOCAL:
+            from kumoai.experimental.rfm.backend.local import LocalSampler
+            self._sampler: Sampler = LocalSampler(graph, verbose)
+        elif graph.backend == DataBackend.SQLITE:
+            from kumoai.experimental.rfm.backend.sqlite import SQLiteSampler
+            self._sampler = SQLiteSampler(graph, verbose, optimize)
+        elif graph.backend == DataBackend.SNOWFLAKE:
+            from kumoai.experimental.rfm.backend.snow import SnowSampler
+            self._sampler = SnowSampler(graph, verbose)
+        else:
+            raise NotImplementedError
+
+        self._client: RFMAPI | None = None
+
+        self._batch_size: int | Literal['max'] | None = None
+        self.num_retries: int = 0
+
+    @property
+    def _api_client(self) -> RFMAPI:
+        if self._client is not None:
+            return self._client
+
+        from kumoai.experimental.rfm import global_state
+        self._client = RFMAPI(global_state.client)
+        return self._client
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}()'
+
+    @contextmanager
+    def batch_mode(
+        self,
+        batch_size: int | Literal['max'] = 'max',
+        num_retries: int = 1,
+    ) -> Generator[None, None, None]:
+        """Context manager to predict in batches.
+
+        .. code-block:: python
+
+            with model.batch_mode(batch_size='max', num_retries=1):
+                df = model.predict(query, indices=...)
+
+        Args:
+            batch_size: The batch size. If set to ``"max"``, will use the
+                maximum applicable batch size for the given task.
+            num_retries: The maximum number of retries for failed queries due
+                to unexpected server issues.
+        """
+        if batch_size != 'max' and batch_size <= 0:
+            raise ValueError(f"'batch_size' must be greater than zero "
+                             f"(got {batch_size})")
+
+        if num_retries < 0:
+            raise ValueError(f"'num_retries' must be greater than or equal to "
+                             f"zero (got {num_retries})")
+
+        self._batch_size = batch_size
+        self.num_retries = num_retries
+        yield
+        self._batch_size = None
+        self.num_retries = 0
+
+    @overload
+    def predict(
+        self,
+        query: str,
+        indices: list[str] | list[float] | list[int] | None = None,
+        *,
+        explain: Literal[False] = False,
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
+        context_anchor_time: pd.Timestamp | None = None,
+        run_mode: RunMode | str = RunMode.FAST,
+        num_neighbors: list[int] | None = None,
+        num_hops: int = 2,
+        max_pq_iterations: int = 10,
+        random_seed: int | None = _RANDOM_SEED,
+        verbose: bool | ProgressLogger = True,
+        use_prediction_time: bool = False,
+    ) -> pd.DataFrame:
+        pass
+
+    @overload
+    def predict(
+        self,
+        query: str,
+        indices: list[str] | list[float] | list[int] | None = None,
+        *,
+        explain: Literal[True] | ExplainConfig | dict[str, Any],
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
+        context_anchor_time: pd.Timestamp | None = None,
+        run_mode: RunMode | str = RunMode.FAST,
+        num_neighbors: list[int] | None = None,
+        num_hops: int = 2,
+        max_pq_iterations: int = 10,
+        random_seed: int | None = _RANDOM_SEED,
+        verbose: bool | ProgressLogger = True,
+        use_prediction_time: bool = False,
+    ) -> Explanation:
+        pass
+
+    def predict(
+        self,
+        query: str,
+        indices: list[str] | list[float] | list[int] | None = None,
+        *,
+        explain: bool | ExplainConfig | dict[str, Any] = False,
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
+        context_anchor_time: pd.Timestamp | None = None,
+        run_mode: RunMode | str = RunMode.FAST,
+        num_neighbors: list[int] | None = None,
+        num_hops: int = 2,
+        max_pq_iterations: int = 10,
+        random_seed: int | None = _RANDOM_SEED,
+        verbose: bool | ProgressLogger = True,
+        use_prediction_time: bool = False,
+    ) -> pd.DataFrame | Explanation:
+        """Returns predictions for a predictive query.
+
+        Args:
+            query: The predictive query.
+            indices: The entity primary keys to predict for. Will override the
+                indices given as part of the predictive query. Predictions will
+                be generated for all indices, independent of whether they
+                fulfill entity filter constraints. To pre-filter entities, use
+                :meth:`~KumoRFM.is_valid_entity`.
+            explain: Configuration for explainability.
+                If set to ``True``, will additionally explain the prediction.
+                Passing in an :class:`ExplainConfig` instance provides control
+                over which parts of explanation are generated.
+                Explainability is currently only supported for single entity
+                predictions with ``run_mode="FAST"``.
+            anchor_time: The anchor timestamp for the prediction. If set to
+                ``None``, will use the maximum timestamp in the data.
+                If set to ``"entity"``, will use the timestamp of the entity.
+            context_anchor_time: The maximum anchor timestamp for context
+                examples. If set to ``None``, ``anchor_time`` will
+                determine the anchor time for context examples.
+            run_mode: The :class:`RunMode` for the query.
+            num_neighbors: The number of neighbors to sample for each hop.
+                If specified, the ``num_hops`` option will be ignored.
+            num_hops: The number of hops to sample when generating the context.
+            max_pq_iterations: The maximum number of iterations to perform to
+                collect valid labels. It is advised to increase the number of
+                iterations in case the predictive query has strict entity
+                filters, in which case, :class:`KumoRFM` needs to sample more
+                entities to find valid labels.
+            random_seed: A manual seed for generating pseudo-random numbers.
+            verbose: Whether to print verbose output.
+            use_prediction_time: Whether to use the anchor timestamp as an
+                additional feature during prediction. This is typically
+                beneficial for time series forecasting tasks.
+
+        Returns:
+            The predictions as a :class:`pandas.DataFrame`.
+            If ``explain`` is provided, returns an :class:`Explanation` object
+            containing the prediction, summary, and details.
+        """
+        explain_config: ExplainConfig | None = None
+        if explain is True:
+            explain_config = ExplainConfig()
+        elif explain is not False:
+            explain_config = ExplainConfig._cast(explain)
+
+        query_def = self._parse_query(query)
+        query_str = query_def.to_string()
+
+        if num_hops != 2 and num_neighbors is not None:
+            warnings.warn(f"Received custom 'num_neighbors' option; ignoring "
+                          f"custom 'num_hops={num_hops}' option")
+
+        if explain_config is not None and run_mode in {
+                RunMode.NORMAL, RunMode.BEST
+        }:
+            warnings.warn(f"Explainability is currently only supported for "
+                          f"run mode 'FAST' (got '{run_mode}'). Provided run "
+                          f"mode has been reset. Please lower the run mode to "
+                          f"suppress this warning.")
+
+        if indices is None:
+            if query_def.rfm_entity_ids is None:
+                raise ValueError("Cannot find entities to predict for. Please "
+                                 "pass them via `predict(query, indices=...)`")
+            indices = query_def.get_rfm_entity_id_list()
+        else:
+            query_def = replace(query_def, rfm_entity_ids=None)
+
+        if len(indices) == 0:
+            raise ValueError("At least one entity is required")
+
+        if explain_config is not None and len(indices) > 1:
+            raise ValueError(
+                f"Cannot explain predictions for more than a single entity "
+                f"(got {len(indices)})")
+
+        query_repr = query_def.to_string(rich=True, exclude_predict=True)
+        if explain_config is not None:
+            msg = f'[bold]EXPLAIN[/bold] {query_repr}'
+        else:
+            msg = f'[bold]PREDICT[/bold] {query_repr}'
+
+        if not isinstance(verbose, ProgressLogger):
+            verbose = ProgressLogger.default(msg=msg, verbose=verbose)
+
+        with verbose as logger:
+
+            batch_size: int | None = None
+            if self._batch_size == 'max':
+                task_type = self._get_task_type(
+                    query=query_def,
+                    edge_types=self._sampler.edge_types,
+                )
+                batch_size = _MAX_PRED_SIZE[task_type]
+            else:
+                batch_size = self._batch_size
+
+            if batch_size is not None:
+                offsets = range(0, len(indices), batch_size)
+                batches = [indices[step:step + batch_size] for step in offsets]
+            else:
+                batches = [indices]
+
+            if len(batches) > 1:
+                logger.log(f"Splitting {len(indices):,} entities into "
+                           f"{len(batches):,} batches of size {batch_size:,}")
+
+            predictions: list[pd.DataFrame] = []
+            summary: str | None = None
+            details: Explanation | None = None
+            for i, batch in enumerate(batches):
+                # TODO Re-use the context for subsequent predictions.
+                context = self._get_context(
+                    query=query_def,
+                    indices=batch,
+                    anchor_time=anchor_time,
+                    context_anchor_time=context_anchor_time,
+                    run_mode=RunMode(run_mode),
+                    num_neighbors=num_neighbors,
+                    num_hops=num_hops,
+                    max_pq_iterations=max_pq_iterations,
+                    evaluate=False,
+                    random_seed=random_seed,
+                    logger=logger if i == 0 else None,
+                )
+                request = RFMPredictRequest(
+                    context=context,
+                    run_mode=RunMode(run_mode),
+                    query=query_str,
+                    use_prediction_time=use_prediction_time,
+                )
+                with warnings.catch_warnings():
+                    warnings.filterwarnings('ignore', message='gencode')
+                    request_msg = request.to_protobuf()
+                    _bytes = request_msg.SerializeToString()
+                if i == 0:
+                    logger.log(f"Generated context of size "
+                               f"{len(_bytes) / (1024*1024):.2f}MB")
+
+                if len(_bytes) > _MAX_SIZE:
+                    stats = Context.get_memory_stats(request_msg.context)
+                    raise ValueError(_SIZE_LIMIT_MSG.format(stats=stats))
+
+                if i == 0 and len(batches) > 1:
+                    verbose.init_progress(
+                        total=len(batches),
+                        description='Predicting',
+                    )
+
+                for attempt in range(self.num_retries + 1):
+                    try:
+                        if explain_config is not None:
+                            resp = self._api_client.explain(
+                                request=_bytes,
+                                skip_summary=explain_config.skip_summary,
+                            )
+                            summary = resp.summary
+                            details = resp.details
+                        else:
+                            resp = self._api_client.predict(_bytes)
+                        df = pd.DataFrame(**resp.prediction)
+
+                        # Cast 'ENTITY' to correct data type:
+                        if 'ENTITY' in df:
+                            table_dict = context.subgraph.table_dict
+                            table = table_dict[query_def.entity_table]
+                            ser = table.df[table.primary_key]
+                            df['ENTITY'] = df['ENTITY'].astype(ser.dtype)
+
+                        # Cast 'ANCHOR_TIMESTAMP' to correct data type:
+                        if 'ANCHOR_TIMESTAMP' in df:
+                            ser = df['ANCHOR_TIMESTAMP']
+                            if not pd.api.types.is_datetime64_any_dtype(ser):
+                                if isinstance(ser.iloc[0], str):
+                                    unit = None
+                                else:
+                                    unit = 'ms'
+                                df['ANCHOR_TIMESTAMP'] = pd.to_datetime(
+                                    ser, errors='coerce', unit=unit)
+
+                        predictions.append(df)
+
+                        if len(batches) > 1:
+                            verbose.step()
+
+                        break
+                    except HTTPException as e:
+                        if attempt == self.num_retries:
+                            try:
+                                msg = json.loads(e.detail)['detail']
+                            except Exception:
+                                msg = e.detail
+                            raise RuntimeError(
+                                f"An unexpected exception occurred. Please "
+                                f"create an issue at "
+                                f"'https://github.com/kumo-ai/kumo-rfm'. {msg}"
+                            ) from None
+
+                        time.sleep(2**attempt)  # 1s, 2s, 4s, 8s, ...
+
+        if len(predictions) == 1:
+            prediction = predictions[0]
+        else:
+            prediction = pd.concat(predictions, ignore_index=True)
+
+        if explain_config is not None:
+            assert len(predictions) == 1
+            assert summary is not None
+            assert details is not None
+            return Explanation(
+                prediction=prediction,
+                summary=summary,
+                details=details,
+            )
+
+        return prediction
+
+    def is_valid_entity(
+        self,
+        query: str,
+        indices: list[str] | list[float] | list[int] | None = None,
+        *,
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
+    ) -> np.ndarray:
+        r"""Returns a mask that denotes which entities are valid for the
+        given predictive query, *i.e.*, which entities fulfill (temporal)
+        entity filter constraints.
+
+        Args:
+            query: The predictive query.
+            indices: The entity primary keys to predict for. Will override the
+                indices given as part of the predictive query.
+            anchor_time: The anchor timestamp for the prediction. If set to
+                ``None``, will use the maximum timestamp in the data.
+                If set to ``"entity"``, will use the timestamp of the entity.
+        """
+        query_def = self._parse_query(query)
+
+        if indices is None:
+            if query_def.rfm_entity_ids is None:
+                raise ValueError("Cannot find entities to predict for. Please "
+                                 "pass them via "
+                                 "`is_valid_entity(query, indices=...)`")
+            indices = query_def.get_rfm_entity_id_list()
+
+        if len(indices) == 0:
+            raise ValueError("At least one entity is required")
+
+        if anchor_time is None:
+            anchor_time = self._get_default_anchor_time(query_def)
+
+        if isinstance(anchor_time, pd.Timestamp):
+            self._validate_time(query_def, anchor_time, None, False)
+        else:
+            assert anchor_time == 'entity'
+            if query_def.entity_table not in self._sampler.time_column_dict:
+                raise ValueError(f"Anchor time 'entity' requires the entity "
+                                 f"table '{query_def.entity_table}' "
+                                 f"to have a time column.")
+
+        raise NotImplementedError
+
+    def evaluate(
+        self,
+        query: str,
+        *,
+        metrics: list[str] | None = None,
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
+        context_anchor_time: pd.Timestamp | None = None,
+        run_mode: RunMode | str = RunMode.FAST,
+        num_neighbors: list[int] | None = None,
+        num_hops: int = 2,
+        max_pq_iterations: int = 10,
+        random_seed: int | None = _RANDOM_SEED,
+        verbose: bool | ProgressLogger = True,
+        use_prediction_time: bool = False,
+    ) -> pd.DataFrame:
+        """Evaluates a predictive query.
+
+        Args:
+            query: The predictive query.
+            metrics: The metrics to use.
+            anchor_time: The anchor timestamp for the prediction. If set to
+                ``None``, will use the maximum timestamp in the data.
+                If set to ``"entity"``, will use the timestamp of the entity.
+            context_anchor_time: The maximum anchor timestamp for context
+                examples. If set to ``None``, ``anchor_time`` will
+                determine the anchor time for context examples.
+            run_mode: The :class:`RunMode` for the query.
+            num_neighbors: The number of neighbors to sample for each hop.
+                If specified, the ``num_hops`` option will be ignored.
+            num_hops: The number of hops to sample when generating the context.
+            max_pq_iterations: The maximum number of iterations to perform to
+                collect valid labels. It is advised to increase the number of
+                iterations in case the predictive query has strict entity
+                filters, in which case, :class:`KumoRFM` needs to sample more
+                entities to find valid labels.
+            random_seed: A manual seed for generating pseudo-random numbers.
+            verbose: Whether to print verbose output.
+            use_prediction_time: Whether to use the anchor timestamp as an
+                additional feature during prediction. This is typically
+                beneficial for time series forecasting tasks.
+
+        Returns:
+            The metrics as a :class:`pandas.DataFrame`
+        """
+        query_def = self._parse_query(query)
+
+        if num_hops != 2 and num_neighbors is not None:
+            warnings.warn(f"Received custom 'num_neighbors' option; ignoring "
+                          f"custom 'num_hops={num_hops}' option")
+
+        if query_def.rfm_entity_ids is not None:
+            query_def = replace(
+                query_def,
+                rfm_entity_ids=None,
+            )
+
+        query_repr = query_def.to_string(rich=True, exclude_predict=True)
+        msg = f'[bold]EVALUATE[/bold] {query_repr}'
+
+        if not isinstance(verbose, ProgressLogger):
+            verbose = ProgressLogger.default(msg=msg, verbose=verbose)
+
+        with verbose as logger:
+            context = self._get_context(
+                query=query_def,
+                indices=None,
+                anchor_time=anchor_time,
+                context_anchor_time=context_anchor_time,
+                run_mode=RunMode(run_mode),
+                num_neighbors=num_neighbors,
+                num_hops=num_hops,
+                max_pq_iterations=max_pq_iterations,
+                evaluate=True,
+                random_seed=random_seed,
+                logger=logger if verbose else None,
+            )
+            if metrics is not None and len(metrics) > 0:
+                self._validate_metrics(metrics, context.task_type)
+                metrics = list(dict.fromkeys(metrics))
+            request = RFMEvaluateRequest(
+                context=context,
+                run_mode=RunMode(run_mode),
+                metrics=metrics,
+                use_prediction_time=use_prediction_time,
+            )
+            with warnings.catch_warnings():
+                warnings.filterwarnings('ignore', message='Protobuf gencode')
+                request_msg = request.to_protobuf()
+                request_bytes = request_msg.SerializeToString()
+            logger.log(f"Generated context of size "
+                       f"{len(request_bytes) / (1024*1024):.2f}MB")
+
+            if len(request_bytes) > _MAX_SIZE:
+                stats_msg = Context.get_memory_stats(request_msg.context)
+                raise ValueError(_SIZE_LIMIT_MSG.format(stats=stats_msg))
+
+            try:
+                resp = self._api_client.evaluate(request_bytes)
+            except HTTPException as e:
+                try:
+                    msg = json.loads(e.detail)['detail']
+                except Exception:
+                    msg = e.detail
+                raise RuntimeError(f"An unexpected exception occurred. "
+                                   f"Please create an issue at "
+                                   f"'https://github.com/kumo-ai/kumo-rfm'. "
+                                   f"{msg}") from None
+
+        return pd.DataFrame.from_dict(
+            resp.metrics,
+            orient='index',
+            columns=['value'],
+        ).reset_index(names='metric')
+
+    def get_train_table(
+        self,
+        query: str,
+        size: int,
+        *,
+        anchor_time: pd.Timestamp | Literal['entity'] | None = None,
+        random_seed: int | None = _RANDOM_SEED,
+        max_iterations: int = 10,
+    ) -> pd.DataFrame:
+        """Returns the labels of a predictive query for a specified anchor
+        time.
+
+        Args:
+            query: The predictive query.
+            size: The maximum number of entities to generate labels for.
+            anchor_time: The anchor timestamp for the query. If set to
+                :obj:`None`, will use the maximum timestamp in the data.
+                If set to :`"entity"`, will use the timestamp of the entity.
+            random_seed: A manual seed for generating pseudo-random numbers.
+            max_iterations: The number of steps to run before aborting.
+
+        Returns:
+            The labels as a :class:`pandas.DataFrame`.
+        """
+        query_def = self._parse_query(query)
+
+        if anchor_time is None:
+            anchor_time = self._get_default_anchor_time(query_def)
+            if query_def.target_ast.date_offset_range is not None:
+                offset = query_def.target_ast.date_offset_range.end_date_offset
+                offset *= query_def.num_forecasts
+                anchor_time -= offset
+
+        assert anchor_time is not None
+        if isinstance(anchor_time, pd.Timestamp):
+            self._validate_time(query_def, anchor_time, None, evaluate=True)
+        else:
+            assert anchor_time == 'entity'
+            if query_def.entity_table not in self._sampler.time_column_dict:
+                raise ValueError(f"Anchor time 'entity' requires the entity "
+                                 f"table '{query_def.entity_table}' "
+                                 f"to have a time column")
+
+        train, test = self._sampler.sample_target(
+            query=query_def,
+            num_train_examples=0,
+            train_anchor_time=anchor_time,
+            num_train_trials=0,
+            num_test_examples=size,
+            test_anchor_time=anchor_time,
+            num_test_trials=max_iterations * size,
+            random_seed=random_seed,
+        )
+
+        return pd.DataFrame({
+            'ENTITY': test.entity_pkey,
+            'ANCHOR_TIMESTAMP': test.anchor_time,
+            'TARGET': test.target,
+        })
+
+    # Helpers #################################################################
+
+    def _parse_query(self, query: str) -> ValidatedPredictiveQuery:
+        if isinstance(query, ValidatedPredictiveQuery):
+            return query
+
+        if isinstance(query, str) and query.strip()[:9].lower() == 'evaluate ':
+            raise ValueError("'EVALUATE PREDICT ...' queries are not "
+                             "supported in the SDK. Instead, use either "
+                             "`predict()` or `evaluate()` methods to perform "
+                             "predictions or evaluations.")
+
+        try:
+            request = RFMParseQueryRequest(
+                query=query,
+                graph_definition=self._graph_def,
+            )
+
+            resp = self._api_client.parse_query(request)
+
+            if len(resp.validation_response.warnings) > 0:
+                msg = '\n'.join([
+                    f'{i+1}. {warning.title}: {warning.message}' for i, warning
+                    in enumerate(resp.validation_response.warnings)
+                ])
+                warnings.warn(f"Encountered the following warnings during "
+                              f"parsing:\n{msg}")
+
+            return resp.query
+        except HTTPException as e:
+            try:
+                msg = json.loads(e.detail)['detail']
+            except Exception:
+                msg = e.detail
+            raise ValueError(f"Failed to parse query '{query}'. "
+                             f"{msg}") from None
+
+    @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 _get_default_anchor_time(
+        self,
+        query: ValidatedPredictiveQuery,
+    ) -> pd.Timestamp:
+        if query.query_type == QueryType.TEMPORAL:
+            aggr_table_names = [
+                aggr._get_target_column_name().split('.')[0]
+                for aggr in query.get_all_target_aggregations()
+            ]
+            return self._sampler.get_max_time(aggr_table_names)
+
+        assert query.query_type == QueryType.STATIC
+        return self._sampler.get_max_time()
+
+    def _validate_time(
+        self,
+        query: ValidatedPredictiveQuery,
+        anchor_time: pd.Timestamp,
+        context_anchor_time: pd.Timestamp | None,
+        evaluate: bool,
+    ) -> None:
+
+        if len(self._sampler.time_column_dict) == 0:
+            return  # Graph without timestamps
+
+        min_time = self._sampler.get_min_time()
+        max_time = self._sampler.get_max_time()
+
+        if anchor_time < min_time:
+            raise ValueError(f"Anchor timestamp '{anchor_time}' is before "
+                             f"the earliest timestamp '{min_time}' in the "
+                             f"data.")
+
+        if context_anchor_time is not None and context_anchor_time < min_time:
+            raise ValueError(f"Context anchor timestamp is too early or "
+                             f"aggregation time range is too large. To make "
+                             f"this prediction, we would need data back to "
+                             f"'{context_anchor_time}', however, your data "
+                             f"only contains data back to '{min_time}'.")
+
+        if query.target_ast.date_offset_range is not None:
+            end_offset = query.target_ast.date_offset_range.end_date_offset
+        else:
+            end_offset = pd.DateOffset(0)
+        end_offset = end_offset * query.num_forecasts
+
+        if (context_anchor_time is not None
+                and context_anchor_time > anchor_time):
+            warnings.warn(f"Context anchor timestamp "
+                          f"(got '{context_anchor_time}') is set to a later "
+                          f"date than the prediction anchor timestamp "
+                          f"(got '{anchor_time}'). Please make sure this is "
+                          f"intended.")
+        elif (query.query_type == QueryType.TEMPORAL
+              and context_anchor_time is not None
+              and context_anchor_time + end_offset > anchor_time):
+            warnings.warn(f"Aggregation for context examples at timestamp "
+                          f"'{context_anchor_time}' will leak information "
+                          f"from the prediction anchor timestamp "
+                          f"'{anchor_time}'. Please make sure this is "
+                          f"intended.")
+
+        elif (context_anchor_time is not None
+              and context_anchor_time - end_offset < min_time):
+            _time = context_anchor_time - end_offset
+            warnings.warn(f"Context anchor timestamp is too early or "
+                          f"aggregation time range is too large. To form "
+                          f"proper input data, we would need data back to "
+                          f"'{_time}', however, your data only contains "
+                          f"data back to '{min_time}'.")
+
+        if not evaluate and anchor_time > max_time + pd.DateOffset(days=1):
+            warnings.warn(f"Anchor timestamp '{anchor_time}' is after the "
+                          f"latest timestamp '{max_time}' in the data. Please "
+                          f"make sure this is intended.")
+
+        if evaluate and anchor_time > max_time - end_offset:
+            raise ValueError(
+                f"Anchor timestamp for evaluation is after the latest "
+                f"supported timestamp '{max_time - end_offset}'.")
+
+    def _get_context(
+        self,
+        query: ValidatedPredictiveQuery,
+        indices: list[str] | list[float] | list[int] | None,
+        anchor_time: pd.Timestamp | Literal['entity'] | None,
+        context_anchor_time: pd.Timestamp | None,
+        run_mode: RunMode,
+        num_neighbors: list[int] | None,
+        num_hops: int,
+        max_pq_iterations: int,
+        evaluate: bool,
+        random_seed: int | None = _RANDOM_SEED,
+        logger: ProgressLogger | None = None,
+    ) -> Context:
+
+        if num_neighbors is not None:
+            num_hops = len(num_neighbors)
+
+        if num_hops < 0:
+            raise ValueError(f"'num_hops' must be non-negative "
+                             f"(got {num_hops})")
+        if num_hops > 6:
+            raise ValueError(f"Cannot predict on subgraphs with more than 6 "
+                             f"hops (got {num_hops}). Please reduce the "
+                             f"number of hops and try again. Please create a "
+                             f"feature request at "
+                             f"'https://github.com/kumo-ai/kumo-rfm' if you "
+                             f"must go beyond this for your use-case.")
+
+        task_type = self._get_task_type(
+            query=query,
+            edge_types=self._sampler.edge_types,
+        )
+
+        if logger is not None:
+            if task_type == TaskType.BINARY_CLASSIFICATION:
+                task_type_repr = 'binary classification'
+            elif task_type == TaskType.MULTICLASS_CLASSIFICATION:
+                task_type_repr = 'multi-class classification'
+            elif task_type == TaskType.REGRESSION:
+                task_type_repr = 'regression'
+            elif task_type == TaskType.TEMPORAL_LINK_PREDICTION:
+                task_type_repr = 'link prediction'
+            else:
+                task_type_repr = str(task_type)
+            logger.log(f"Identified {query.query_type} {task_type_repr} task")
+
+        if task_type.is_link_pred and num_hops < 2:
+            raise ValueError(f"Cannot perform link prediction on subgraphs "
+                             f"with less than 2 hops (got {num_hops}) since "
+                             f"historical target entities need to be part of "
+                             f"the context. Please increase the number of "
+                             f"hops and try again.")
+
+        if num_neighbors is None:
+            if run_mode == RunMode.DEBUG:
+                num_neighbors = [16, 16, 4, 4, 1, 1][:num_hops]
+            elif run_mode == RunMode.FAST or task_type.is_link_pred:
+                num_neighbors = [32, 32, 8, 8, 4, 4][:num_hops]
+            else:
+                num_neighbors = [64, 64, 8, 8, 4, 4][:num_hops]
+
+        if query.target_ast.date_offset_range is None:
+            step_offset = pd.DateOffset(0)
+        else:
+            step_offset = query.target_ast.date_offset_range.end_date_offset
+        end_offset = step_offset * query.num_forecasts
+
+        if anchor_time is None:
+            anchor_time = self._get_default_anchor_time(query)
+
+            if evaluate:
+                anchor_time = anchor_time - end_offset
+
+            if logger is not None:
+                assert isinstance(anchor_time, pd.Timestamp)
+                if anchor_time == pd.Timestamp.min:
+                    pass  # Static graph
+                elif (anchor_time.hour == 0 and anchor_time.minute == 0
+                      and anchor_time.second == 0
+                      and anchor_time.microsecond == 0):
+                    logger.log(f"Derived anchor time {anchor_time.date()}")
+                else:
+                    logger.log(f"Derived anchor time {anchor_time}")
+
+        assert anchor_time is not None
+        if isinstance(anchor_time, pd.Timestamp):
+            if context_anchor_time == 'entity':
+                raise ValueError("Anchor time 'entity' needs to be shared "
+                                 "for context and prediction examples")
+            if context_anchor_time is None:
+                context_anchor_time = anchor_time - end_offset
+            self._validate_time(query, anchor_time, context_anchor_time,
+                                evaluate)
+        else:
+            assert anchor_time == 'entity'
+            if query.query_type != QueryType.STATIC:
+                raise ValueError("Anchor time 'entity' is only valid for "
+                                 "static predictive queries")
+            if query.entity_table not in self._sampler.time_column_dict:
+                raise ValueError(f"Anchor time 'entity' requires the entity "
+                                 f"table '{query.entity_table}' to "
+                                 f"have a time column")
+            if isinstance(context_anchor_time, pd.Timestamp):
+                raise ValueError("Anchor time 'entity' needs to be shared "
+                                 "for context and prediction examples")
+            context_anchor_time = 'entity'
+
+        num_train_examples = _MAX_CONTEXT_SIZE[run_mode]
+        if evaluate:
+            num_test_examples = _MAX_TEST_SIZE[run_mode]
+            if task_type.is_link_pred:
+                num_test_examples = num_test_examples // 5
+        else:
+            num_test_examples = 0
+
+        train, test = self._sampler.sample_target(
+            query=query,
+            num_train_examples=num_train_examples,
+            train_anchor_time=context_anchor_time,
+            num_train_trials=max_pq_iterations * num_train_examples,
+            num_test_examples=num_test_examples,
+            test_anchor_time=anchor_time,
+            num_test_trials=max_pq_iterations * num_test_examples,
+            random_seed=random_seed,
+        )
+        train_pkey, train_time, y_train = train
+        test_pkey, test_time, y_test = test
+
+        if evaluate and logger is not None:
+            if task_type == TaskType.BINARY_CLASSIFICATION:
+                pos = 100 * int((y_test > 0).sum()) / len(y_test)
+                msg = (f"Collected {len(y_test):,} test examples with "
+                       f"{pos:.2f}% positive cases")
+            elif task_type == TaskType.MULTICLASS_CLASSIFICATION:
+                msg = (f"Collected {len(y_test):,} test examples holding "
+                       f"{y_test.nunique()} classes")
+            elif task_type == TaskType.REGRESSION:
+                _min, _max = float(y_test.min()), float(y_test.max())
+                msg = (f"Collected {len(y_test):,} test examples with targets "
+                       f"between {format_value(_min)} and "
+                       f"{format_value(_max)}")
+            elif task_type == TaskType.TEMPORAL_LINK_PREDICTION:
+                num_rhs = y_test.explode().nunique()
+                msg = (f"Collected {len(y_test):,} test examples with "
+                       f"{num_rhs:,} unique items")
+            else:
+                raise NotImplementedError
+            logger.log(msg)
+
+        if not evaluate:
+            assert indices is not None
+            if len(indices) > _MAX_PRED_SIZE[task_type]:
+                raise ValueError(f"Cannot predict for more than "
+                                 f"{_MAX_PRED_SIZE[task_type]:,} entities at "
+                                 f"once (got {len(indices):,}). Use "
+                                 f"`KumoRFM.batch_mode` to process entities "
+                                 f"in batches")
+
+            test_pkey = pd.Series(indices, dtype=train_pkey.dtype)
+            if isinstance(anchor_time, pd.Timestamp):
+                test_time = pd.Series([anchor_time]).repeat(
+                    len(indices)).reset_index(drop=True)
+            else:
+                train_time = test_time = 'entity'
+
+        if logger is not None:
+            if task_type == TaskType.BINARY_CLASSIFICATION:
+                pos = 100 * int((y_train > 0).sum()) / len(y_train)
+                msg = (f"Collected {len(y_train):,} in-context examples with "
+                       f"{pos:.2f}% positive cases")
+            elif task_type == TaskType.MULTICLASS_CLASSIFICATION:
+                msg = (f"Collected {len(y_train):,} in-context examples "
+                       f"holding {y_train.nunique()} classes")
+            elif task_type == TaskType.REGRESSION:
+                _min, _max = float(y_train.min()), float(y_train.max())
+                msg = (f"Collected {len(y_train):,} in-context examples with "
+                       f"targets between {format_value(_min)} and "
+                       f"{format_value(_max)}")
+            elif task_type == TaskType.TEMPORAL_LINK_PREDICTION:
+                num_rhs = y_train.explode().nunique()
+                msg = (f"Collected {len(y_train):,} in-context examples with "
+                       f"{num_rhs:,} unique items")
+            else:
+                raise NotImplementedError
+            logger.log(msg)
+
+        entity_table_names: tuple[str, ...]
+        if task_type.is_link_pred:
+            final_aggr = query.get_final_target_aggregation()
+            assert final_aggr is not None
+            edge_fkey = final_aggr._get_target_column_name()
+            for edge_type in self._sampler.edge_types:
+                if edge_fkey == f'{edge_type[0]}.{edge_type[1]}':
+                    entity_table_names = (
+                        query.entity_table,
+                        edge_type[2],
+                    )
+        else:
+            entity_table_names = (query.entity_table, )
+
+        # Exclude the entity anchor time from the feature set to prevent
+        # running out-of-distribution between in-context and test examples:
+        exclude_cols_dict = query.get_exclude_cols_dict()
+        if entity_table_names[0] in self._sampler.time_column_dict:
+            if entity_table_names[0] not in exclude_cols_dict:
+                exclude_cols_dict[entity_table_names[0]] = []
+            time_column = self._sampler.time_column_dict[entity_table_names[0]]
+            exclude_cols_dict[entity_table_names[0]].append(time_column)
+
+        subgraph = self._sampler.sample_subgraph(
+            entity_table_names=entity_table_names,
+            entity_pkey=pd.concat(
+                [train_pkey, test_pkey],
+                axis=0,
+                ignore_index=True,
+            ),
+            anchor_time=pd.concat(
+                [train_time, test_time],
+                axis=0,
+                ignore_index=True,
+            ) if isinstance(train_time, pd.Series) else 'entity',
+            num_neighbors=num_neighbors,
+            exclude_cols_dict=exclude_cols_dict,
+        )
+
+        if len(subgraph.table_dict) >= 15:
+            raise ValueError(f"Cannot query from a graph with more than 15 "
+                             f"tables (got {len(subgraph.table_dict)}). "
+                             f"Please create a feature request at "
+                             f"'https://github.com/kumo-ai/kumo-rfm' if you "
+                             f"must go beyond this for your use-case.")
+
+        return Context(
+            task_type=task_type,
+            entity_table_names=entity_table_names,
+            subgraph=subgraph,
+            y_train=y_train,
+            y_test=y_test if evaluate else None,
+            top_k=query.top_k,
+            step_size=None,
+        )
+
+    @staticmethod
+    def _validate_metrics(
+        metrics: list[str],
+        task_type: TaskType,
+    ) -> None:
+
+        if task_type == TaskType.BINARY_CLASSIFICATION:
+            supported_metrics = [
+                'acc', 'precision', 'recall', 'f1', 'auroc', 'auprc', 'ap'
+            ]
+        elif task_type == TaskType.MULTICLASS_CLASSIFICATION:
+            supported_metrics = ['acc', 'precision', 'recall', 'f1', 'mrr']
+        elif task_type == TaskType.REGRESSION:
+            supported_metrics = ['mae', 'mape', 'mse', 'rmse', 'smape', 'r2']
+        elif task_type == TaskType.TEMPORAL_LINK_PREDICTION:
+            supported_metrics = [
+                'map@', 'ndcg@', 'mrr@', 'precision@', 'recall@', 'f1@',
+                'hit_ratio@'
+            ]
+        else:
+            raise NotImplementedError
+
+        for metric in metrics:
+            if '@' in metric:
+                metric_split = metric.split('@')
+                if len(metric_split) != 2:
+                    raise ValueError(f"Unsupported metric '{metric}'. "
+                                     f"Available metrics "
+                                     f"are {supported_metrics}.")
+
+                name, top_k = f'{metric_split[0]}@', metric_split[1]
+
+                if not top_k.isdigit():
+                    raise ValueError(f"Metric '{metric}' does not define a "
+                                     f"valid 'top_k' value (got '{top_k}').")
+
+                if int(top_k) <= 0:
+                    raise ValueError(f"Metric '{metric}' needs to define a "
+                                     f"positive 'top_k' value (got '{top_k}')")
+
+                if int(top_k) > 100:
+                    raise ValueError(f"Metric '{metric}' defines a 'top_k' "
+                                     f"value greater than 100 "
+                                     f"(got '{top_k}'). Please create a "
+                                     f"feature request at "
+                                     f"'https://github.com/kumo-ai/kumo-rfm' "
+                                     f"if you must go beyond this for your "
+                                     f"use-case.")
+
+                metric = name
+
+            if metric not in supported_metrics:
+                raise ValueError(f"Unsupported metric '{metric}'. Available "
+                                 f"metrics are {supported_metrics}. If you "
+                                 f"feel a metric is missing, please create a "
+                                 f"feature request at "
+                                 f"'https://github.com/kumo-ai/kumo-rfm'.")
+
+
+def format_value(value: int | float) -> str:
+    if value == int(value):
+        return f'{int(value):,}'
+    if abs(value) >= 1000:
+        return f'{value:,.0f}'
+    if abs(value) >= 10:
+        return f'{value:.1f}'
+    return f'{value:.2f}'