Add string formatter

Author: timsaucer

python/datafusion/html_formatter.py

@@ -26,6 +26,8 @@
     runtime_checkable,
 )
 
+from datafusion._internal import DataFrame as DataFrameInternal
+
 
 def _validate_positive_int(value: Any, param_name: str) -> None:
     """Validate that a parameter is a positive integer.
@@ -345,6 +347,32 @@ def format_html(
 
         return "\n".join(html)
 
+    def format_str(
+        self,
+        batches: list,
+        schema: Any,
+        has_more: bool = False,
+        table_uuid: str | None = None,
+    ) -> str:
+        """Format record batches as a string.
+
+        This method is used by DataFrame's __repr__ implementation and can be
+        called directly when string rendering is needed.
+
+        Args:
+            batches: List of Arrow RecordBatch objects
+            schema: Arrow Schema object
+            has_more: Whether there are more batches not shown
+            table_uuid: Unique ID for the table, used for JavaScript interactions
+
+        Returns:
+            String representation of the data
+
+        Raises:
+            TypeError: If schema is invalid and no batches are provided
+        """
+        return DataFrameInternal.default_str_repr(batches, schema, has_more, table_uuid)
+
     def _build_html_header(self) -> list[str]:
         """Build the HTML header with CSS styles."""
         html = []

src/dataframe.rs

@@ -24,6 +24,7 @@ use arrow::compute::can_cast_types;
 use arrow::error::ArrowError;
 use arrow::ffi::FFI_ArrowSchema;
 use arrow::ffi_stream::FFI_ArrowArrayStream;
+use arrow::pyarrow::FromPyArrow;
 use datafusion::arrow::datatypes::Schema;
 use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow};
 use datafusion::arrow::util::pretty;
@@ -295,6 +296,46 @@ impl PyDataFrame {
     pub fn new(df: DataFrame) -> Self {
         Self { df: Arc::new(df) }
     }
+
+    fn prepare_repr_string(&self, py: Python, as_html: bool) -> PyDataFusionResult<String> {
+        // Get the Python formatter and config
+        let PythonFormatter { formatter, config } = get_python_formatter_with_config(py)?;
+        let (batches, has_more) = wait_for_future(
+            py,
+            collect_record_batches_to_display(self.df.as_ref().clone(), config),
+        )??;
+        if batches.is_empty() {
+            // This should not be reached, but do it for safety since we index into the vector below
+            return Ok("No data to display".to_string());
+        }
+
+        let table_uuid = uuid::Uuid::new_v4().to_string();
+
+        // Convert record batches to PyObject list
+        let py_batches = batches
+            .into_iter()
+            .map(|rb| rb.to_pyarrow(py))
+            .collect::<PyResult<Vec<PyObject>>>()?;
+
+        let py_schema = self.schema().into_pyobject(py)?;
+
+        let kwargs = pyo3::types::PyDict::new(py);
+        let py_batches_list = PyList::new(py, py_batches.as_slice())?;
+        kwargs.set_item("batches", py_batches_list)?;
+        kwargs.set_item("schema", py_schema)?;
+        kwargs.set_item("has_more", has_more)?;
+        kwargs.set_item("table_uuid", table_uuid)?;
+
+        let method_name = match as_html {
+            true => "format_html",
+            false => "format_str",
+        };
+
+        let html_result = formatter.call_method(method_name, (), Some(&kwargs))?;
+        let html_str: String = html_result.extract()?;
+
+        Ok(html_str)
+    }
 }
 
 #[pymethods]
@@ -321,18 +362,27 @@ impl PyDataFrame {
     }
 
     fn __repr__(&self, py: Python) -> PyDataFusionResult<String> {
-        // Get the Python formatter config
-        let PythonFormatter {
-            formatter: _,
-            config,
-        } = get_python_formatter_with_config(py)?;
-        let (batches, has_more) = wait_for_future(
-            py,
-            collect_record_batches_to_display(self.df.as_ref().clone(), config),
-        )??;
+        self.prepare_repr_string(py, false)
+    }
+
+    #[staticmethod]
+    #[expect(unused_variables)]
+    fn default_str_repr<'py>(
+        batches: Vec<Bound<'py, PyAny>>,
+        schema: &Bound<'py, PyAny>,
+        has_more: bool,
+        table_uuid: &str,
+    ) -> PyResult<String> {
+        let batches = batches
+            .into_iter()
+            .map(|batch| RecordBatch::from_pyarrow_bound(&batch))
+            .collect::<PyResult<Vec<RecordBatch>>>()?
+            .into_iter()
+            .filter(|batch| batch.num_rows() > 0)
+            .collect::<Vec<_>>();
+
         if batches.is_empty() {
-            // This should not be reached, but do it for safety since we index into the vector below
-            return Ok("No data to display".to_string());
+            return Ok("No data to display".to_owned());
         }
 
         let batches_as_displ =
@@ -347,38 +397,7 @@ impl PyDataFrame {
     }
 
     fn _repr_html_(&self, py: Python) -> PyDataFusionResult<String> {
-        // Get the Python formatter and config
-        let PythonFormatter { formatter, config } = get_python_formatter_with_config(py)?;
-        let (batches, has_more) = wait_for_future(
-            py,
-            collect_record_batches_to_display(self.df.as_ref().clone(), config),
-        )??;
-        if batches.is_empty() {
-            // This should not be reached, but do it for safety since we index into the vector below
-            return Ok("No data to display".to_string());
-        }
-
-        let table_uuid = uuid::Uuid::new_v4().to_string();
-
-        // Convert record batches to PyObject list
-        let py_batches = batches
-            .into_iter()
-            .map(|rb| rb.to_pyarrow(py))
-            .collect::<PyResult<Vec<PyObject>>>()?;
-
-        let py_schema = self.schema().into_pyobject(py)?;
-
-        let kwargs = pyo3::types::PyDict::new(py);
-        let py_batches_list = PyList::new(py, py_batches.as_slice())?;
-        kwargs.set_item("batches", py_batches_list)?;
-        kwargs.set_item("schema", py_schema)?;
-        kwargs.set_item("has_more", has_more)?;
-        kwargs.set_item("table_uuid", table_uuid)?;
-
-        let html_result = formatter.call_method("format_html", (), Some(&kwargs))?;
-        let html_str: String = html_result.extract()?;
-
-        Ok(html_str)
+        self.prepare_repr_string(py, true)
     }
 
     /// Calculate summary statistics for a DataFrame