Re: [PR] Set DataFusion runtime configurations through SQL interface [datafusion]

[via GitHub]

berkaysynnada commented on code in PR #15594:
URL: https://github.com/apache/datafusion/pull/15594#discussion_r2033296553


##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -1036,13 +1040,82 @@ impl SessionContext {
             variable, value, ..
         } = stmt;
 
-        let mut state = self.state.write();
-        state.config_mut().options_mut().set(&variable, &value)?;
-        drop(state);
+        // Check if this is a runtime configuration
+        if variable.starts_with("datafusion.runtime.") {
+            self.set_runtime_variable(&variable, &value)?;
+        } else {
+            let mut state = self.state.write();
+            state.config_mut().options_mut().set(&variable, &value)?;
+            drop(state);
+        }
 
         self.return_empty_dataframe()
     }
 
+    fn set_runtime_variable(&self, variable: &str, value: &str) -> Result<()> {
+        let key = variable.strip_prefix("datafusion.runtime.").unwrap();
+
+        match key {
+            "memory_limit" => {
+                let memory_limit = self.parse_memory_limit(value)?;
+
+                let current_runtime = {
+                    let state = self.state.read();
+                    Arc::<RuntimeEnv>::clone(state.runtime_env())
+                };
+
+                let mut builder = 
RuntimeEnvBuilder::from_runtime_env(&current_runtime);
+                builder = builder.with_memory_limit(memory_limit, 1.0);
+
+                let new_runtime = builder.build()?;
+                self.update_runtime_env(Arc::new(new_runtime))?;
+            }
+            _ => {
+                return Err(DataFusionError::Plan(format!(
+                    "Unknown runtime configuration: {}",
+                    variable
+                )));
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Parse memory limit from string to number of bytes
+    /// e.g. '1.5G', '100M'
+    fn parse_memory_limit(&self, limit: &str) -> Result<usize> {

Review Comment:
   Maybe we can open this function to external uses?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org