Function format_batches_with_maxrows

datafusion-cli/src/print_format.rs:109–154

(
    writer: &mut W,
    batches: &[RecordBatch],
    maxrows: MaxRows,
    format_options: &FormatOptions,
)

}

fn format_batches_with_maxrows<W: std::io::Write>(
    writer: &mut W,
    batches: &[RecordBatch],
    maxrows: MaxRows,
    format_options: &FormatOptions,
) -> Result<()> {
    let options: arrow::util::display::FormatOptions = format_options.try_into()?;

    match maxrows {
        MaxRows::Limited(maxrows) => {
            // Filter batches to meet the maxrows condition
            let mut filtered_batches = Vec::new();
            let mut row_count: usize = 0;
            let mut over_limit = false;
            for batch in batches {
                if row_count + batch.num_rows() > maxrows {
                    // If adding this batch exceeds maxrows, slice the batch
                    let limit = maxrows - row_count;
                    let sliced_batch = batch.slice(0, limit);
                    filtered_batches.push(sliced_batch);
                    over_limit = true;
                    break;
                } else {
                    filtered_batches.push(batch.clone());
                    row_count += batch.num_rows();
                }
            }

            let formatted =
                pretty_format_batches_with_options(&filtered_batches, &options)?;
            if over_limit {
                let mut formatted_str = format!("{formatted}");
                formatted_str = keep_only_maxrows(&formatted_str, maxrows);
                writeln!(writer, "{formatted_str}")?;
            } else {
                writeln!(writer, "{formatted}")?;
            }
        }
        MaxRows::Unlimited => {
            let formatted = pretty_format_batches_with_options(batches, &options)?;
            writeln!(writer, "{formatted}")?;
        }
    }

    Ok(())
}

impl PrintFormat {
    /// Print the batches to a writer using the specified format