cudf-polars PythonScan support with rank-aware IO sources

[madsbk]

Description

This PR has two parts:

1. It implements IR::PythonScan, making it possible to use a Python function as an IO source for a cudf-polars query via polars.io.plugins.register_io_source(). On main, this path just raises NotImplementedError.
2. It introduces a rank-aware, GPU-backed IO source through a new abstract base class, RankAwareSource. This lets you implement a custom IO source that returns cudf-polars DataFrames based on the worker's rank.

The motivation is likewise two-fold:

1. It is generally useful to generate data directly within a query, e.g. to support a custom file format.
2. RankAwareSource lets us implement Polars' new LazyFrame.execute() API without collecting all the GPU data to the client's host memory. We may eventually want dedicated IR in upstream Polars for this, but a prototype built on just this PR already works very well.

Limitation

We don't support a pushed-down row limit. Polars folds limit / head / tail into the PythonScan as n_rows; cudf-polars rejects that during translation (raising NotImplementedError, which falls back to CPU), because a single global row count can't be enforced across independent per-rank sources. Note the current versions of Polars never folds a limit into the scan when a predicate is also pushed so this limitation is only observed when user use something like limit() or head() with also filtering.

Example

A minimal end-to-end example: generate data from a Python function and run a GPU query over it.

import polars as pl
from polars.io.plugins import register_io_source

# 1. An IO source: a callable yielding polars DataFrames.
def source(with_columns, predicate, n_rows, batch_size):
    yield pl.DataFrame({"a": [1, 2, 3, 4, 5]})

# 2. Register it as a LazyFrame with a declared schema.
lf = register_io_source(source, schema={"a": pl.Int64})

# 3. Build and collect a query on the GPU.
result = lf.filter(pl.col("a") > 2).collect(engine=pl.GPUEngine())
print(result)
# shape: (3, 1)
# ┌─────┐
# │ a   │
# │ --- │
# │ i64 │
# ╞═════╡
# │ 3   │
# │ 4   │
# │ 5   │
# └─────┘

What happens internally:

1. register_io_source produces a LazyFrame whose root is a Polars PythonScan node wrapping source.
2. .collect(engine=...) hands the optimized plan to cudf-polars, which translates the PythonScan into IR::PythonScan (the predicate col("a") > 2 is pushed into it).
3. PythonScan.do_evaluate calls source, moves each yielded frame to the GPU , validates the result against the declared schema, and applies the pushed predicate on the device.
4. The filtered GPU frame is returned and copied back to host as the result.

xref #22917 for the gaps in Polars we need TODO

[Matt711]

Thanks @madsbk I think this is in a good place. I mostly left minor suggestions. But since there are a bunch, I'll approve after you look them over.

FYI: I also opened up some sub-issues for the gaps you found in polars in #22917.

[wence-]

Minor nits, looks good. Thanks!

[wence-]

No part of the introductory documentation has previously talked about these cudf_polars.containers.DataFrame objects?

[madsbk]

Good point, updated:

Chunks are normally returned as regular host-resident `polars.DataFrame`
objects, but cudf-polars internally represents GPU-resident data using
`cudf_polars.containers.DataFrame`.

An IO source may return `cudf_polars.containers.DataFrame` objects directly.
Doing so avoids the host-to-device copy, but restricts the source to
cudf-polars engines, since the default Polars CPU engine cannot consume
`cudf_polars.containers.DataFrame` objects.

[wence-]

Do we need this section?

[madsbk]

I am not sure. I think we do. This isn't obvious from the API, and users may
reasonably expect things like sqlite3.Connection objects or other bridging
libraries to work. I think it's worth stating explicitly?

[wence-]

This seems wrong for the GPU-resident input? Surely you need to reserve for the output in both cases, so:

# Reserve for the filtered output if it exists, and moving the input to device if it not already there
reservation = chunk.estimated_size() * (1 + (ir.predicate is not None) - isinstance(chunk, DataFrame))

By an abuse of casting.

[madsbk]

good catch, fixed: 7660aae

[wence-]

nitpick: Perhaps we should just provide the communicator to the source? That would be necessary for limit in multi-rank cases, I think.

[wence-]

Suggested change

	while True:
	while !ch_out.is_shutdown():

So that if the consumer doesn't want any more we don't keep making chunks.

[madsbk]

I don't think this is correct, if the downstream shutsdown cout_out we should fail, which we already do.

[madsbk]

/merge