How to change the row_count of RecordBatch returned from reading a parquet file

[twitu]

pub struct RecordBatch {
    schema: SchemaRef,
    columns: Vec<Arc<dyn Array>>,

    /// The number of rows in this RecordBatch
    ///
    /// This is stored separately from the columns to handle the case of no columns
    row_count: usize,
}

RecordBatch has row_count and by default it returns 1000 rows at a time. When reading a parquet from a parquet file using an sql query. Is there anyway to configure this?

    let config = SessionConfig::new();
    let config = SessionConfig::with_batch_size(config, 5000);
    let ctx = SessionContext::with_config(config);

    // execute the query
    let df = ctx.sql("SELECT * FROM alltypes_plain").await?;

    // println!("{:?}", &df.schema());
    // let batches = df.clone().collect();
    // // print the results
    // df.show().await?;
    let mut stream = df.execute_stream().await.unwrap();
    ....
    Some(batch) = stream.next().await // use record batch here
    let batch: RecordBatch = batch.unwrap();
    dbg!(&batch.num_rows());

I've tried changing the ExecutionOptions batch_size but it still returns a 1000 rows per batch. I could not find any other options that can configure this. However changing the query to "SELECT * FROM alltypes_plain ORDER BY <field>" makes row_count 5000.

[alamb]

I think the batch_size https://docs.rs/datafusion/17.0.0/datafusion/config/struct.ExecutionOptions.html#structfield.batch_size is the correct setting. I don't think the size is guaranteed. The 1000 rows is likely coming from the underlying reader (how are you registering alltypes_plain?)

[alamb]

One way to understand why the ORDER BY makes a difference is to use EXPLAIN <query>

So in your case

    let df = ctx.sql("EXPLAIN SELECT * FROM alltypes_plain ORDER BY c1").await?;

You should see that the plans are different for different queries. The ORDER BY query will almost certainly have a SortExec in it. The SortExec buffers data (to sort it) and then produces batches of the configured size . Since they were buffered anyways

If you want to get fancy, you can try to add a CoaelsceBatches exec to your plan -- which is responsible for resizing RecordBatches

[twitu]

Hmm understood. Datafusion is quiiite powerful that way, because the logical plan is extensible, which means a lot to read up on 😄. For my specific question on configuring the options for the row_count of a record batch while reading a parquet file is a no, right?

[alamb]

I think that is correct -- sounds like a useful feature to add, however if you have time :)

[tustvold]

The parquet reader will respect the batch size setting with one caveat, it will not read across row groups. Is it possible the parquet file you are reading has a row group size of 1000 rows?

The odd bit is that using the ORDER BY clause starts using the larger batch_size of 5000.

This is expected, the sort requires buffering up records to the larger batch size

[twitu]

Hmm yeah the the file I'm testing on I wrote it with row groups of 1000. So it's picking that up.