Add test for logical plan pretty-printing

Author: tschwarzinger

datafusion/core/tests/extension_types/pretty_print_data_frames.rs

@@ -4,16 +4,18 @@ use arrow_schema::{DataType, Field, Schema, SchemaRef};
 use datafusion::dataframe::DataFrame;
 use datafusion::error::Result;
 use datafusion::prelude::SessionContext;
+use datafusion_common::metadata::FieldMetadata;
+use datafusion_common::ScalarValue;
+use datafusion_expr::{col, lit_with_metadata};
 use insta::assert_snapshot;
 use std::sync::Arc;
 
 fn test_schema() -> SchemaRef {
-    Arc::new(Schema::new(vec![Field::new(
-        "my_uuids",
-        DataType::FixedSizeBinary(16),
-        false,
-    )
-    .with_extension_type(Uuid)]))
+    Arc::new(Schema::new(vec![uuid_field()]))
+}
+
+fn uuid_field() -> Field {
+    Field::new("my_uuids", DataType::FixedSizeBinary(16), false).with_extension_type(Uuid)
 }
 
 async fn create_test_table() -> Result<DataFrame> {
@@ -36,7 +38,7 @@ async fn create_test_table() -> Result<DataFrame> {
 }
 
 #[tokio::test]
-async fn test_pretty_print_extension_types() -> Result<()> {
+async fn test_pretty_print_logical_plan() -> Result<()> {
     let result = create_test_table().await?.to_string().await?;
 
     assert_snapshot!(
@@ -53,3 +55,27 @@ async fn test_pretty_print_extension_types() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_pretty_print_extension_types() -> Result<()> {
+    let scalar = ScalarValue::FixedSizeBinary(
+        16,
+        Some(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 5, 6]),
+    );
+    let literal =
+        lit_with_metadata(scalar, Some(FieldMetadata::new_from_field(&uuid_field())));
+
+    let df = create_test_table()
+        .await?
+        .filter(col("my_uuids").eq(literal))?;
+
+    assert_snapshot!(
+        df.logical_plan(),
+        @r#"
+    Filter: test.my_uuids = FixedSizeBinary(16, "0,1,2,3,4,5,6,7,8,9,0,1,2,3,5,6") FieldMetadata { inner: {"ARROW:extension:name": "arrow.uuid"} }
+      TableScan: test
+    "#
+    );
+
+    Ok(())
+}