Skip to content

Commit cc6ea44

Browse files
timsaucertimsaucer
committed
Add trait to simplify pattern for accessing FFI extension options
1 parent 238ba00 commit cc6ea44

2 files changed

Lines changed: 65 additions & 66 deletions

File tree

‎datafusion/ffi/src/config/mod.rs‎

Lines changed: 43 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,18 @@ pub mod extension_options;
1919

2020
use abi_stable::StableAbi;
2121
use abi_stable::std_types::{RHashMap, RString};
22-
use datafusion_common::config::{ConfigOptions, ExtensionOptions, TableOptions};
22+
use datafusion_common::config::{
23+
ConfigExtension, ConfigOptions, ExtensionOptions, TableOptions,
24+
};
2325
use datafusion_common::{DataFusionError, Result};
2426

2527
use crate::config::extension_options::FFI_ExtensionOptions;
2628

27-
// TODO(tsaucer) add text about how extension options will require user to convert to concrete type.
29+
/// A stable struct for sharing [`ConfigOptions`] across FFI boundaries.
30+
///
31+
/// Accessing FFI extension options require a slightly different pattern
32+
/// than local extensions. The trait [`ExtensionOptionsFFIProvider`] can
33+
/// be used to simplify accessing FFI extensions.
2834
#[repr(C)]
2935
#[derive(Debug, Clone, StableAbi)]
3036
pub struct FFI_ConfigOptions {
@@ -74,7 +80,41 @@ impl TryFrom<FFI_ConfigOptions> for ConfigOptions {
7480
}
7581
}
7682

77-
// TODO(tsaucer) add text about how extension options will require user to convert to concrete type.
83+
pub trait ExtensionOptionsFFIProvider {
84+
fn ffi_extension<C: ConfigExtension + Clone + Default>(&self) -> Option<C>;
85+
}
86+
87+
impl ExtensionOptionsFFIProvider for ConfigOptions {
88+
fn ffi_extension<C: ConfigExtension + Clone + Default>(&self) -> Option<C> {
89+
self.extensions
90+
.get::<C>()
91+
.map(|v| v.to_owned())
92+
.or_else(|| {
93+
self.extensions
94+
.get::<FFI_ExtensionOptions>()
95+
.and_then(|ffi_ext| ffi_ext.to_extension().ok())
96+
})
97+
}
98+
}
99+
100+
impl ExtensionOptionsFFIProvider for TableOptions {
101+
fn ffi_extension<C: ConfigExtension + Clone + Default>(&self) -> Option<C> {
102+
self.extensions
103+
.get::<C>()
104+
.map(|v| v.to_owned())
105+
.or_else(|| {
106+
self.extensions
107+
.get::<FFI_ExtensionOptions>()
108+
.and_then(|ffi_ext| ffi_ext.to_extension().ok())
109+
})
110+
}
111+
}
112+
113+
/// A stable struct for sharing [`TableOptions`] across FFI boundaries.
114+
///
115+
/// Accessing FFI extension options require a slightly different pattern
116+
/// than local extensions. The trait [`ExtensionOptionsFFIProvider`] can
117+
/// be used to simplify accessing FFI extensions.
78118
#[repr(C)]
79119
#[derive(Debug, Clone, StableAbi)]
80120
pub struct FFI_TableOptions {

‎datafusion/ffi/tests/ffi_config.rs‎

Lines changed: 22 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ mod tests {
2323
use datafusion_common::ScalarValue;
2424
use datafusion_common::config::{ConfigOptions, TableOptions};
2525
use datafusion_execution::config::SessionConfig;
26-
use datafusion_ffi::config::extension_options::FFI_ExtensionOptions;
26+
use datafusion_ffi::config::ExtensionOptionsFFIProvider;
2727
use datafusion_ffi::tests::config::ExternalConfig;
2828
use datafusion_ffi::tests::utils::get_module;
2929

@@ -42,30 +42,16 @@ mod tests {
4242
let mut config = ConfigOptions::new();
4343
config.extensions.insert(extension_options);
4444

45-
fn extract_config(config: &ConfigOptions) -> ExternalConfig {
46-
// For our use case of this test, we do not need to check
47-
// using .get::<ExternalConfig>() but it is left here as an
48-
// example to users of this crate.
49-
config
50-
.extensions
51-
.get::<ExternalConfig>()
52-
.map(|v| v.to_owned())
53-
.or_else(|| {
54-
config
55-
.extensions
56-
.get::<FFI_ExtensionOptions>()
57-
.and_then(|ext| ext.to_extension().ok())
58-
})
59-
.expect("Should be able to get ExternalConfig")
60-
}
61-
6245
// Verify default values are as expected
63-
let returned_config = extract_config(&config);
64-
46+
let returned_config: ExternalConfig = config
47+
.ffi_extension()
48+
.expect("should have external config extension");
6549
assert_eq!(returned_config, ExternalConfig::default());
6650

6751
config.set("external_config.is_enabled", "false")?;
68-
let returned_config = extract_config(&config);
52+
let returned_config: ExternalConfig = config
53+
.ffi_extension()
54+
.expect("should have external config extension");
6955
assert!(!returned_config.is_enabled);
7056

7157
Ok(())
@@ -86,31 +72,18 @@ mod tests {
8672
let mut table_options = TableOptions::new();
8773
table_options.extensions.insert(extension_options);
8874

89-
fn extract_options(options: &TableOptions) -> ExternalConfig {
90-
// For our use case of this test, we do not need to check
91-
// using .get::<ExternalConfig>() but it is left here as an
92-
// example to users of this crate.
93-
options
94-
.extensions
95-
.get::<ExternalConfig>()
96-
.map(|v| v.to_owned())
97-
.or_else(|| {
98-
options
99-
.extensions
100-
.get::<FFI_ExtensionOptions>()
101-
.and_then(|ext| ext.to_extension().ok())
102-
})
103-
.expect("Should be able to get ExternalConfig")
104-
}
105-
10675
// Verify default values are as expected
107-
let returned_options = extract_options(&table_options);
76+
let returned_options: ExternalConfig = table_options
77+
.ffi_extension()
78+
.expect("should have external config extension");
10879

10980
assert_eq!(returned_options, ExternalConfig::default());
11081

11182
table_options.set("external_config.is_enabled", "false")?;
112-
let returned_config = extract_options(&table_options);
113-
assert!(!returned_config.is_enabled);
83+
let returned_options: ExternalConfig = table_options
84+
.ffi_extension()
85+
.expect("should have external config extension");
86+
assert!(!returned_options.is_enabled);
11487

11588
Ok(())
11689
}
@@ -129,35 +102,21 @@ mod tests {
129102

130103
let mut config = SessionConfig::new().with_option_extension(extension_options);
131104

132-
fn extract_config(config: &SessionConfig) -> ExternalConfig {
133-
// For our use case of this test, we do not need to check
134-
// using .get::<ExternalConfig>() but it is left here as an
135-
// example to users of this crate.
136-
137-
config
138-
.options()
139-
.extensions
140-
.get::<ExternalConfig>()
141-
.map(|v| v.to_owned())
142-
.or_else(|| {
143-
config
144-
.options()
145-
.extensions
146-
.get::<FFI_ExtensionOptions>()
147-
.and_then(|ext| ext.to_extension().ok())
148-
})
149-
.expect("Should be able to get ExternalConfig")
150-
}
151-
152105
// Verify default values are as expected
153-
let returned_config = extract_config(&config);
106+
let returned_config: ExternalConfig = config
107+
.options()
108+
.ffi_extension()
109+
.expect("should have external config extension");
154110
assert_eq!(returned_config, ExternalConfig::default());
155111

156112
config = config.set(
157113
"external_config.is_enabled",
158114
&ScalarValue::Boolean(Some(false)),
159115
);
160-
let returned_config = extract_config(&config);
116+
let returned_config: ExternalConfig = config
117+
.options()
118+
.ffi_extension()
119+
.expect("should have external config extension");
161120
assert!(!returned_config.is_enabled);
162121

163122
Ok(())

0 commit comments

Comments
 (0)