Refactor DataFrame printing to ensure headers are displayed for empty DataFrames

Author: kosiew

python/tests/test_dataframe.py

@@ -1496,11 +1496,12 @@ def test_show_no_batches(capsys):
 
 
 def test_show_empty_dataframe(df, capsys):
-    """Ensure showing an empty DataFrame prints a helpful message."""
+    """Ensure showing an empty DataFrame still prints headers."""
     empty_df = df.limit(0)
     empty_df.show()
     captured = capsys.readouterr()
-    assert "Empty DataFrame" in captured.out
+    assert "| a | b | c |" in captured.out
+    assert "Empty DataFrame" not in captured.out
 
 
 def test_to_polars(df):

src/dataframe.rs

@@ -1097,10 +1097,24 @@ impl Iterator for ArrowStreamReader {
 
 /// Print DataFrame
 fn print_dataframe(py: Python, df: DataFrame) -> PyDataFusionResult<()> {
+    // Get the schema before consuming the DataFrame
+    let schema: SchemaRef = Arc::new(df.schema().clone().into());
+
     // Get string representation of record batches
-    let batches = wait_for_future(py, df.collect())??;
-    let is_empty = batches.is_empty() || batches.iter().all(|b| b.num_rows() == 0);
-    let result = if is_empty {
+    let collected_batches = wait_for_future(py, df.collect())??;
+
+    let batches =
+        if collected_batches.is_empty() || collected_batches.iter().all(|b| b.num_rows() == 0) {
+            if schema.fields().is_empty() {
+                vec![]
+            } else {
+                vec![RecordBatch::new_empty(schema.clone())]
+            }
+        } else {
+            collected_batches
+        };
+
+    let result = if batches.is_empty() {
         "Empty DataFrame".to_string()
     } else {
         match pretty::pretty_format_batches(&batches) {