From 61f8f9974a29e2124167e36e36093176bd5df2ce Mon Sep 17 00:00:00 2001
From: Andrew Lamb <andrew@nerdnetworks.org>
Date: Thu, 27 Feb 2025 20:23:39 -0500
Subject: [PATCH] Add additional protobuf tests

---
 .../tests/cases/roundtrip_physical_plan.rs    | 56 ++++++++++++++++++-
 1 file changed, 55 insertions(+), 1 deletion(-)

diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
index a2506bb318d2..54c4946a2c9a 100644
--- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
+++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
@@ -87,7 +87,7 @@ use datafusion::physical_plan::windows::{
 use datafusion::physical_plan::{
     displayable, ExecutionPlan, InputOrderMode, Partitioning, PhysicalExpr, Statistics,
 };
-use datafusion::prelude::SessionContext;
+use datafusion::prelude::{ParquetReadOptions, SessionContext};
 use datafusion::scalar::ScalarValue;
 use datafusion_common::config::TableParquetOptions;
 use datafusion_common::file_options::csv_writer::CsvWriterOptions;
@@ -157,6 +157,31 @@ fn roundtrip_test_with_context(
     Ok(())
 }
 
+/// Perform a serde roundtrip for the specified sql query, and  assert that
+/// query results are identical.
+async fn roundtrip_test_sql_with_context(sql: &str, ctx: &SessionContext) -> Result<()> {
+    let codec = DefaultPhysicalExtensionCodec {};
+    let initial_plan = ctx.sql(sql).await?.create_physical_plan().await?;
+
+    roundtrip_test_and_return(initial_plan, ctx, &codec)?;
+    Ok(())
+}
+
+/// returns a SessionContext with `alltypes_plain` registered
+async fn all_types_context() -> Result<SessionContext> {
+    let ctx = SessionContext::new();
+
+    let testdata = datafusion::test_util::parquet_test_data();
+    ctx.register_parquet(
+        "alltypes_plain",
+        &format!("{testdata}/alltypes_plain.parquet"),
+        ParquetReadOptions::default(),
+    )
+    .await?;
+
+    Ok(ctx)
+}
+
 #[test]
 fn roundtrip_empty() -> Result<()> {
     roundtrip_test(Arc::new(EmptyExec::new(Arc::new(Schema::empty()))))
@@ -1579,3 +1604,32 @@ async fn roundtrip_coalesce() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn roundtrip_parquet_select_star() -> Result<()> {
+    let ctx = all_types_context().await?;
+    let sql = "select * from alltypes_plain";
+    roundtrip_test_sql_with_context(sql, &ctx).await
+}
+
+#[tokio::test]
+async fn roundtrip_parquet_select_projection() -> Result<()> {
+    let ctx = all_types_context().await?;
+    let sql = "select string_col, timestamp_col from alltypes_plain";
+    roundtrip_test_sql_with_context(sql, &ctx).await
+}
+
+#[tokio::test]
+async fn roundtrip_parquet_select_star_predicate() -> Result<()> {
+    let ctx = all_types_context().await?;
+    let sql = "select * from alltypes_plain where id > 4";
+    roundtrip_test_sql_with_context(sql, &ctx).await
+}
+
+#[ignore = "Test failing due to https://github.com/apache/datafusion/issues/14679"]
+#[tokio::test]
+async fn roundtrip_parquet_select_projection_predicate() -> Result<()> {
+    let ctx = all_types_context().await?;
+    let sql = "select string_col, timestamp_col from alltypes_plain where id > 4";
+    roundtrip_test_sql_with_context(sql, &ctx).await
+}