Skip to content

Commit f47b560

Browse files
authored
Add additional protobuf tests (apache#14924)
1 parent b4c3784 commit f47b560

File tree

1 file changed

+55
-1
lines changed

1 file changed

+55
-1
lines changed

datafusion/proto/tests/cases/roundtrip_physical_plan.rs

+55-1
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ use datafusion::physical_plan::windows::{
8787
use datafusion::physical_plan::{
8888
displayable, ExecutionPlan, InputOrderMode, Partitioning, PhysicalExpr, Statistics,
8989
};
90-
use datafusion::prelude::SessionContext;
90+
use datafusion::prelude::{ParquetReadOptions, SessionContext};
9191
use datafusion::scalar::ScalarValue;
9292
use datafusion_common::config::TableParquetOptions;
9393
use datafusion_common::file_options::csv_writer::CsvWriterOptions;
@@ -157,6 +157,31 @@ fn roundtrip_test_with_context(
157157
Ok(())
158158
}
159159

160+
/// Perform a serde roundtrip for the specified sql query, and assert that
161+
/// query results are identical.
162+
async fn roundtrip_test_sql_with_context(sql: &str, ctx: &SessionContext) -> Result<()> {
163+
let codec = DefaultPhysicalExtensionCodec {};
164+
let initial_plan = ctx.sql(sql).await?.create_physical_plan().await?;
165+
166+
roundtrip_test_and_return(initial_plan, ctx, &codec)?;
167+
Ok(())
168+
}
169+
170+
/// returns a SessionContext with `alltypes_plain` registered
171+
async fn all_types_context() -> Result<SessionContext> {
172+
let ctx = SessionContext::new();
173+
174+
let testdata = datafusion::test_util::parquet_test_data();
175+
ctx.register_parquet(
176+
"alltypes_plain",
177+
&format!("{testdata}/alltypes_plain.parquet"),
178+
ParquetReadOptions::default(),
179+
)
180+
.await?;
181+
182+
Ok(ctx)
183+
}
184+
160185
#[test]
161186
fn roundtrip_empty() -> Result<()> {
162187
roundtrip_test(Arc::new(EmptyExec::new(Arc::new(Schema::empty()))))
@@ -1579,3 +1604,32 @@ async fn roundtrip_coalesce() -> Result<()> {
15791604

15801605
Ok(())
15811606
}
1607+
1608+
#[tokio::test]
1609+
async fn roundtrip_parquet_select_star() -> Result<()> {
1610+
let ctx = all_types_context().await?;
1611+
let sql = "select * from alltypes_plain";
1612+
roundtrip_test_sql_with_context(sql, &ctx).await
1613+
}
1614+
1615+
#[tokio::test]
1616+
async fn roundtrip_parquet_select_projection() -> Result<()> {
1617+
let ctx = all_types_context().await?;
1618+
let sql = "select string_col, timestamp_col from alltypes_plain";
1619+
roundtrip_test_sql_with_context(sql, &ctx).await
1620+
}
1621+
1622+
#[tokio::test]
1623+
async fn roundtrip_parquet_select_star_predicate() -> Result<()> {
1624+
let ctx = all_types_context().await?;
1625+
let sql = "select * from alltypes_plain where id > 4";
1626+
roundtrip_test_sql_with_context(sql, &ctx).await
1627+
}
1628+
1629+
#[ignore = "Test failing due to https://github.com/apache/datafusion/issues/14679"]
1630+
#[tokio::test]
1631+
async fn roundtrip_parquet_select_projection_predicate() -> Result<()> {
1632+
let ctx = all_types_context().await?;
1633+
let sql = "select string_col, timestamp_col from alltypes_plain where id > 4";
1634+
roundtrip_test_sql_with_context(sql, &ctx).await
1635+
}

0 commit comments

Comments
 (0)