Skip to content

Commit

Permalink
add test to verify switching off configuration
Browse files Browse the repository at this point in the history
  • Loading branch information
milenkovicm committed Feb 19, 2025
1 parent be8979c commit 3c4413e
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 11 deletions.
32 changes: 32 additions & 0 deletions ballista/client/tests/context_checks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -365,4 +365,36 @@ mod supported {

Ok(())
}

// test checks if this view types have been disabled in the configuration
//
// `datafusion.execution.parquet.schema_force_view_types` have been disabled
// temporary as they could break shuffle reader/writer.
#[rstest]
#[case::standalone(standalone_context())]
#[case::remote(remote_context())]
#[tokio::test]
async fn should_disable_view_types(
#[future(awt)]
#[case]
ctx: SessionContext,
) -> datafusion::error::Result<()> {
let result = ctx
.sql("select name, value from information_schema.df_settings where name like 'datafusion.execution.parquet.schema_force_view_types' order by name limit 1")
.await?
.collect()
.await?;
//
let expected = [
"+------------------------------------------------------+-------+",
"| name | value |",
"+------------------------------------------------------+-------+",
"| datafusion.execution.parquet.schema_force_view_types | false |",
"+------------------------------------------------------+-------+",
];

assert_batches_eq!(expected, &result);

Ok(())
}
}
38 changes: 27 additions & 11 deletions ballista/core/src/extension.rs
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,9 @@ pub trait SessionConfigHelperExt {
fn update_from_key_value_pair(self, key_value_pairs: &[KeyValuePair]) -> Self;
/// updates mut [SessionConfig] from proto
fn update_from_key_value_pair_mut(&mut self, key_value_pairs: &[KeyValuePair]);
/// changes some of default datafusion configuration
/// in order to make it suitable for ballista
fn ballista_restricted_configuration(self) -> Self;
}

impl SessionStateExt for SessionState {
Expand Down Expand Up @@ -168,16 +171,7 @@ impl SessionStateExt for SessionState {
.config()
.clone()
.with_option_extension(new_config.clone())
// Ballista disables this option
.with_round_robin_repartition(false)
// There is issue with Utv8View(s) where Arrow IPC will generate
// frames which would be too big to send using Arrow Flight.
// We disable this option temporary
// TODO: enable this option once we get to root of the problem
.set_bool(
"datafusion.execution.parquet.schema_force_view_types",
false,
);
.ballista_restricted_configuration();

let builder = SessionStateBuilder::new_from_existing(self)
.with_config(session_config)
Expand Down Expand Up @@ -205,7 +199,7 @@ impl SessionConfigExt for SessionConfig {
.with_option_extension(BallistaConfig::default())
.with_information_schema(true)
.with_target_partitions(16)
.with_round_robin_repartition(false)
.ballista_restricted_configuration()
}
fn with_ballista_logical_extension_codec(
self,
Expand Down Expand Up @@ -367,6 +361,28 @@ impl SessionConfigHelperExt for SessionConfig {
}
}
}

fn ballista_restricted_configuration(self) -> Self {
self
// round robbin repartition does not work well with ballista.
// this setting it will also be enforced by the scheduler
// thus user will not be able to override it.
.with_round_robin_repartition(false)
// There is issue with Utv8View(s) where Arrow IPC will generate
// frames which would be too big to send using Arrow Flight.
//
// This configuration option will be disabled temporary.
//
// This configuration is not enforced by the scheduler, thus
// user could override this setting using `SET` operation.
//
// TODO: enable this option once we get to root of the problem
// between `IpcWriter` and `ViewTypes`
.set_bool(
"datafusion.execution.parquet.schema_force_view_types",
false,
)
}
}

/// Wrapper for [SessionConfig] extension
Expand Down

0 comments on commit 3c4413e

Please sign in to comment.