|
| 1 | +use datafusion::common::{config_err, DataFusionError}; |
| 2 | +use datafusion::config::{ConfigEntry, ConfigExtension, ConfigField, ExtensionOptions, Visit}; |
| 3 | +use pyo3::{pyclass, pymethods, Bound, PyResult, Python}; |
| 4 | +use std::any::Any; |
| 5 | +use datafusion_ffi::config::extension_options::FFI_ExtensionOptions; |
| 6 | +use pyo3::exceptions::PyRuntimeError; |
| 7 | +use pyo3::types::PyCapsule; |
| 8 | + |
| 9 | +/// My own config options. |
| 10 | +#[pyclass(name = "MyConfig", module = "datafusion_ffi_example", subclass)] |
| 11 | +#[derive(Clone, Debug)] |
| 12 | +pub struct MyConfig { |
| 13 | + /// Should "foo" be replaced by "bar"? |
| 14 | + pub foo_to_bar: bool, |
| 15 | + |
| 16 | + /// How many "baz" should be created? |
| 17 | + pub baz_count: usize, |
| 18 | +} |
| 19 | + |
| 20 | +#[pymethods] |
| 21 | +impl MyConfig { |
| 22 | + #[new] |
| 23 | + fn new() -> Self { |
| 24 | + Self::default() |
| 25 | + } |
| 26 | + |
| 27 | + |
| 28 | + fn __datafusion_extension_options__<'py>( |
| 29 | + &self, |
| 30 | + py: Python<'py>, |
| 31 | + ) -> PyResult<Bound<'py, PyCapsule>> { |
| 32 | + let name = cr"datafusion_extension_options".into(); |
| 33 | + |
| 34 | + let mut config = FFI_ExtensionOptions::default(); |
| 35 | + config.add_config(self).map_err(|e| PyRuntimeError::new_err(e.to_string()))?; |
| 36 | + |
| 37 | + PyCapsule::new(py, config, Some(name)) |
| 38 | + } |
| 39 | +} |
| 40 | + |
| 41 | +impl Default for MyConfig { |
| 42 | + fn default() -> Self { |
| 43 | + Self { |
| 44 | + foo_to_bar: true, |
| 45 | + baz_count: 1337, |
| 46 | + } |
| 47 | + } |
| 48 | +} |
| 49 | + |
| 50 | +impl ConfigExtension for MyConfig { |
| 51 | + const PREFIX: &'static str = "my_config"; |
| 52 | +} |
| 53 | + |
| 54 | +impl ExtensionOptions for MyConfig { |
| 55 | + fn as_any(&self) -> &dyn Any { |
| 56 | + self |
| 57 | + } |
| 58 | + |
| 59 | + fn as_any_mut(&mut self) -> &mut dyn Any { |
| 60 | + self |
| 61 | + } |
| 62 | + |
| 63 | + fn cloned(&self) -> Box<dyn ExtensionOptions> { |
| 64 | + Box::new(self.clone()) |
| 65 | + } |
| 66 | + |
| 67 | + fn set(&mut self, key: &str, value: &str) -> datafusion::common::Result<()> { |
| 68 | + datafusion::config::ConfigField::set(self, key, value) |
| 69 | + } |
| 70 | + |
| 71 | + fn entries(&self) -> Vec<ConfigEntry> { |
| 72 | + vec![ |
| 73 | + ConfigEntry { |
| 74 | + key: "foo_to_bar".to_owned(), |
| 75 | + value: Some(format!("{}", self.foo_to_bar)), |
| 76 | + description: "foo to bar", |
| 77 | + }, |
| 78 | + ConfigEntry { |
| 79 | + key: "baz_count".to_owned(), |
| 80 | + value: Some(format!("{}", self.baz_count)), |
| 81 | + description: "baz count", |
| 82 | + }, |
| 83 | + ] |
| 84 | + } |
| 85 | +} |
| 86 | + |
| 87 | +impl ConfigField for MyConfig { |
| 88 | + fn visit<V: Visit>(&self, v: &mut V, _key: &str, _description: &'static str) { |
| 89 | + let key = "foo_to_bar"; |
| 90 | + let desc = "foo to bar"; |
| 91 | + self.foo_to_bar.visit(v, key, desc); |
| 92 | + |
| 93 | + let key = "baz_count"; |
| 94 | + let desc = "baz count"; |
| 95 | + self.baz_count.visit(v, key, desc); |
| 96 | + } |
| 97 | + |
| 98 | + fn set(&mut self, key: &str, value: &str) -> Result<(), DataFusionError> { |
| 99 | + let (key, rem) = key.split_once('.').unwrap_or((key, "")); |
| 100 | + match key { |
| 101 | + "foo_to_bar" => self.foo_to_bar.set(rem, value.as_ref()), |
| 102 | + "baz_count" => self.baz_count.set(rem, value.as_ref()), |
| 103 | + |
| 104 | + _ => config_err!("Config value \"{}\" not found on MyConfig", key), |
| 105 | + } |
| 106 | + |
| 107 | + } |
| 108 | +} |
0 commit comments