[Enh]: add LazyFrame.collect_batches

[danielgafni]

We would like to learn about your use case. For example, if this feature is needed to adopt Narwhals in an open source project, could you please enter the link to it below?

The user should be given an option to ORDER BY deterministically first, while the default behavior may produce batches in non-deterministic ordering (including rows within batches of course).

Please describe the purpose of the new feature or describe the problem to solve.

The corresponding Polars functionality is incredibly useful for parallel batch processing. Often there is a main entrypoint that can collect these batches and send them to (remote) parallel workers. I'm typically doing this kind of work with Ray.

Suggest a solution if possible.

For backends that do not have a corresponding method natively, I assume it's reasonable to implement it via with_row_index.

If you have tried alternatives, please describe them below.

No response

Additional information that may help us understand your needs.

Happy to implement this feature as my first Narwhals contribution given it makes sense to have it

[MarcoGorelli]

thanks for the request - not sure about this one tbh, what would it look like for other backends?

[danielgafni]

This is a rough implementation I just drafted:

def collect_batches(
    df: Frame,
    chunk_size: int | None = None,
    order_by: Sequence[str] | None = None,
    **kwargs
) -> Iterator[nw.DataFrame[Any]]:
    if isinstance(df, nw.DataFrame):
        df = df.lazy()

    assert isinstance(df, nw.LazyFrame)

    if df.implementation == nw.Implementation.POLARS:
        df_polars = cast(pl.LazyFrame, df.to_native())

        if chunk_size is None:
            # No chunking requested, collect everything at once
            yield nw.from_native(df_polars.collect(**kwargs))
        else:
            if order_by is not None:
                # the user requests a specific ordering
                # and we assume the DataFrame has already been sorted by these columns
                # so all we need to do is not to break it
                kwargs["maintain_order"] = True

            for batch in df_polars.collect_batches(chunk_size=chunk_size, **kwargs):
                yield nw.from_native(batch)

    # other backends don't have a native implementation for this operations,
    # so we have to think about ordering
    # the use must provide a valid order_by argument
    # and actually order the table in advance (or maybe we can do it automatically, not sure?)
    else:
        if chunk_size is None:
            yield df.collect()
        else:
            drop_cols = ["__collect__batches_index__"]
            if order_by is None:
                # there is no ordering requested
                # but with_row_index needs it
                # so we'll just make it up with minimal performance overhead
                order_by = ["__collect_batches_order__"]
                df = df.with_columns(nw.lit(True).alias("__collect_batches_order__"))
                drop_cols.append("__collect_batches_order__")

            df = df.with_row_index(
                "__collect__batches_index__",
                order_by=order_by
            )

            total_len = df.select(nw.len()).collect().item()

            for start in range(0, total_len, chunk_size):
                yield df.filter(
                    (nw.col("__collect__batches_index__") >= start)
                    & (nw.col("__collect__batches_index__") < start + chunk_size)
                ).drop(*drop_cols).collect()

This differs just a bit from Polars' API with the new order_by argument.
Polars has maintain_order=True instead.

Note that we can also add a lazy=True/False argument to yield lazy frames (see the current .collect() invocation, we don't really have to do it for lazy frames)

[MarcoGorelli]

thanks!

for now my inclination is we shouldn't support this, and that this can be handled in user-code as you've demonstrated

if any other backend supports it natively, we can reconsider

[danielgafni]

For future, what is the decision algorithm you have been following here? What are the factors I should be looking at? Is it implementation complexity and ability to delegate to native APIs?

[MarcoGorelli]

one of the main ones is that the api should be honest, in the sense that if code looks efficient, it shouldn't use inefficient fallbacks behind the scenes

doing

            total_len = df.select(nw.len()).collect().item()

            for start in range(0, total_len, chunk_size):
                yield df.filter(
                    (nw.col("__collect__batches_index__") >= start)
                    & (nw.col("__collect__batches_index__") < start + chunk_size)
                ).drop(*drop_cols).collect()

looks potentially very expensive, and a recipe for repeating the same computation multiple times over

[dangotbanned]

Is this not what these APIs are for?

• https://duckdb.org/docs/stable/clients/python/relational_api#fetch_arrow_reader
• https://duckdb.org/docs/stable/clients/python/relational_api#fetch_arrow_table
  • If you still wanted tables
• https://github.com/eakmanrq/sqlframe/blob/5ec7d3530865d430c05096dcec92cfa57d8046a8/sqlframe/duckdb/dataframe.py#L55-L65
• https://ibis-project.org/reference/expression-tables#ibis.expr.types.relations.Table.to_pyarrow_batches

IIUC, we'd just be plumbing these into a single function that handles the pyarrow.RecordBatchReader -> nw.DataFrame wrapping on each iteration, no?

[danielgafni]

@MarcoGorelli I see, so Narwhals APIs should be thin wrappers around native ones without introducing too much logic on top of them? Do you primarily implement APIs that just delegate to backend-specific ones?

I think performance-wise there is nothing much to do in this case, Polars just inherently has more guarantees than SQL because it execute the whole query while yielding rows as they are materialized on the fly.

[danielgafni]

@dangotbanned oh, I wasn't aware of these. Seems like the only difference between these and how Polars implements this is that Polars collects all the rows/batches concurrently and yields them as they become available, while Ibis would return a sequential batch reader? Unfortunately I don't have a good enough understanding of Ibis.

But this seems like a much better way forward with this functionality?

[MarcoGorelli]

thin wrappers around native ones without introducing too much logic on top of them

extra logic is fine, extra hidden collects aren't

Is this not what these APIs are for?

from a quick look, yeah that might be what we're after, thanks!

[danielgafni]

extra logic is fine, extra hidden collects aren't

Not sure if this is applicable here, the method is called **.collect**_batches. It's supposed to materialize things into memory.

[MarcoGorelli]

Sure, but not this way

If df contains "blocking" operations like window functions, then the filters at the end won't get pushed down and computations will get repeated in a way that I think would be unexpected to the user

[dangotbanned]

@dangotbanned oh, I wasn't aware of these.

Yeah I only learned about this through (#3015), (duckdb/duckdb#18642) and then some recent work with Acero which can make use of these iterators

But this seems like a much better way forward with this functionality?

Nice one @danielgafni, yeah I hope so!

Seems like the only difference between these and how Polars implements this is that Polars collects all the rows/batches concurrently and yields them as they become available, while Ibis would return a sequential batch reader? Unfortunately I don't have a good enough understanding of Ibis.

What I'd suggest is try to put together some repro(s) that look like your use case to compare between polars, duckdb and ibis.
I think that if you put together a good case for why ibis should support things concurrently here (demonstrating the real-world benefit) - then you might able to convince them that this integration with Arrow aligns with their goals

---

On the flipside, if we can make do with just using pyarrow.RecordBatchReader then my (optimistic) assumption is we can support LazyFrame.collect_batches for:

• Natively
  • polars
• With a thin, shared wrapper
  • duckdb
  • ibis
  • sqlframe (possibly only when backed by duckdb)
• Fake lazy
  • I'm not sure if this is a good idea, but I'm leaving it here for API consistency
  • pyarrow.Table.from_batches
    • E.g. Table.from_batches([RecordBatchReader(...).read_next_batch()])
  • pyarrow.RecordBatch.to_pandas
    • E.g. RecordBatchReader(...).read_next_batch().to_pandas()