Implementation:Lance format Lance DataFrameExt

Knowledge Sources	Lance
Domains	DataFusion_Integration, Query_Execution
Last Updated	2026-02-08 19:33 GMT

Overview

The DataFrameExt module extends DataFusion's DataFrame with a group_by_stream method that produces grouped record batch streams partitioned by a column value.

Description

This module provides:

• DataFrameExt trait -- An async trait that extends datafusion::dataframe::DataFrame with a group_by_stream method. Given a set of partition columns (currently limited to one), it executes the DataFrame and wraps the result in a BatchStreamGrouper.
• BatchStreamGrouper -- A stream struct that pulls batches from an input SendableRecordBatchStream and groups them by a partition column value. It buffers batches for the current partition and emits a complete group (partition key + collected batches) each time a new partition value is encountered. The partition column is removed from the output schema. The input data must already be sorted by the partition column.

The grouper uses Arrow's partition function to detect distinct value ranges within each batch, handles cross-batch partition boundaries, and supports streaming output via the futures::Stream trait.

Usage

Use DataFrameExt when you need to process DataFrame results grouped by a column, such as:

• Writing partitioned datasets where each partition is handled separately
• Computing per-group aggregations over sorted data
• Splitting sorted scan results by a key column for downstream processing

Code Reference

Source Location

rust/lance-datafusion/src/dataframe.rs

Signature

#[async_trait]
pub trait DataFrameExt {
    async fn group_by_stream(self, partition_columns: &[&str]) -> DFResult<BatchStreamGrouper>;
}

pub struct BatchStreamGrouper {
    input: SendableRecordBatchStream,
    partition_column: String,
    schema: Arc<Schema>,
    buffer: Vec<RecordBatch>,
    current_partition: Option<ScalarValue>,
    unprocessed: Option<(Vec<GroupRange>, RecordBatch)>,
}

Import

use lance_datafusion::dataframe::DataFrameExt;

I/O Contract

Input	Type	Description
self	DataFrame	A DataFusion DataFrame (data must be sorted by partition columns)
partition_columns	&[&str]	Column names to partition by (currently only one column is supported)

Output	Type	Description
BatchStreamGrouper	Stream<Item = DFResult<(Vec<ScalarValue>, Vec<RecordBatch>)>>	A stream yielding tuples of (partition key values, grouped record batches) with the partition column removed from the schema

Usage Examples

use lance_datafusion::dataframe::DataFrameExt;
use datafusion::prelude::*;

// Assuming `df` is a sorted DataFrame
let grouper = df.group_by_stream(&["category"]).await?;
let schema = grouper.schema().clone();

// Iterate over groups
use futures::StreamExt;
let mut grouper = grouper;
while let Some(result) = grouper.next().await {
    let (partition_keys, batches) = result?;
    // Process each group of batches
}

Related Pages

• Lance_format_Lance_ExecPlans -- Execution plan infrastructure used by DataFrames
• Lance_format_Lance_FilterPlanner -- SQL filter planning that may precede grouped execution