[Config] - Support for key-value pair configuration settings

datafusion/__init__.py

@@ -23,22 +23,21 @@
 except ImportError:
     import importlib_metadata
 
-
 import pyarrow as pa
 
 from ._internal import (
     AggregateUDF,
+    Config,
     DataFrame,
     SessionContext,
     Expression,
     ScalarUDF,
 )
 
-
 __version__ = importlib_metadata.version(__name__)
 
-
 __all__ = [
+    "Config",
     "DataFrame",
     "SessionContext",
     "Expression",

datafusion/tests/test_config.py

@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from datafusion import Config
+import pytest
+
+
+@pytest.fixture
+def config():
+    return Config()
+
+
+def test_get_then_set(config):
+    config_key = "datafusion.optimizer.filter_null_join_keys"
+
+    assert config.get(config_key).as_py() is False
+
+    config.set(config_key, True)
+    assert config.get(config_key).as_py() is True
+
+
+def test_get_all(config):
+    config.get_all()
+
+
+def test_get_invalid_config(config):
+    assert config.get("not.valid.key") is None

src/config.rs

@@ -0,0 +1,82 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use pyo3::prelude::*;
+use pyo3::types::*;
+
+use datafusion::config::ConfigOptions;
+use datafusion_common::ScalarValue;
+
+#[pyclass(name = "Config", module = "datafusion", subclass)]
+#[derive(Clone)]
+pub(crate) struct PyConfig {
+    config: ConfigOptions,
+}
+
+#[pymethods]
+impl PyConfig {
+    #[new]
+    fn py_new() -> Self {
+        Self {
+            config: ConfigOptions::new(),
+        }
+    }
+
+    /// Get configurations from environment variables
+    #[staticmethod]
+    pub fn from_env() -> Self {
+        Self {
+            config: ConfigOptions::from_env(),
+        }
+    }
+
+    /// Get a configuration option
+    pub fn get(&mut self, key: &str, py: Python) -> PyResult<PyObject> {
+        Ok(self.config.get(key).into_py(py))
+    }
+
+    /// Set a configuration option
+    pub fn set(&mut self, key: &str, value: PyObject, py: Python) {
+        self.config.set(key, py_obj_to_scalar_value(py, value))
+    }
+
+    /// Get all configuration options
+    pub fn get_all(&mut self, py: Python) -> PyResult<PyObject> {
+        let dict = PyDict::new(py);
+        for (key, value) in self.config.options() {
+            dict.set_item(key, value.clone().into_py(py))?;
+        }
+        Ok(dict.into())
+    }
+}
+
+/// Convert a python object to a ScalarValue
+fn py_obj_to_scalar_value(py: Python, obj: PyObject) -> ScalarValue {
+    if let Ok(value) = obj.extract::<bool>(py) {
+        ScalarValue::Boolean(Some(value))
+    } else if let Ok(value) = obj.extract::<i64>(py) {
+        ScalarValue::Int64(Some(value))
+    } else if let Ok(value) = obj.extract::<u64>(py) {
+        ScalarValue::UInt64(Some(value))
+    } else if let Ok(value) = obj.extract::<f64>(py) {
+        ScalarValue::Float64(Some(value))
+    } else if let Ok(value) = obj.extract::<String>(py) {
+        ScalarValue::Utf8(Some(value))
+    } else {
+        panic!("Unsupported value type")
+    }
+}

src/expression.rs

@@ -19,7 +19,7 @@ use pyo3::{basic::CompareOp, prelude::*};
 use std::convert::{From, Into};
 
 use datafusion::arrow::datatypes::DataType;
-use datafusion::logical_plan::{col, lit, Expr};
+use datafusion_expr::{col, lit, Expr};
 
 use datafusion::scalar::ScalarValue;
 

src/lib.rs

@@ -22,6 +22,8 @@ use pyo3::prelude::*;
 #[allow(clippy::borrow_deref_ref)]
 pub mod catalog;
 #[allow(clippy::borrow_deref_ref)]
+mod config;
+#[allow(clippy::borrow_deref_ref)]
 mod context;
 #[allow(clippy::borrow_deref_ref)]
 mod dataframe;
@@ -58,6 +60,7 @@ fn _internal(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<expression::PyExpr>()?;
     m.add_class::<udf::PyScalarUDF>()?;
     m.add_class::<udaf::PyAggregateUDF>()?;
+    m.add_class::<config::PyConfig>()?;
 
     // Register the functions as a submodule
     let funcs = PyModule::new(py, "functions")?;