implement display for ColumnarValue (apache#14220)

* implement display for ColumnarValue * add tests * avoid panic when impl Display * align the strings in tests --------- Co-authored-by: Andrew Lamb <[email protected]>
himadripal · Jan 24, 2025 · 774d3cb · 774d3cb
1 parent ab16fd5
commit 774d3cb
Showing 1 changed file with 61 additions and 1 deletion.
diff --git a/datafusion/expr-common/src/columnar_value.rs b/datafusion/expr-common/src/columnar_value.rs
@@ -20,8 +20,10 @@
 use arrow::array::{Array, ArrayRef, NullArray};
 use arrow::compute::{kernels, CastOptions};
 use arrow::datatypes::DataType;
+use arrow::util::pretty::pretty_format_columns;
 use datafusion_common::format::DEFAULT_CAST_OPTIONS;
 use datafusion_common::{internal_err, Result, ScalarValue};
+use std::fmt;
 use std::sync::Arc;
 
 /// The result of evaluating an expression.
@@ -218,9 +220,34 @@ impl ColumnarValue {
     }
 }
 
+// Implement Display trait for ColumnarValue
+impl fmt::Display for ColumnarValue {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let formatted = match self {
+            ColumnarValue::Array(array) => {
+                pretty_format_columns("ColumnarValue(ArrayRef)", &[Arc::clone(array)])
+            }
+            ColumnarValue::Scalar(_) => {
+                if let Ok(array) = self.to_array(1) {
+                    pretty_format_columns("ColumnarValue(ScalarValue)", &[array])
+                } else {
+                    return write!(f, "Error formatting columnar value");
+                }
+            }
+        };
+
+        if let Ok(formatted) = formatted {
+            write!(f, "{}", formatted)
+        } else {
+            write!(f, "Error formatting columnar value")
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
+    use arrow::array::Int32Array;
 
     #[test]
     fn values_to_arrays() {
@@ -329,6 +356,39 @@ mod tests {
 
     /// Makes an array of length `len` with all elements set to `val`
     fn make_array(val: i32, len: usize) -> ArrayRef {
-        Arc::new(arrow::array::Int32Array::from(vec![val; len]))
+        Arc::new(Int32Array::from(vec![val; len]))
+    }
+
+    #[test]
+    fn test_display_scalar() {
+        let column = ColumnarValue::from(ScalarValue::from("foo"));
+        assert_eq!(
+            column.to_string(),
+            concat!(
+                "+----------------------------+\n",
+                "| ColumnarValue(ScalarValue) |\n",
+                "+----------------------------+\n",
+                "| foo                        |\n",
+                "+----------------------------+"
+            )
+        );
+    }
+
+    #[test]
+    fn test_display_array() {
+        let array: ArrayRef = Arc::new(Int32Array::from_iter_values(vec![1, 2, 3]));
+        let column = ColumnarValue::from(array);
+        assert_eq!(
+            column.to_string(),
+            concat!(
+                "+-------------------------+\n",
+                "| ColumnarValue(ArrayRef) |\n",
+                "+-------------------------+\n",
+                "| 1                       |\n",
+                "| 2                       |\n",
+                "| 3                       |\n",
+                "+-------------------------+"
+            )
+        );
     }
 }