[Story][FEA] Support as-of joins in libcudf, pylibcudf, cuDF Python, and cudf-polars

[rjzamora]

This issue supersedes/unifies #2231 and #20727. Some older (but useful) discussion can be found in those issues.

Summary

Add GPU support for as-of joins, with a libcudf primitive that can be used by both cuDF Python (merge_asof) and cudf-polars (LazyFrame.join_asof).

An as-of join is a left-preserving time/order join. For every row on the left, find at most one row on the right whose ordered key is the closest valid key in the requested direction, optionally within matching equality groups such as entity or (entity, source).

This is a common operation in time-series and event-stream analytics (e.g. market data). It is especially useful for aligning timestamped events with the latest known state, derived state snapshots, or cumulative per-entity state.

This issue is intended to be a long-lived story for as-of join support. The v1 section below defines the first delivery milestone.

v1 Delivery Scope

The v1 target is the backward as-of join used to enrich timestamped events with the latest known state for the same entity/source.

Required Semantics and Types

v1 must support:

• strategy="backward": match the last right row where right_on <= left_on within the same group.
• Ungrouped joins on a single ordered key.
• Grouped joins with one equality key, for example by="entity".
• Grouped joins with multiple equality keys, for example by=["entity", "source"].
• Ordered key types needed for Polars time-series workloads: integer, Date, Datetime, Duration, and Time.
• Equality group key types needed for real query shapes, especially strings and fixed-width scalar columns.
• Default exact-match behavior (allow_exact_matches=True).
• Duplicate right-side as-of keys with Polars-compatible behavior.
• Null ordered-key and grouping-key behavior compatible with Polars.

Useful v1 extensions, if they fit naturally into the primitive:

• allow_exact_matches=False: for backward matching, choose right_on < left_on instead of right_on <= left_on.
• strategy="forward": match the first right row where right_on >= left_on within the same group.
• Fixed-distance tolerance: after selecting the candidate row, reject it if the ordered-key distance exceeds a numeric or fixed temporal tolerance.

Representative Polars Query Shapes

These snippets are illustrative query shapes, not standalone reproducers.

Single equality key:

events.join_asof(
    state_updates,
    left_on="event_ts",
    right_on="state_ts",
    by="entity",
    strategy="backward",
)

Multiple equality keys:

events.join_asof(
    source_state,
    left_on="event_ts",
    right_on="state_ts",
    by=["entity", "source"],
    strategy="backward",
)

Ungrouped after filtering to one entity:

entity_events.join_asof(
    entity_state,
    left_on="event_ts",
    right_on="state_ts",
    strategy="backward",
)

libcudf Primitive Contract

The exact C++ API shape is up to the libcudf implementation. cudf-polars can work with separate on/by arguments, packed key tables, or a reusable sort-merge-style object. The v1 primitive contract we need is:

• One ordered as-of key for each side.
• Zero or more equality grouping keys for each side, compared left-to-right.
• Corresponding left/right key columns have matching or otherwise compatible dtypes.
• v1 group key columns support strings and fixed-width scalar columns.
• v1 as-of key columns support integer, Date, Datetime, Duration, and Time.
• Inputs are sorted ascending by (group_keys..., asof_key) for v1, with nulls ordered before non-nulls.
• The result is a right-side gather map with one entry per left row, in left input row order.
• Each gather-map value is either the matching right row index or cudf::JoinNoMatch.
• JoinNoMatch is returned when the left as-of key is null, when no right row exists in the same group, or when the selected candidate is rejected by an enabled tolerance policy.
• Right rows with null as-of keys are not valid match candidates.
• With compare_group_nulls == null_equality::UNEQUAL, null group keys do not match. This is the required default for Polars semantics.
• For strategy == BACKWARD, choose the last right row in the same group where right_on <= left_on.
• If multiple right rows have the selected as-of key, backward matching picks the last eligible duplicate.

Returning only the right gather map avoids allocating an unnecessary identity left gather map for this left-preserving operation. A pair-returning API would also be workable if that better matches existing libcudf join conventions.

NEAREST can be added to the strategy enum/API in a follow-up once backward and forward semantics are stable. v1 should not require nearest support.

pylibcudf / cudf-polars Integration

The pylibcudf binding should expose the libcudf primitive directly enough that cudf-polars can request the v1 contract above and receive right-side gather indices.

At the cudf-polars layer, we would materialize Polars join keys, call pylibcudf, gather right-side payload columns with NULLIFY for JoinNoMatch, and then apply Polars-specific output rules: suffixing, coalescing, join-key handling, and final column assembly.

v1 Acceptance Criteria

• libcudf exposes a v1 as-of join primitive satisfying the contract above.
• pylibcudf exposes the primitive and enum/options needed by cudf-polars.
• cudf-polars can execute the v1 LazyFrame.join_asof query shapes above with strategy="backward", including ungrouped, single-key grouped, and multi-key grouped joins with same-name and different-name as-of keys.
• v1 tests cover no-match rows, duplicate right-side as-of keys, null ordered keys, and null grouping keys.

Follow-Up Work

• [P0] Add streaming and multi-GPU cudf-polars support.
• [P1] Add allow_exact_matches=False if it is not included in v1.
• [P1] Add strategy="forward" if it is not included in v1.
• [P2] Add broader tolerance support, including calendar-duration strings such as months, quarters, and years whose physical duration depends on the timestamp being evaluated.
• [P3] Add strategy="nearest" unless specifically requested earlier by users.
• [P3] Add higher-level cuDF Python merge_asof support.
• Fill in broader Polars parity details after the initial primitive and cudf-polars V1 integration are in place.

[wence-]

Just to note. I did some survey of the semantics of asof joins in polars, pandas, and SQL (duckdb, etc...). This is probably helpful so that we don't paint ourselves into a corner with the API choices.

Here's what my spelunking found:

Direction and inequality mapping

For a left ordering value L and right ordering value R:

Public strategy	SQL predicate	Selected right value
BACKWARD, exact allowed	L >= R	Greatest R <= L
BACKWARD, exact disallowed	L > R	Greatest R < L
FORWARD, exact allowed	L <= R	Smallest R >= L
FORWARD, exact disallowed	L < R	Smallest R > L
NEAREST	Not a native DuckDB SQL as-of form	R minimizing the distance from L

Every equality key must also match. Each left row selects at most one right row.

Ordering contracts

• pandas requires both inputs to be globally ascending on the as-of key. The
  by keys do not need to be sorted.
• Polars requires the as-of key to be ascending within each by group.
• DuckDB accepts unsorted inputs and performs the necessary organization
  internally.

Duplicate selection for nearest

Nearest behaviour needs two categories of policy:

1. Which side wins when the preceding and following values are equally distant.
2. Which duplicate within the selected following value is returned.

Observed and documented behaviour differs:

Behaviour	pandas	Polars
Equal-distance nearest tie	preceding	following
Duplicate at preceding value	last	last
Duplicate at following value	first	last
Duplicate at exact value	last	last

[shrshi]

Thank you for raising this issue @rjzamora . I'm interested in taking up the libcudf as-of join primitive.

API design

What do you think of the following class-based API where, similar to sort_merge_join, we preprocess the right side once if necessary, and support repeated probing?

  enum class asof_join_strategy {
    BACKWARD,
    FORWARD,
    NEAREST
  };

  class asof_join {
   public:
    asof_join(
      table_view const& right_by,
      column_view const& right_on,
      cuda::stream_ref stream);

    std::unique_ptr<rmm::device_uvector<size_type>> join(
      table_view const& left_by,
      column_view const& left_on,
      asof_join_strategy strategy,
      bool allow_exact_matches,
      cuda::stream_ref stream,
      rmm::device_async_resource_ref mr) const;

   private:
    table_view _right_by;
    column_view _right_on;
    rmm::device_uvector<size_type> _right_group_offsets;
    size_type _num_right_groups;
  };

Each left row would locate its matching right-side group using the rows referenced by _right_group_offsets, then binary-search right_on within that group’s offset range. For BACKWARD, upper_bound - 1 provides the candidate and last-duplicate behavior. This would be a single batched GPU operation. OTOH, for ungrouped joins, the right side is treated as one range and group lookup is unnecessary.

Staging work

To re-iterate the scope for v1 based on what you've mentioned above, the first libcudf PR in this series would support:

• BACKWARD, with exact matches enabled;
• ungrouped, single-key grouped, and multi-key grouped joins;
• inputs sorted ascending by (by..., on);
• null_equality::UNEQUAL for grouping keys;
• null_order::BEFORE.

Apart from implementing the other strategies (FORWARD, NEAREST) and supporting allow_exact_matches=false, I think follow-on PRs would also need to add the configurability supported by other libcudf ordered operations:

• whether the right input is sorted, with internal sorting when necessary;
• group-key null_equality;
• per-column null ordering for (by..., on).

[rjzamora]

Awesome, thanks @shrshi !

For V1, requiring the right side to be sorted by (by..., on) may be reasonable for the libcudf primitive. However, in the market-data-shaped workloads we care about, datasets are typically ordered by time, with groups interleaved. This is probably why Polars requires the on key to be sorted within each by group, but does not necessarily require physical contiguity by group. That said, contiguous groups do seem like a sensible precondition for the libcudf primitive.

@wence- , this makes me wonder if we should prioritize support for the tolerance option in V1, if feasible. When the input is not already sorted by (by..., on), boundary exchange is much more painful with tolerance=None: for a backward as-of join, the latest preceding right row for the same group may be many chunks/ranks away. The initial cudf-polars implementation could force a global sort by (by..., on) for this case, but it would be nice to avoid that for naturally time-ordered data.