This is an automated email from the ASF dual-hosted git repository.
milenkovicm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new d92903f8 chore: add read/write roundtrip tests (#1249)
d92903f8 is described below
commit d92903f8f335eee3cd77fd1a645d7c558de30ad8
Author: Marko Milenković <[email protected]>
AuthorDate: Tue Apr 29 09:53:08 2025 +0100
chore: add read/write roundtrip tests (#1249)
* chore: add read/write roundtrip tests
* add additional tests to cover with_state cases
---
ballista/client/tests/common/mod.rs | 21 ++++-
ballista/client/tests/context_checks.rs | 139 +++++++++++++++++++++++++++++++-
2 files changed, 155 insertions(+), 5 deletions(-)
diff --git a/ballista/client/tests/common/mod.rs
b/ballista/client/tests/common/mod.rs
index 1d2bca94..800c21ea 100644
--- a/ballista/client/tests/common/mod.rs
+++ b/ballista/client/tests/common/mod.rs
@@ -25,7 +25,7 @@ use ballista_core::serde::{
};
use ballista_core::{ConfigProducer, RuntimeProducer};
use ballista_scheduler::SessionBuilder;
-use datafusion::execution::SessionState;
+use datafusion::execution::{SessionState, SessionStateBuilder};
use datafusion::prelude::{SessionConfig, SessionContext};
/// Returns the parquet test data directory, which is by default
@@ -234,14 +234,27 @@ pub async fn remote_context() -> SessionContext {
.unwrap()
}
+#[allow(dead_code)]
+pub async fn standalone_context_with_state() -> SessionContext {
+ let state = SessionStateBuilder::new().with_default_features().build();
+ SessionContext::standalone_with_state(state).await.unwrap()
+}
+
+#[allow(dead_code)]
+pub async fn remote_context_with_state() -> SessionContext {
+ let state = SessionStateBuilder::new().with_default_features().build();
+ let (host, port) = setup_test_cluster_with_state(state.clone()).await;
+ SessionContext::remote_with_state(&format!("df://{host}:{port}"), state)
+ .await
+ .unwrap()
+}
+
#[ctor::ctor]
fn init() {
// Enable RUST_LOG logging configuration for test
let _ = env_logger::builder()
.filter_level(log::LevelFilter::Info)
- .parse_filters(
-
"ballista=debug,ballista_scheduler-rs=debug,ballista_executor=debug",
- )
+
.parse_filters("ballista=debug,ballista_scheduler=debug,ballista_executor=debug")
//.parse_filters("ballista=debug,ballista_scheduler-rs=debug,ballista_executor=debug,datafusion=debug")
.is_test(true)
.try_init();
diff --git a/ballista/client/tests/context_checks.rs
b/ballista/client/tests/context_checks.rs
index f60559ac..026b83c2 100644
--- a/ballista/client/tests/context_checks.rs
+++ b/ballista/client/tests/context_checks.rs
@@ -19,7 +19,10 @@ mod common;
#[cfg(test)]
mod supported {
- use crate::common::{remote_context, standalone_context};
+ use crate::common::{
+ remote_context, remote_context_with_state, standalone_context,
+ standalone_context_with_state,
+ };
use ballista_core::config::BallistaConfig;
use datafusion::physical_plan::collect;
use datafusion::prelude::*;
@@ -612,4 +615,138 @@ mod supported {
assert_batches_eq!(expected, &result);
}
+
+ #[rstest]
+ #[case::standalone(standalone_context())]
+ #[case::remote(remote_context())]
+ #[case::standalone_state(standalone_context_with_state())]
+ #[case::remote_state(remote_context_with_state())]
+ #[tokio::test]
+ #[cfg(not(windows))] // test is failing at windows, can't debug it
+ async fn should_execute_sql_write_read_roundtrip(
+ #[future(awt)]
+ #[case]
+ ctx: SessionContext,
+ test_data: String,
+ ) -> datafusion::error::Result<()> {
+ ctx.register_parquet(
+ "test",
+ &format!("{test_data}/alltypes_plain.parquet"),
+ Default::default(),
+ )
+ .await?;
+
+ let expected = [
+ "+----+------------+",
+ "| id | string_col |",
+ "+----+------------+",
+ "| 5 | 31 |",
+ "| 6 | 30 |",
+ "| 7 | 31 |",
+ "+----+------------+",
+ ];
+
+ let write_dir = tempfile::tempdir().expect("temporary directory to be
created");
+ let write_dir_path = write_dir
+ .path()
+ .to_str()
+ .expect("path to be converted to str");
+
+ ctx.sql("select * from test")
+ .await?
+ .write_parquet(write_dir_path, Default::default(),
Default::default())
+ .await?;
+
+ ctx.register_parquet("p_written_table", write_dir_path,
Default::default())
+ .await?;
+
+ let result = ctx
+ .sql("select id, string_col from p_written_table where id > 4")
+ .await?
+ .collect()
+ .await?;
+
+ assert_batches_eq!(expected, &result);
+
+ ctx.sql("select * from test")
+ .await?
+ .write_csv(write_dir_path, Default::default(), Default::default())
+ .await?;
+
+ ctx.register_csv("c_written_table", write_dir_path, Default::default())
+ .await?;
+
+ let result = ctx
+ .sql("select id, string_col from c_written_table where id > 4")
+ .await?
+ .collect()
+ .await?;
+
+ assert_batches_eq!(expected, &result);
+
+ //
+ // TODO: enable when json supports lands in datafusion 47
+ //
+ // ctx.sql("select * from test")
+ // .await?
+ // .write_json(write_dir_path, Default::default(),
Default::default())
+ // .await?;
+
+ // ctx.register_json("j_written_table", write_dir_path,
Default::default())
+ // .await?;
+
+ // let result = ctx
+ // .sql("select id, string_col from j_written_table where id > 4")
+ // .await?
+ // .collect()
+ // .await?;
+
+ // assert_batches_eq!(expected, &result);
+
+ Ok(())
+ }
+
+ #[rstest]
+ #[case::standalone(standalone_context())]
+ #[case::remote(remote_context())]
+ #[case::standalone_state(standalone_context_with_state())]
+ #[case::remote_state(remote_context_with_state())]
+ #[tokio::test]
+ async fn should_execute_sql_show_multiple_times(
+ #[future(awt)]
+ #[case]
+ ctx: SessionContext,
+ test_data: String,
+ ) -> datafusion::error::Result<()> {
+ ctx.register_parquet(
+ "test",
+ &format!("{test_data}/alltypes_plain.parquet"),
+ Default::default(),
+ )
+ .await?;
+
+ let expected = [
+ "+------------+---------------------+",
+ "| string_col | timestamp_col |",
+ "+------------+---------------------+",
+ "| 31 | 2009-03-01T00:01:00 |",
+ "| 30 | 2009-04-01T00:00:00 |",
+ "| 31 | 2009-04-01T00:01:00 |",
+ "+------------+---------------------+",
+ ];
+ // there were cases when we break session context
+ // with standalone setup. so subsequent query for the
+ // same table fails with table does not exist
+ for _ in 0..5 {
+ let result = ctx
+ .sql("select string_col, timestamp_col from test where id > 4")
+ .await?
+ .collect()
+ .await?;
+
+ assert_batches_eq!(expected, &result);
+ }
+
+ Ok(())
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
