Implement config options

Author: timsaucer

crates/core/src/context.rs

@@ -49,6 +49,7 @@ use datafusion::prelude::{
 };
 use datafusion_ffi::catalog_provider::FFI_CatalogProvider;
 use datafusion_ffi::catalog_provider_list::FFI_CatalogProviderList;
+use datafusion_ffi::config::extension_options::FFI_ExtensionOptions;
 use datafusion_ffi::execution::FFI_TaskContextProvider;
 use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
 use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory;
@@ -185,6 +186,27 @@ impl PySessionConfig {
     fn set(&self, key: &str, value: &str) -> Self {
         Self::from(self.config.clone().set_str(key, value))
     }
+
+    pub fn with_extension(&self, extension: Bound<PyAny>) -> PyResult<Self> {
+        let capsule = extension.call_method0("__datafusion_extension_options__")?;
+        let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+
+        validate_pycapsule(capsule, "datafusion_extension_options")?;
+
+        let mut extension = unsafe { capsule.reference::<FFI_ExtensionOptions>() }.clone();
+
+        let mut config = self.config.clone();
+        let options = config.options_mut();
+        if let Some(prior_extension) = options.extensions.get::<FFI_ExtensionOptions>() {
+            extension
+                .merge(prior_extension)
+                .map_err(py_datafusion_err)?;
+        }
+
+        options.extensions.insert(extension);
+
+        Ok(Self::from(config))
+    }
 }
 
 /// Runtime options for a SessionContext

crates/core/src/dataset_exec.rs

@@ -111,7 +111,7 @@ impl DatasetExec {
 
         let scanner = dataset.call_method("scanner", (), Some(&kwargs))?;
 
-        let schema = Arc::new(
+        let schema: SchemaRef = Arc::new(
             scanner
                 .getattr("projected_schema")?
                 .extract::<PyArrowType<_>>()?

examples/datafusion-ffi-example/python/tests/_test_config.py

@@ -0,0 +1,8 @@
+import pyarrow as pa
+from datafusion import SessionConfig
+from datafusion_ffi_example import MyConfig
+
+def test_catalog_provider():
+    config = MyConfig()
+    config = SessionConfig().with_extension(config)
+    config.set("my_config.baz_count", "42")

examples/datafusion-ffi-example/src/config.rs

@@ -0,0 +1,108 @@
+use datafusion::common::{config_err, DataFusionError};
+use datafusion::config::{ConfigEntry, ConfigExtension, ConfigField, ExtensionOptions, Visit};
+use pyo3::{pyclass, pymethods, Bound, PyResult, Python};
+use std::any::Any;
+use datafusion_ffi::config::extension_options::FFI_ExtensionOptions;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::types::PyCapsule;
+
+/// My own config options.
+#[pyclass(name = "MyConfig", module = "datafusion_ffi_example", subclass)]
+#[derive(Clone, Debug)]
+pub struct MyConfig {
+    /// Should "foo" be replaced by "bar"?
+    pub foo_to_bar: bool,
+
+    /// How many "baz" should be created?
+    pub baz_count: usize,
+}
+
+#[pymethods]
+impl MyConfig {
+    #[new]
+    fn new() -> Self {
+        Self::default()
+    }
+
+
+    fn __datafusion_extension_options__<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let name = cr"datafusion_extension_options".into();
+
+        let mut config = FFI_ExtensionOptions::default();
+        config.add_config(self).map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+
+        PyCapsule::new(py, config, Some(name))
+    }
+}
+
+impl Default for MyConfig {
+    fn default() -> Self {
+        Self {
+            foo_to_bar: true,
+            baz_count: 1337,
+        }
+    }
+}
+
+impl ConfigExtension for MyConfig {
+    const PREFIX: &'static str = "my_config";
+}
+
+impl ExtensionOptions for MyConfig {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        self
+    }
+
+    fn cloned(&self) -> Box<dyn ExtensionOptions> {
+        Box::new(self.clone())
+    }
+
+    fn set(&mut self, key: &str, value: &str) -> datafusion::common::Result<()> {
+        datafusion::config::ConfigField::set(self, key, value)
+    }
+
+    fn entries(&self) -> Vec<ConfigEntry> {
+        vec![
+            ConfigEntry {
+                key: "foo_to_bar".to_owned(),
+                value: Some(format!("{}", self.foo_to_bar)),
+                description: "foo to bar",
+            },
+            ConfigEntry {
+                key: "baz_count".to_owned(),
+                value: Some(format!("{}", self.baz_count)),
+                description: "baz count",
+            },
+        ]
+    }
+}
+
+impl ConfigField for MyConfig {
+    fn visit<V: Visit>(&self, v: &mut V, _key: &str, _description: &'static str) {
+        let key = "foo_to_bar";
+        let desc = "foo to bar";
+        self.foo_to_bar.visit(v, key, desc);
+
+        let key = "baz_count";
+        let desc = "baz count";
+        self.baz_count.visit(v, key, desc);
+    }
+
+    fn set(&mut self, key: &str, value: &str) -> Result<(), DataFusionError> {
+        let (key, rem) = key.split_once('.').unwrap_or((key, ""));
+        match key {
+            "foo_to_bar" => self.foo_to_bar.set(rem, value.as_ref()),
+            "baz_count" => self.baz_count.set(rem, value.as_ref()),
+
+            _ => config_err!("Config value \"{}\" not found on MyConfig", key),
+        }
+
+    }
+}

python/datafusion/context.py

@@ -296,6 +296,10 @@ def set(self, key: str, value: str) -> SessionConfig:
         self.config_internal = self.config_internal.set(key, value)
         return self
 
+    def with_extension(self, extension: Any) -> SessionConfig:
+        self.config_internal = self.config_internal.with_extension(extension)
+        return self
+
 
 class RuntimeEnvBuilder:
     """Runtime configuration options."""