fix AQE planning with S3 (#1962)

Author: milenkovicm

ballista/scheduler/src/state/aqe/planner.rs

@@ -31,6 +31,7 @@ use datafusion::common;
 use datafusion::common::{HashMap, exec_err};
 use datafusion::error::DataFusionError;
 use datafusion::execution::context::SessionContext;
+use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::execution::{SessionState, SessionStateBuilder};
 use datafusion::logical_expr::LogicalPlan;
 use datafusion::physical_optimizer::PhysicalOptimizerRule;
@@ -87,6 +88,7 @@ impl AdaptivePlanner {
     /// # Arguments:
     ///
     /// * `session_config` - The session configuration for the job.
+    /// * `runtime_env` - runtime environment
     /// * `plan` - The physical execution plan for the job.
     /// * `job_name` - The name of the job.
     /// * `physical_optimizer_rules` - A list of physical optimizer rules to apply.
@@ -95,12 +97,16 @@ impl AdaptivePlanner {
     /// A new instance of `AdaptivePlanner` or an error if the initialization fails.
     pub fn try_new_with_optimizers(
         session_config: &SessionConfig,
+        runtime_env: Arc<RuntimeEnv>,
         plan: Arc<dyn ExecutionPlan>,
         job_name: String,
         physical_optimizer_rules: Vec<PhysicalOptimizerRuleRef>,
     ) -> common::Result<Self> {
-        let session_state =
-            Self::create_session_state(session_config, physical_optimizer_rules);
+        let session_state = Self::create_session_state(
+            session_config,
+            runtime_env,
+            physical_optimizer_rules,
+        );
         let planner = DefaultPhysicalPlanner::default();
         let plan = planner.optimize_physical_plan(plan, &session_state, |_, _| {})?;
 
@@ -133,6 +139,7 @@ impl AdaptivePlanner {
         let plan_id_generator = Arc::new(AtomicUsize::new(0));
         Self::try_new_with_optimizers(
             session_config,
+            RuntimeEnv::default().into(),
             plan,
             job_name,
             Self::default_optimizers(plan_id_generator),
@@ -159,12 +166,16 @@ impl AdaptivePlanner {
         // running standard set of optimizers, which will
         // after each stage.
         let plan_id_generator = Arc::new(AtomicUsize::new(0));
-        let state = Self::create_session_state(
+        let runtime_env = ctx.runtime_env();
+        let plan_preparation_stage = Self::create_session_state(
             ctx.state().config(),
+            ctx.runtime_env(),
             Self::plan_preparation_optimizers(plan_id_generator.clone()),
         );
 
-        let plan = state.create_physical_plan(logical_plan).await?;
+        let plan = plan_preparation_stage
+            .create_physical_plan(logical_plan)
+            .await?;
 
         // Note: the signature requires a JobId, but we are passing a JobName. The below is a
         // dirty fix but this seems like a bug or a design flaw.
@@ -175,6 +186,7 @@ impl AdaptivePlanner {
 
         Self::try_new_with_optimizers(
             ctx.state().config(),
+            runtime_env,
             plan,
             job_name,
             Self::default_optimizers(plan_id_generator),
@@ -546,11 +558,13 @@ impl AdaptivePlanner {
     /// A new `SessionState` instance.
     fn create_session_state(
         session_config: &SessionConfig,
+        runtime_env: Arc<RuntimeEnv>,
         physical_optimizers: Vec<PhysicalOptimizerRuleRef>,
     ) -> SessionState {
         SessionStateBuilder::new_with_default_features()
             .with_physical_optimizer_rules(physical_optimizers)
             .with_config(session_config.clone())
+            .with_runtime_env(runtime_env)
             .build()
     }
     /// Recursively finds runnable exchanges in the execution plan.

examples/tests/object_store.rs

@@ -371,6 +371,149 @@ mod custom_s3_config {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn should_configure_aqe_s3_execute_sql_write_remote()
+    -> datafusion::error::Result<()> {
+        let test_data = examples_test_data();
+
+        //
+        // Minio cluster setup
+        //
+        let container = crate::common::create_minio_container();
+        let node = container.start().await.unwrap();
+
+        node.exec(crate::common::create_bucket_command())
+            .await
+            .unwrap();
+
+        let endpoint_host = node.get_host().await.unwrap();
+        let endpoint_port = node.get_host_port_ipv4(9000).await.unwrap();
+
+        log::info!(
+            "MINIO testcontainers host: {}, port: {}",
+            endpoint_host,
+            endpoint_port
+        );
+
+        //
+        // Session Context and Ballista cluster setup
+        //
+
+        // Setting up configuration producer
+        //
+        // configuration producer registers user defined config extension
+        // S3Option with relevant S3 configuration
+        let config_producer =
+            Arc::new(ballista_core::object_store::session_config_with_s3_support);
+        // Setting up runtime producer
+        //
+        // Runtime producer creates object store registry
+        // which can create object store connecter based on
+        // S3Option configuration.
+        let runtime_producer: RuntimeProducer =
+            Arc::new(ballista_core::object_store::runtime_env_with_s3_support);
+
+        // Session builder creates SessionState
+        //
+        // which is configured using runtime and configuration producer,
+        // producing same runtime environment, and providing same
+        // object store registry.
+
+        let session_builder =
+            Arc::new(ballista_core::object_store::session_state_with_s3_support);
+
+        let state = session_builder(config_producer())?;
+
+        // setting up ballista cluster with new runtime, configuration, and session state producers
+        let (host, port) = crate::common::setup_test_cluster_with_builders(
+            config_producer,
+            runtime_producer,
+            session_builder,
+        )
+        .await;
+        let url = format!("df://{host}:{port}");
+
+        // establishing cluster connection,
+        let ctx: SessionContext = SessionContext::remote_with_state(&url, state).await?;
+
+        ctx.sql("SET ballista.planner.adaptive.enabled = true")
+            .await?
+            .show()
+            .await?;
+        // setting up relevant S3 options
+        ctx.sql("SET s3.allow_http = true").await?.show().await?;
+        ctx.sql(&format!("SET s3.access_key_id = '{}'", ACCESS_KEY_ID))
+            .await?
+            .show()
+            .await?;
+        ctx.sql(&format!("SET s3.secret_access_key = '{}'", SECRET_KEY))
+            .await?
+            .show()
+            .await?;
+        ctx.sql(&format!(
+            "SET s3.endpoint = 'http://{}:{}'",
+            endpoint_host, endpoint_port
+        ))
+        .await?
+        .show()
+        .await?;
+        ctx.sql("SET s3.allow_http = true").await?.show().await?;
+
+        // verifying that we have set S3Options
+        ctx.sql("select name, value from information_schema.df_settings where name like 's3.%'").await?.show().await?;
+
+        ctx.register_parquet(
+            "test",
+            &format!("{test_data}/alltypes_plain.parquet"),
+            Default::default(),
+        )
+        .await?;
+
+        let write_dir_path =
+            &format!("s3://{}/write_test.parquet", crate::common::BUCKET);
+
+        ctx.sql("select * from test")
+            .await?
+            .write_parquet(write_dir_path, Default::default(), Default::default())
+            .await?;
+
+        ctx.register_parquet("written_table", write_dir_path, Default::default())
+            .await?;
+
+        let result = ctx
+            .sql("select id, string_col, timestamp_col from written_table where id > 4")
+            .await?
+            .collect()
+            .await?;
+        let expected = [
+            "+----+------------+---------------------+",
+            "| id | string_col | timestamp_col       |",
+            "+----+------------+---------------------+",
+            "| 5  | 31         | 2009-03-01T00:01:00 |",
+            "| 6  | 30         | 2009-04-01T00:00:00 |",
+            "| 7  | 31         | 2009-04-01T00:01:00 |",
+            "+----+------------+---------------------+",
+        ];
+
+        assert_batches_eq!(expected, &result);
+
+        let result = ctx
+            .sql("select max(id) as max, min(id) as min from written_table")
+            .await?
+            .collect()
+            .await?;
+        let expected = [
+            "+-----+-----+",
+            "| max | min |",
+            "+-----+-----+",
+            "| 7   | 0   |",
+            "+-----+-----+",
+        ];
+
+        assert_batches_eq!(expected, &result);
+        Ok(())
+    }
+
     // this test shows how to register external ObjectStoreRegistry and configure it
     // using infrastructure provided by ballista standalone.
     //