feat: Add max_table_rows_in_repr to control row display in DataFrame

- Updated DataFrame class to include max_table_rows_in_repr parameter for display configuration.
- Enhanced configure_display method to accept max_table_rows_in_repr.
- Modified DisplayConfig struct to include max_table_rows_in_repr with a default value of 10.
- Added tests to verify the functionality of max_table_rows_in_repr in both configuration and display output.

Author: kosiew

python/datafusion/dataframe.py

@@ -818,6 +818,7 @@ def configure_display(
         max_table_bytes: Optional[int] = None,
         min_table_rows: Optional[int] = None,
         max_cell_length: Optional[int] = None,
+        max_table_rows_in_repr: Optional[int] = None,
     ) -> None:
         """Configure display options for DataFrame representation.
 
@@ -828,17 +829,26 @@ def configure_display(
                             This is used for initial display and in notebooks.
             max_cell_length: Maximum length of a cell before it gets minimized (default: 25).
                              Longer cells will be truncated with an expand button.
+            max_table_rows_in_repr: Maximum number of rows to display in string representation
+                                   (default: 10).
 
         Raises:
             ValueError: If any of the provided values are less than or equal to 0.
         """
         if any(
             value is not None and value <= 0
-            for value in (max_table_bytes, min_table_rows, max_cell_length)
+            for value in (
+                max_table_bytes,
+                min_table_rows,
+                max_cell_length,
+                max_table_rows_in_repr,
+            )
         ):
             raise ValueError("All values must be greater than 0.")
 
-        self.df.configure_display(max_table_bytes, min_table_rows, max_cell_length)
+        self.df.configure_display(
+            max_table_bytes, min_table_rows, max_cell_length, max_table_rows_in_repr
+        )
 
     def reset_display_config(self) -> None:
         """Reset display configuration to default values."""

python/tests/test_dataframe.py

@@ -1271,27 +1271,33 @@ def test_display_config(df):
     assert config.max_table_bytes == 2 * 1024 * 1024  # 2 MB
     assert config.min_table_rows == 20
     assert config.max_cell_length == 25
+    assert config.max_table_rows_in_repr == 10  # Verify the new property
 
 
 def test_configure_display(df):
     """Test setting display configuration properties."""
     # Modify the display configuration
     df.configure_display(
-        max_table_bytes=1024 * 1024, min_table_rows=10, max_cell_length=50  # 1 MB
+        max_table_bytes=1024 * 1024,
+        min_table_rows=10,
+        max_cell_length=50,
+        max_table_rows_in_repr=15,  # Add test for the new property
     )
 
     # Verify the changes took effect
     config = df.display_config
     assert config.max_table_bytes == 1024 * 1024  # 1 MB
     assert config.min_table_rows == 10
     assert config.max_cell_length == 50
+    assert config.max_table_rows_in_repr == 15
 
     # Test partial update (only changing one property)
-    df.configure_display(min_table_rows=5)
+    df.configure_display(max_table_rows_in_repr=5)
     config = df.display_config
     assert config.max_table_bytes == 1024 * 1024  # previous value retained
-    assert config.min_table_rows == 5  # only this value changed
+    assert config.min_table_rows == 10  # previous value retained
     assert config.max_cell_length == 50  # previous value retained
+    assert config.max_table_rows_in_repr == 5  # only this value changed
 
     # Test with extreme values (still valid, but potentially problematic)
     # Zero values
@@ -1490,3 +1496,40 @@ def _create_numeric_test_df(ctx, rows) -> DataFrame:
     data = list(range(rows))
     batch = pa.RecordBatch.from_arrays([pa.array(data)], names=["values"])
     return ctx.create_dataframe([[batch]])
+
+
+def test_max_table_rows_in_repr(ctx):
+    """Test that max_table_rows_in_repr controls the number of rows in string representation."""
+    # Create a dataframe with more rows than the default max_table_rows_in_repr (10)
+    rows = 20
+    df = _create_numeric_test_df(ctx, rows)
+
+    # First test with default setting (should limit to 10 rows)
+    repr_str = df.__repr__()
+    lines_default = _count_lines_in_str(repr_str)
+
+    # Default should be 10 rows max
+    assert lines_default <= 10
+    assert "Data truncated" in repr_str
+
+    # Now set a custom max_table_rows_in_repr value
+    custom_max_rows = 15
+    df.configure_display(max_table_rows_in_repr=custom_max_rows)
+
+    # Get the string representation with new configuration
+    repr_str_more = df.__repr__()
+    lines_custom = _count_lines_in_str(repr_str_more)
+
+    # Should show more rows than default but not more than configured max
+    assert lines_custom > lines_default
+    assert lines_custom <= custom_max_rows
+    assert "Data truncated" in repr_str_more
+
+    # Now set max_rows higher than total rows - should show all rows
+    df.configure_display(max_table_rows_in_repr=25)
+    repr_str_all = df.__repr__()
+    lines_all = _count_lines_in_str(repr_str_all)
+
+    # Should show all rows (20)
+    assert lines_all == rows
+    assert "Data truncated" not in repr_str_all

src/dataframe.rs

@@ -86,22 +86,28 @@ pub struct DisplayConfig {
     /// Maximum length of a cell before it gets minimized (default: 25)
     #[pyo3(get, set)]
     pub max_cell_length: usize,
+    /// Maximum number of rows to display in repr string output (default: 10)
+    #[pyo3(get, set)]
+    pub max_table_rows_in_repr: usize,
 }
 
 #[pymethods]
 impl DisplayConfig {
     #[new]
-    #[pyo3(signature = (max_table_bytes=None, min_table_rows=None, max_cell_length=None))]
+    #[pyo3(signature = (max_table_bytes=None, min_table_rows=None, max_cell_length=None, max_table_rows_in_repr=None))]
     fn new(
         max_table_bytes: Option<usize>,
         min_table_rows: Option<usize>,
         max_cell_length: Option<usize>,
+        max_table_rows_in_repr: Option<usize>,
     ) -> Self {
         let default = DisplayConfig::default();
         Self {
             max_table_bytes: max_table_bytes.unwrap_or(default.max_table_bytes),
             min_table_rows: min_table_rows.unwrap_or(default.min_table_rows),
             max_cell_length: max_cell_length.unwrap_or(default.max_cell_length),
+            max_table_rows_in_repr: max_table_rows_in_repr
+                .unwrap_or(default.max_table_rows_in_repr),
         }
     }
 }
@@ -112,6 +118,7 @@ impl Default for DisplayConfig {
             max_table_bytes: 2 * 1024 * 1024, // 2 MB
             min_table_rows: 20,
             max_cell_length: 25,
+            max_table_rows_in_repr: 10,
         }
     }
 }
@@ -165,7 +172,7 @@ impl PyDataFrame {
             collect_record_batches_to_display(
                 self.df.as_ref().clone(),
                 self.config.min_table_rows,
-                10,
+                self.config.max_table_rows_in_repr,
                 &self.config,
             ),
         )?;
@@ -858,12 +865,18 @@ impl PyDataFrame {
     }
 
     /// Update display configuration
-    #[pyo3(signature = (max_table_bytes=None, min_table_rows=None, max_cell_length=None))]
+    #[pyo3(signature = (
+        max_table_bytes=None,
+        min_table_rows=None,
+        max_cell_length=None,
+        max_table_rows_in_repr=None
+    ))]
     fn configure_display(
         &mut self,
         max_table_bytes: Option<usize>,
         min_table_rows: Option<usize>,
         max_cell_length: Option<usize>,
+        max_table_rows_in_repr: Option<usize>,
     ) {
         let mut new_config = (*self.config).clone();
 
@@ -879,6 +892,10 @@ impl PyDataFrame {
             new_config.max_cell_length = length;
         }
 
+        if let Some(rows) = max_table_rows_in_repr {
+            new_config.max_table_rows_in_repr = rows;
+        }
+
         self.config = Arc::new(new_config);
     }