feat: Add method to configure DataFrame display options in SessionConfig (python)

- Introduced `with_dataframe_display_config` method in `SessionConfig` to allow customization of DataFrame display settings.
- Parameters include `max_table_bytes`, `min_table_rows`, `max_cell_length`, and `max_table_rows_in_repr` for flexible display configurations.
- Utilizes `DataframeDisplayConfig` for internal management of display settings.

Author: kosiew

python/datafusion/context.py

@@ -36,6 +36,7 @@
 from ._internal import SessionConfig as SessionConfigInternal
 from ._internal import SessionContext as SessionContextInternal
 from ._internal import SQLOptions as SQLOptionsInternal
+from ._internal import DataframeDisplayConfig as DataframeDisplayConfigInternal
 
 if TYPE_CHECKING:
     import pathlib
@@ -89,6 +90,37 @@ def __init__(self, config_options: dict[str, str] | None = None) -> None:
         """
         self.config_internal = SessionConfigInternal(config_options)
 
+    def with_dataframe_display_config(
+        self,
+        max_table_bytes: int = None,
+        min_table_rows: int = None,
+        max_cell_length: int = None,
+        max_table_rows_in_repr: int = None,
+    ) -> SessionConfig:
+        """Configure the display options for DataFrames.
+
+        Args:
+            max_table_bytes: Maximum bytes to display for table presentation (default: 2MB)
+            min_table_rows: Minimum number of table rows to display (default: 20)
+            max_cell_length: Maximum length of a cell before it gets minimized (default: 25)
+            max_table_rows_in_repr: Maximum number of rows to display in repr string output (default: 10)
+
+        Returns:
+            A new :py:class:`SessionConfig` object with the updated display settings.
+        """
+
+        display_config = DataframeDisplayConfigInternal(
+            max_table_bytes=max_table_bytes,
+            min_table_rows=min_table_rows,
+            max_cell_length=max_cell_length,
+            max_table_rows_in_repr=max_table_rows_in_repr,
+        )
+
+        self.config_internal = self.config_internal.with_dataframe_display_config(
+            display_config
+        )
+        return self
+
     def with_create_default_catalog_and_schema(
         self, enabled: bool = True
     ) -> SessionConfig:
@@ -806,9 +838,11 @@ def register_parquet(
             file_extension,
             skip_metadata,
             schema,
-            [sort_list_to_raw_sort_list(exprs) for exprs in file_sort_order]
-            if file_sort_order is not None
-            else None,
+            (
+                [sort_list_to_raw_sort_list(exprs) for exprs in file_sort_order]
+                if file_sort_order is not None
+                else None
+            ),
         )
 
     def register_csv(

src/context.rs

@@ -73,54 +73,6 @@ use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple, PyType};
 use tokio::task::JoinHandle;
 
 /// Configuration for displaying DataFrames
-#[pyclass(name = "DataframeDisplayConfig", module = "datafusion", subclass)]
-#[derive(Clone)]
-pub struct DataframeDisplayConfig {
-    /// Maximum bytes to display for table presentation (default: 2MB)
-    #[pyo3(get, set)]
-    pub max_table_bytes: usize,
-    /// Minimum number of table rows to display (default: 20)
-    #[pyo3(get, set)]
-    pub min_table_rows: usize,
-    /// 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 DataframeDisplayConfig {
-    #[new]
-    #[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 = Self::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),
-        }
-    }
-}
-
-impl Default for DataframeDisplayConfig {
-    fn default() -> Self {
-        Self {
-            max_table_bytes: 2 * 1024 * 1024, // 2 MB
-            min_table_rows: 20,
-            max_cell_length: 25,
-            max_table_rows_in_repr: 10,
-        }
-    }
-}
 
 /// Configuration options for a SessionContext
 #[pyclass(name = "SessionConfig", module = "datafusion", subclass)]
@@ -229,6 +181,55 @@ impl PySessionConfig {
     }
 }
 
+#[pyclass(name = "DataframeDisplayConfig", module = "datafusion", subclass)]
+#[derive(Clone)]
+pub struct DataframeDisplayConfig {
+    /// Maximum bytes to display for table presentation (default: 2MB)
+    #[pyo3(get, set)]
+    pub max_table_bytes: usize,
+    /// Minimum number of table rows to display (default: 20)
+    #[pyo3(get, set)]
+    pub min_table_rows: usize,
+    /// 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 DataframeDisplayConfig {
+    #[new]
+    #[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 = Self::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),
+        }
+    }
+}
+
+impl Default for DataframeDisplayConfig {
+    fn default() -> Self {
+        Self {
+            max_table_bytes: 2 * 1024 * 1024, // 2 MB
+            min_table_rows: 20,
+            max_cell_length: 25,
+            max_table_rows_in_repr: 10,
+        }
+    }
+}
+
 /// Runtime options for a SessionContext
 #[pyclass(name = "RuntimeEnvBuilder", module = "datafusion", subclass)]
 #[derive(Clone)]

src/lib.rs

@@ -82,6 +82,7 @@ fn _internal(py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<context::PyRuntimeEnvBuilder>()?;
     m.add_class::<context::PySessionConfig>()?;
     m.add_class::<context::PySessionContext>()?;
+    m.add_class::<context::DataframeDisplayConfig>()?;
     m.add_class::<context::PySQLOptions>()?;
     m.add_class::<dataframe::PyDataFrame>()?;
     m.add_class::<udf::PyScalarUDF>()?;