From 3c4413e600655220abef390344e6a6af64217a51 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marko=20Milenkovi=C4=87?=
 <milenkovicm@users.noreply.github.com>
Date: Wed, 19 Feb 2025 20:20:55 +0000
Subject: [PATCH] add test to verify switching off configuration

---
 ballista/client/tests/context_checks.rs | 32 +++++++++++++++++++++
 ballista/core/src/extension.rs          | 38 ++++++++++++++++++-------
 2 files changed, 59 insertions(+), 11 deletions(-)

diff --git a/ballista/client/tests/context_checks.rs b/ballista/client/tests/context_checks.rs
index a46f93803..d1ecf2f65 100644
--- a/ballista/client/tests/context_checks.rs
+++ b/ballista/client/tests/context_checks.rs
@@ -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(())
+    }
 }
diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs
index a0a1a67c5..4ff1af98b 100644
--- a/ballista/core/src/extension.rs
+++ b/ballista/core/src/extension.rs
@@ -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 {
@@ -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)
@@ -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,
@@ -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