NVIDIA / NVIDIA/cudf

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

Open
#24,110 3 comments 0 reactions 2 assignees Claimed by @shrshi View on GitHub
cudf-polars feature request libcudf
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

This issue supersedes/unifies [#2231](https://github.com/NVIDIA/cudf/issues/2231) and [#20727](https://github.com/NVIDIA/cudf/issues/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`](https://pandas.pydata.org/docs/reference/api/pandas.merge_asof.html)) and cudf-polars ([`LazyFrame.join_asof`](https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.join_asof.html)).

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:

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

Multiple equality keys:

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

Ungrouped after filtering to one entity:

```python
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.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.