feat: add opt-in readable column labels for SQL results

datafusion-cli/src/print_format.rs

@@ -28,6 +28,7 @@ use arrow::record_batch::RecordBatch;
 use arrow::util::pretty::pretty_format_batches_with_options;
 use datafusion::config::FormatOptions;
 use datafusion::error::Result;
+use datafusion_common::metadata::batches_with_display_names;
 
 /// Allow records to be printed in different formats
 #[derive(Debug, PartialEq, Eq, clap::ValueEnum, Clone, Copy)]
@@ -183,6 +184,9 @@ impl PrintFormat {
                 if maxrows == MaxRows::Limited(0) {
                     return Ok(());
                 }
+                // Only the table format shows column labels. Programs parse the
+                // other formats, which keep column names.
+                let batches = batches_with_display_names(&batches)?;
                 format_batches_with_maxrows(writer, &batches, maxrows, format_options)
             }
             Self::Json => batches_to_json!(ArrayWriter, writer, &batches),
@@ -204,8 +208,9 @@ impl PrintFormat {
                     format_options.try_into()?;
 
                 let empty_batch = RecordBatch::new_empty(schema);
+                let batches = batches_with_display_names(&[empty_batch])?;
                 let formatted =
-                    pretty_format_batches_with_options(&[empty_batch], &format_options)?;
+                    pretty_format_batches_with_options(&batches, &format_options)?;
                 writeln!(writer, "{formatted}")?;
             }
             _ => {}
@@ -502,6 +507,29 @@ mod tests {
         assert_eq!(output, "")
     }
 
+    #[test]
+    fn test_print_batches_zero_column_batch_with_rows() {
+        let options = arrow::array::RecordBatchOptions::new().with_row_count(Some(2));
+        let batch = RecordBatch::try_new_with_options(
+            Arc::new(Schema::empty()),
+            vec![],
+            &options,
+        )
+        .unwrap();
+
+        // Table format must not fail on a batch that has rows but no columns
+        let output = PrintBatchesTest::new()
+            .with_format(PrintFormat::Table)
+            .with_schema(Arc::new(Schema::empty()))
+            .with_batches(vec![batch])
+            .run();
+        assert_snapshot!(output, @r"
+        ++
+        ++
+        ++
+        ");
+    }
+
     #[derive(Debug)]
     struct PrintBatchesTest {
         format: PrintFormat,
@@ -651,6 +679,72 @@ mod tests {
         ");
     }
 
+    #[test]
+    fn print_column_labels() {
+        let output = PrintBatchesTest::new()
+            .with_format(PrintFormat::Table)
+            .with_schema(labeled_schema())
+            .with_batches(vec![labeled_batch()])
+            .run();
+        assert_snapshot!(output, @r"
+        +-------+---+
+        | a + 1 | b |
+        +-------+---+
+        | 1     |   |
+        | 2     |   |
+        | 3     |   |
+        +-------+---+
+        ");
+
+        let output = PrintBatchesTest::new()
+            .with_format(PrintFormat::Table)
+            .with_schema(labeled_schema())
+            .with_batches(vec![])
+            .run();
+        assert_snapshot!(output, @r"
+        +-------+---+
+        | a + 1 | b |
+        +-------+---+
+        +-------+---+
+        ");
+
+        // Formats that programs parse keep column names
+        let output = PrintBatchesTest::new()
+            .with_format(PrintFormat::Csv)
+            .with_batches(vec![labeled_batch()])
+            .with_header(WithHeader::Yes)
+            .run();
+        assert_snapshot!(output, @r#"
+        t.a + Int64(1),b
+        1,
+        2,
+        3,
+        "#);
+    }
+
+    /// Return a schema whose first column has a label
+    fn labeled_schema() -> SchemaRef {
+        let name = "t.a + Int64(1)";
+        let labeled = datafusion_common::metadata::column_label_metadata("a + 1", name)
+            .add_to_field(Field::new(name, DataType::Int32, false));
+        Arc::new(Schema::new(vec![
+            labeled,
+            Field::new("b", DataType::Int32, true),
+        ]))
+    }
+
+    /// Return a batch for [`labeled_schema`] with three rows
+    fn labeled_batch() -> RecordBatch {
+        RecordBatch::try_new(
+            labeled_schema(),
+            vec![
+                Arc::new(Int32Array::from(vec![1, 2, 3])),
+                Arc::new(Int32Array::from(vec![None, None, None])),
+            ],
+        )
+        .unwrap()
+    }
+
     /// return a schema with many columns (to exercise wide table formatting)
     fn wide_column_schema() -> SchemaRef {
         let fields: Vec<Field> = (0..10)

datafusion/common/src/config.rs

@@ -324,6 +324,21 @@ config_namespace! {
         /// Disable this option to preserve explicit subquery ordering in the
         /// planned query.
         pub enable_subquery_sort_elimination: bool, default = true
+
+        /// When set to true, each unaliased output column of a top-level SQL
+        /// query gets a readable label, such as `a + 1` for a column named
+        /// `t.a + Int64(1)`. Labels leave out type wrappers and table
+        /// qualifiers, and window functions use `OVER (...)` syntax without
+        /// default parts. The label is stored in the field metadata key
+        /// `datafusion.label`, and `DataFrame::show` and `datafusion-cli`
+        /// table output display it as the column header.
+        ///
+        /// Column names don't change, so name resolution and code that reads
+        /// columns by name work as before. Columns that you name in the
+        /// outermost `SELECT` list, and statements that store column names
+        /// (`CREATE VIEW`, `CREATE TABLE AS`, `INSERT`, `COPY`), get no
+        /// labels.
+        pub column_labels: bool, default = false
     }
 }
 

datafusion/common/src/metadata.rs

@@ -17,7 +17,8 @@
 
 use std::{collections::BTreeMap, sync::Arc};
 
-use arrow::datatypes::{DataType, Field, FieldRef, Metadata};
+use arrow::array::{RecordBatch, RecordBatchOptions};
+use arrow::datatypes::{DataType, Field, FieldRef, Metadata, Schema, SchemaRef};
 use hashbrown::HashMap;
 
 use crate::{DataFusionError, ScalarValue, error::_plan_err};
@@ -135,6 +136,84 @@ pub fn format_type_and_metadata(
     }
 }
 
+/// Field metadata key for the readable label of a query result column, such
+/// as `a + 1` for a column named `t.a + Int64(1)`.
+///
+/// The SQL planner adds this key when `datafusion.sql_parser.column_labels`
+/// is enabled. Tools that display results show the label instead of the
+/// column name. See [`display_name`].
+pub const COLUMN_LABEL_KEY: &str = "datafusion.label";
+
+/// Field metadata key for the column name that [`COLUMN_LABEL_KEY`] labels.
+///
+/// Field metadata follows a column through later aliases, so a renamed
+/// column can still carry an old label. The label applies only while the
+/// field name equals this value.
+pub const COLUMN_LABEL_OF_KEY: &str = "datafusion.label_of";
+
+/// Returns the metadata that labels the column `name` with `label`.
+pub fn column_label_metadata(label: &str, name: &str) -> FieldMetadata {
+    FieldMetadata::from(BTreeMap::from([
+        (COLUMN_LABEL_KEY.to_string(), label.to_string()),
+        (COLUMN_LABEL_OF_KEY.to_string(), name.to_string()),
+    ]))
+}
+
+/// Returns the name to display for `field`: its label (see
+/// [`COLUMN_LABEL_KEY`]) when the label applies to the field's current name,
+/// and the field name otherwise.
+pub fn display_name(field: &Field) -> &str {
+    let metadata = field.metadata();
+    match (
+        metadata.get(COLUMN_LABEL_KEY),
+        metadata.get(COLUMN_LABEL_OF_KEY),
+    ) {
+        (Some(label), Some(label_of)) if label_of == field.name() => label,
+        _ => field.name(),
+    }
+}
+
+/// Returns `schema` with each field renamed to its [`display_name`].
+pub fn schema_with_display_names(schema: &SchemaRef) -> SchemaRef {
+    if schema
+        .fields()
+        .iter()
+        .all(|field| display_name(field) == field.name())
+    {
+        return Arc::clone(schema);
+    }
+    let fields: Vec<FieldRef> = schema
+        .fields()
+        .iter()
+        .map(|field| Arc::new(field.as_ref().clone().with_name(display_name(field))))
+        .collect();
+    Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))
+}
+
+/// Returns `batches` with each column renamed to its [`display_name`], for
+/// printing. The arrays are shared, not copied.
+pub fn batches_with_display_names(
+    batches: &[RecordBatch],
+) -> Result<Vec<RecordBatch>, DataFusionError> {
+    batches
+        .iter()
+        .map(|batch| {
+            let schema = schema_with_display_names(&batch.schema());
+            if Arc::ptr_eq(&schema, &batch.schema()) {
+                return Ok(batch.clone());
+            }
+            // Keep the row count so zero-column batches survive the rebuild.
+            let options =
+                RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+            Ok(RecordBatch::try_new_with_options(
+                schema,
+                batch.columns().to_vec(),
+                &options,
+            )?)
+        })
+        .collect()
+}
+
 /// Literal metadata
 ///
 /// Stores metadata associated with a literal expressions
@@ -392,3 +471,72 @@ impl From<&HashMap<String, String>> for FieldMetadata {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn labeled_field(name: &str, label: &str, label_of: &str) -> Field {
+        let metadata = column_label_metadata(label, label_of);
+        metadata.add_to_field(Field::new(name, DataType::Int64, true))
+    }
+
+    #[test]
+    fn display_name_uses_label_for_its_column() {
+        let field = labeled_field("t.a + Int64(1)", "a + 1", "t.a + Int64(1)");
+        assert_eq!(display_name(&field), "a + 1");
+    }
+
+    #[test]
+    fn display_name_ignores_label_after_rename() {
+        let field = labeled_field("total", "sum(a)", "sum(t.a)");
+        assert_eq!(display_name(&field), "total");
+    }
+
+    #[test]
+    fn display_name_without_label() {
+        let field = Field::new("a", DataType::Int64, true);
+        assert_eq!(display_name(&field), "a");
+    }
+
+    #[test]
+    fn schema_with_display_names_keeps_unlabeled_schema() {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)]));
+        assert!(Arc::ptr_eq(&schema_with_display_names(&schema), &schema));
+    }
+
+    #[test]
+    fn batches_with_display_names_keeps_zero_column_rows() {
+        let options = RecordBatchOptions::new().with_row_count(Some(3));
+        let batch = RecordBatch::try_new_with_options(
+            Arc::new(Schema::empty()),
+            vec![],
+            &options,
+        )
+        .unwrap();
+        let out = batches_with_display_names(&[batch]).unwrap();
+        assert_eq!(out.len(), 1);
+        assert_eq!(out[0].num_columns(), 0);
+        assert_eq!(out[0].num_rows(), 3);
+    }
+
+    #[test]
+    fn batches_with_display_names_keeps_unlabeled_schema_arc() {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)]));
+        let batch = RecordBatch::new_empty(Arc::clone(&schema));
+        let out = batches_with_display_names(&[batch]).unwrap();
+        assert!(Arc::ptr_eq(&out[0].schema(), &schema));
+    }
+
+    #[test]
+    fn batches_with_display_names_relabels_labeled_columns() {
+        let field = labeled_field("t.a + Int64(1)", "a + 1", "t.a + Int64(1)");
+        let schema = Arc::new(Schema::new(vec![field]));
+        let column: arrow::array::ArrayRef =
+            Arc::new(arrow::array::Int64Array::from(vec![1, 2]));
+        let batch = RecordBatch::try_new(schema, vec![column]).unwrap();
+        let out = batches_with_display_names(&[batch]).unwrap();
+        assert_eq!(out[0].schema().field(0).name(), "a + 1");
+        assert_eq!(out[0].num_rows(), 2);
+    }
+}

datafusion/core/src/dataframe/mod.rs

@@ -51,6 +51,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use arrow::util::display::{ArrayFormatter, FormatOptions};
 use arrow_schema::FieldRef;
 use datafusion_common::config::{CsvOptions, JsonOptions};
+use datafusion_common::metadata::batches_with_display_names;
 use datafusion_common::{
     Column, DFSchema, DataFusionError, ParamValues, ScalarValue, SchemaError,
     TableReference, UnnestOptions, exec_err, internal_datafusion_err, not_impl_err,
@@ -1643,7 +1644,7 @@ impl DataFrame {
         let arrow_options =
             arrow_options.with_formatter_factory(Some(&formatter_factory));
 
-        let results = self.collect().await?;
+        let results = batches_with_display_names(&self.collect().await?)?;
         Ok(
             pretty::pretty_format_batches_with_options(&results, &arrow_options)?
                 .to_string(),
@@ -1669,7 +1670,9 @@ impl DataFrame {
     /// ```
     pub async fn show_limit(self, num: usize) -> Result<()> {
         let results = self.limit(0, Some(num))?.collect().await?;
-        Ok(pretty::print_batches(&results)?)
+        Ok(pretty::print_batches(&batches_with_display_names(
+            &results,
+        )?)?)
     }
 
     /// Return a new [`TaskContext`] which would be used to execute this DataFrame