Python IO Sources#

Polars lets you feed a query from a custom Python function instead of a file, by registering an IO plugin with polars.io.plugins.register_io_source. cudf-polars executes such sources on the GPU, so you can generate or load data directly into a query, for example from a custom file format, a remote store, or data produced on the fly.

See the Polars IO plugins guide for the full description of the IO-source contract. This page covers the cudf-polars-specific behavior.

Plain IO Sources#

An IO source is a callable with the signature (with_columns, predicate, n_rows, batch_size) that yields one or more chunks (i.e. polars DataFrames):

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


def source(with_columns, predicate, n_rows, batch_size):
    df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
    if predicate is not None:
        df = df.filter(predicate)
    if n_rows is not None:
        df = df.head(n_rows)
    if with_columns is not None:
        df = df.select(with_columns)
    yield df


lf = register_io_source(source, schema={"a": pl.Int64, "b": pl.Int64})
result = lf.select("a").collect(engine=pl.GPUEngine())

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.

The source may yield multiple chunks, which cudf-polars combines into the scan output (under a streaming engine the chunks are forwarded individually; see Streaming Chunks With a Known Count).

Schema Validation#

The columns a source emits (after applying with_columns) must match the registered schema in name, order, and dtype. cudf-polars validates the produced output against that schema and raises polars.exceptions.SchemaError on a mismatch.

Polars itself only validates when register_io_source(..., validate_schema=True) is used, but that flag is not carried into the GPU plan, so cudf-polars validates unconditionally. A source that deliberately yields a dtype different from its declared schema (only valid with validate_schema=False) therefore cannot run on the GPU and must be collected with the default Polars CPU engine. This is tracked in cudf#22917.

Rank-Aware Sources#

With a multi-GPU streaming engine (see Engines), every rank runs the scan function. Plain Python sources are not rank-aware, so cudf-polars executes them on rank 0 only.

For distributed loading, use a rank-aware source that subclasses RankAwareSource. Its __call__ method follows the register_io_source contract and adds two optional trailing arguments:

rank: the zero-based rank running the source.
nranks: the total number of ranks in the query.

Both default to rank=0, nranks=1 for single-rank execution, the in-memory cudf-polars engine, and the default Polars CPU engine.

import polars as pl
from polars.io.plugins import register_io_source
from cudf_polars.streaming.rank_aware_source import RankAwareSource


class PartitionedFrame(RankAwareSource):
    def __init__(self, frame: pl.DataFrame):
        self.frame = frame

    def __call__(
        self,
        with_columns: list[str] | None,
        predicate: pl.Expr | None,
        n_rows: int | None,
        batch_size: int | None,
        rank: int = 0,
        nranks: int = 1,
    ):
        df = self.frame.gather_every(nranks, offset=rank)
        if predicate is not None:
            df = df.filter(predicate)
        if n_rows is not None:
            # Only reached on the Polars CPU engine (see the Row-Limit Pushdown section).
            df = df.head(n_rows)
        if with_columns is not None:
            df = df.select(with_columns)
        yield df


source = PartitionedFrame(pl.DataFrame({"a": range(10)}))
lf = register_io_source(source, schema={"a": pl.Int64})

Pass the RankAwareSource instance directly to register_io_source; do not wrap it in another callable. To inject rank information, cudf-polars must recover the RankAwareSource instance from the registered callable. Only an unwrapped instance is recognized; wrapping it in anything else (for example, a functools.partial, closure, lambda, or decorator) hides the instance, in which case the source is treated as rank-unaware and executes on rank 0 only.

This limitation is tracked in cudf#22917.

Streaming Chunks With a Known Count#

Under a streaming executor, cudf-polars forwards each chunk produced by an IO source through the pipeline individually. However, a Polars source is a lazy iterator whose length is unknown until fully consumed, so cudf-polars normally has to drain the source before it can start forwarding chunks.

A source that already knows its chunk count can avoid this by returning a SizedChunks: a thin iterator wrapper that also reports its length. cudf-polars can then stream one chunk at a time, keeping only a single chunk resident on the GPU at once. Since SizedChunks is still a normal iterator, the default Polars CPU engine handles it unchanged.

import polars as pl
from polars.io.plugins import register_io_source
from cudf_polars.streaming.rank_aware_source import SizedChunks


def source(with_columns, predicate, n_rows, batch_size):
    def chunks():
        # Produced lazily, but the total count is known up front.
        for path in ("a.parquet", "b.parquet"):
            yield pl.read_parquet(path)

    return SizedChunks(2, chunks())


lf = register_io_source(source, schema={"a": pl.Int64})

Row-Limit Pushdown#

Polars can push head / limit operations into a Python scan via the n_rows parameter. cudf-polars does not currently support this pushdown and rejects such plans during translation. Standard GPU fallback behavior applies: by default the query falls back to the Polars CPU engine, and in raise-on-fail mode cudf-polars raises NotImplementedError.

As a result, n_rows is always None on GPU engines. However, IO sources intended to also work with the default Polars CPU engine must still handle n_rows correctly.

Distributed support for a global row limit is tracked in cudf#22918.

Threading#

Under a streaming engine, cudf-polars runs IO sources on a worker thread pool. A source is created on a worker thread, and successive chunks are pulled on worker threads that may differ from the one that created the source and from each other. A source must therefore not depend on thread-affine state that is created up front and reused across chunks, for example a sqlite3.Connection (which by default may only be used on the thread that opened it). Open such resources inside the function that produces each chunk, or use a thread-safe equivalent.

API#

class cudf_polars.streaming.rank_aware_source.RankAwareSource[source]#

Base class for rank-aware Polars IO-plugin sources.

Subclasses must implement __call__(), which follows the standard Polars io_source contract with additional rank and nranks keyword arguments supplied by cudf-polars streaming engines.

During multi-rank streaming execution, every rank calls the IO source once. The implementation should therefore use rank and nranks to produce the rank-local rows of the global dataframe.

This is an interim cudf-polars API. It exists because Polars’ IO-source contract has no way to thread rank information into a source. It may be removed if Polars exposes a supported mechanism. Tracked in rapidsai/cudf#22917.

The rank and nranks arguments default to 0 and 1 for the in-memory cudf-polars engine, the default Polars engine, and single-rank streaming.

A single RankAwareSource instance may be invoked concurrently (the same source can appear in more than one scan of a query, and the streaming engine runs scans on separate threads). Implementations must therefore be safe to call concurrently. The streaming engine also advances a single returned iterator on a worker thread pool, so successive chunks may be produced on different threads thus do not hold thread-affine state across chunks (see the Threading section of the IO-plugins guide).

Pass the RankAwareSource instance directly to register_io_source; do not wrap it in another callable. To inject rank information, cudf-polars must locate the RankAwareSource instance inside the registered callable. Only an unwrapped instance is recognized; wrapping it in anything else (a functools.partial, closure, lambda, or decorator) hides it, in which case the source is treated as rank-unaware and runs on rank 0 only. This limitation is tracked in rapidsai/cudf#22917.

Methods

__call__(with_columns, predicate, n_rows, ...)

Produce the rank-local chunks of this IO source.

Examples

>>> import polars as pl
>>> from polars.io.plugins import register_io_source
>>> from cudf_polars.streaming.rank_aware_source import RankAwareSource
>>>
>>> class StripedSource(RankAwareSource):
...     def __init__(self, df):
...         self.df = df
...
...     def __call__(
...         self,
...         with_columns,
...         predicate,
...         n_rows,
...         batch_size,
...         rank=0,
...         nranks=1,
...     ):
...         out = self.df.gather_every(nranks, offset=rank)
...         if with_columns is not None:
...             out = out.select(with_columns)
...         if predicate is not None:
...             out = out.filter(predicate)
...         yield out
>>>
>>> lf = register_io_source(
...     StripedSource(pl.DataFrame({"a": [1, 2, 3]})),
...     schema={"a": pl.Int64},
... )

Produce the rank-local chunks of this IO source.

Parameters:

with_columns: Projected column names. If not None, the source should return only these columns.
predicate: Polars expression. The reader must filter their rows accordingly.
n_rows: Maximum number of rows requested from the scan. cudf-polars does not support a pushed-down row limit and rejects it during translation, so this is always None on a GPU engine (see rapidsai/cudf#22918). It may be non-None on the default Polars CPU engine, where the source must honor it.
batch_size: Optional hint for the number of rows to yield per chunk.
rank: Rank running this scan function, bound by the streaming engine. Defaults to 0 for single-rank / in-memory / default Polars-engine execution.
nranks: Total number of ranks (the world size), bound by the streaming engine. Defaults to 1 for single-rank execution.

Returns:

Chunks for this rank. Return a SizedChunks if the total chunk count
is known upfront to enable lazy streaming.

Notes

Returning one or more cudf_polars.containers.DataFrame objects restricts the source to cudf-polars engines. To support the default (host) Polars engine, an IO source should only return polars.DataFrame.

The emitted columns, after applying with_columns, must match the registered schema in name, order, and dtype. cudf-polars always validates this and raises polars.exceptions.SchemaError on mismatch. Polars only performs this validation when register_io_source(validate_schema=True) is used, but that flag is not exposed to the GPU execution path.

class cudf_polars.streaming.rank_aware_source.SizedChunks(count: int, chunks: Iterable[pl.DataFrame | DataFrame])[source]#

An iterator that also reports how many chunks it will produce.

A Polars IO source returns an iterator, so the total number of chunks is generally unknown until the iterator is fully consumed. When a source already knows its chunk count, wrapping the chunks in SizedChunks lets the cudf-polars streaming engine forward chunks lazily, keeping only a single device chunk resident at a time instead of materializing all chunks up front to learn the count.

SizedChunks remains an ordinary iterator, so the default (host) Polars engine handles it unchanged.

This is an interim cudf-polars API. It exists because Polars’ IO-source contract returns a plain iterator with no length. It may be removed if Polars exposes a supported way to report the chunk count. Tracked in rapidsai/cudf#22917.

Parameters:

count: Number of chunks chunks will yield.
chunks: Iterable of host polars.DataFrame or GPU-resident cudf_polars.containers.DataFrame chunks.

Examples

>>> import polars as pl
>>> from polars.io.plugins import register_io_source
>>> from cudf_polars.streaming.rank_aware_source import SizedChunks
>>>
>>> def source(with_columns, predicate, n_rows, batch_size):
...     def chunks():
...         yield pl.DataFrame({"a": [1, 2]})
...         yield pl.DataFrame({"a": [3, 4]})
...
...     return SizedChunks(2, chunks())
>>>
>>> lf = register_io_source(source, schema={"a": pl.Int64})