_repr_ and _html_repr_ show '... and additional rows' message

[Spaarsh]

…ncated outputs

Which issue does this PR close?

Closes #1026

Rationale for this change

Current implementation of _repr_ and _html_repr_ has no indication for whether the output was truncated or not.

What changes are included in this PR?

The changes override the default output for the _repr_ and _html_repr_ functions and, instead, check if the table has more than 10 rows. If it does, then it prints the first 10 rows along with a ... and additional rows message.

A thing to note is that these changes flatten the batches of the dataframe in the _repr_ function (as already done for the _repr_html_ function). This was done primarily since using .collect twice was causing major performance degradation. Batch flattening does this at 1/10 th of the time.

Are there any user-facing changes?

The users shall now see the ... and additional rows message along with the first 10 rows if the table has more than 10 rows.

[Spaarsh]

This is the new output:

For _repr_

+---------+---------+
| letters | numbers |
+---------+---------+
| A       | 1       |
| B       | 2       |
| C       | 3       |
| D       | 4       |
| E       | 5       |
| F       | 6       |
| G       | 7       |
| H       | 8       |
| I       | 9       |
| J       | 10      |
+---------+---------+
and more...

For _repr_html_

<table border='1'>
<tr><th>letters</th><th>numbers</th></tr>
<tr><td>A</td><td>1</td></tr>
<tr><td>B</td><td>2</td></tr>
<tr><td>C</td><td>3</td></tr>
<tr><td>D</td><td>4</td></tr>
<tr><td>E</td><td>5</td></tr>
<tr><td>F</td><td>6</td></tr>
<tr><td>G</td><td>7</td></tr>
<tr><td>H</td><td>8</td></tr>
<tr><td>I</td><td>9</td></tr>
<tr><td>J</td><td>10</td></tr>
<tr><td colspan="100%">and more...</td></tr>
</table>

Performance Changes

Implementation	_repr_	_html_repr_
Current	3.16ms	1.35ms
New	5.5ms	4.5ms

Note: I have manually ran the command multiple times and observed this. If required, I will run a script and produce the average for both cases

[kosiew]

You can simplify batches_as_string to:

        // First get just the first 10 rows
        let preview_df = self.df.as_ref().clone().limit(0, Some(10))?;
        let preview_batches = wait_for_future(py, preview_df.collect())?;

        // Check if there are more rows by trying to get the 11th row
        let has_more_rows = {
            let check_df = self.df.as_ref().clone().limit(10, Some(1))?;
            let check_batch = wait_for_future(py, check_df.collect())?;
            !check_batch.is_empty()
        };

        let batches_as_string = pretty::pretty_format_batches(&preview_batches);

This directly retrieves just the first 10 rows, eliminating the need for manual row tracking and slicing.

[Spaarsh]

I did try this initiatially but calling collect twice led to a severe performance degradation. It used to take 50ms. With the manual slicing, it dropped to 5ms.

You can check my initial suggestion for the same here

[kosiew]

calling collect twice led to a severe performance degradation

I ran this test to compare the performance:

import pyarrow as pa
from datafusion import (
    SessionContext,
)
import time


def run_dataframe_repr_long() -> None:
    ctx = SessionContext()
    # Create a DataFrame with more than 10 rows
    batch = pa.RecordBatch.from_arrays(
        [
            pa.array(list(range(15))),
            pa.array([x * 2 for x in range(15)]),
            pa.array([x * 3 for x in range(15)]),
        ],
        names=["a", "b", "c"],
    )
    df = ctx.create_dataframe([[batch]])

    output = repr(df)


def average_runtime(func, runs=100):
    total_time = 0
    for _ in range(runs):
        start_time = time.time()
        func()
        end_time = time.time()
        total_time += end_time - start_time
    return total_time / runs


average_time = average_runtime(run_dataframe_repr_long)
print(f"Average runtime over {100} runs: {average_time:.6f} seconds")

and found no significant difference:

pr_1041 - is the branch with one collect
amended_pr_1041 - is the branch with two collect

[Spaarsh]

That's weird. Maybe some artifact of my system settings? If there is no performance issues than I'll use your approach. But then why was the _repr_html_ using batch manipulation at the first place? I took the idea from that function!

[kosiew]

using has_more_rows from above

+                if has_more_rows {

[kosiew]

How about simplifying to:

            let rows_remaining = 10 - total_rows;
            let rows_in_batch = batch.num_rows().min(rows_remaining);

            for row in 0..rows_in_batch {
                html_str.push_str("<tr>");
                for col in batch.columns() {
                    let formatter =
                        ArrayFormatter::try_new(col.as_ref(), &FormatOptions::default())?;
                    html_str.push_str("<td>");
                    html_str.push_str(&formatter.value(row).to_string());
                    html_str.push_str("</td>");
                }
                html_str.push_str("</tr>\n");
            }

            total_rows += rows_in_batch;

Reasons:

1. More Accurate Row Limiting:

Before: total_rows was updated before checking the row limit, which could result in processing extra rows unnecessarily.
After: rows_remaining = 10 - total_rows ensures that we never exceed the row limit.

2. Avoids Redundant Vec Allocation:

Before: Each row was constructed using a Vec, and format!() was used for each cell.
After: Directly appends elements to html_str, eliminating unnecessary heap allocations.

3. Simplified and More Efficient Row Processing:

Before:
Used .map() and .collect() to create a list of ArrayFormatters before processing rows.
After:
Retrieves and formats values inside the loop, reducing redundant processing.

4. Avoids Unnecessary break Condition:

Before: Explicit if total_rows >= 10 { break; } was used to stop processing.
After: The min(rows_remaining, batch.num_rows()) logic naturally prevents extra iterations

[kevinjqliu]

btw #1036 also changes _repr_html_