Skip to content

Commit 44e80b1

Browse files
timsaucerdavisp
authored andcommitted
Update examples and tests
1 parent 51346f7 commit 44e80b1

4 files changed

Lines changed: 35 additions & 11 deletions

File tree

examples/datafusion-ffi-example/Cargo.toml

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,3 +51,11 @@ pyo3-build-config = { workspace = true }
5151
[lib]
5252
name = "datafusion_ffi_example"
5353
crate-type = ["cdylib", "rlib"]
54+
55+
[patch.crates-io]
56+
datafusion-catalog = { git = "https://github.com/timsaucer/datafusion", rev = "33ea91c5fa37aa0cfa787d58c6bf84ec4628db88" }
57+
datafusion-common = { git = "https://github.com/timsaucer/datafusion", rev = "33ea91c5fa37aa0cfa787d58c6bf84ec4628db88" }
58+
datafusion-functions-aggregate = { git = "https://github.com/timsaucer/datafusion", rev = "33ea91c5fa37aa0cfa787d58c6bf84ec4628db88" }
59+
datafusion-functions-window = { git = "https://github.com/timsaucer/datafusion", rev = "33ea91c5fa37aa0cfa787d58c6bf84ec4628db88" }
60+
datafusion-expr = { git = "https://github.com/timsaucer/datafusion", rev = "33ea91c5fa37aa0cfa787d58c6bf84ec4628db88" }
61+
datafusion-ffi = { git = "https://github.com/timsaucer/datafusion", rev = "33ea91c5fa37aa0cfa787d58c6bf84ec4628db88" }
Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,18 @@
1-
import pyarrow as pa
2-
from datafusion import SessionConfig
1+
from datafusion import SessionConfig, SessionContext
32
from datafusion_ffi_example import MyConfig
43

4+
55
def test_catalog_provider():
66
config = MyConfig()
7-
config = SessionConfig().with_extension(config)
7+
config = SessionConfig(
8+
{"datafusion.catalog.information_schema": "true"}
9+
).with_extension(config)
810
config.set("my_config.baz_count", "42")
11+
ctx = SessionContext(config)
12+
13+
result = ctx.sql("SHOW my_config.baz_count;").collect()
14+
assert result[0][1][0].as_py() == "42"
15+
16+
ctx.sql("SET my_config.baz_count=1;")
17+
result = ctx.sql("SHOW my_config.baz_count;").collect()
18+
assert result[0][1][0].as_py() == "1"

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

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
1-
use datafusion::common::{config_err, DataFusionError};
2-
use datafusion::config::{ConfigEntry, ConfigExtension, ConfigField, ExtensionOptions, Visit};
3-
use pyo3::{pyclass, pymethods, Bound, PyResult, Python};
41
use std::any::Any;
2+
3+
use datafusion_common::config::{
4+
ConfigEntry, ConfigExtension, ConfigField, ExtensionOptions, Visit,
5+
};
6+
use datafusion_common::{DataFusionError, config_err};
57
use datafusion_ffi::config::extension_options::FFI_ExtensionOptions;
68
use pyo3::exceptions::PyRuntimeError;
79
use pyo3::types::PyCapsule;
10+
use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
811

912
/// My own config options.
1013
#[pyclass(name = "MyConfig", module = "datafusion_ffi_example", subclass)]
@@ -24,15 +27,16 @@ impl MyConfig {
2427
Self::default()
2528
}
2629

27-
2830
fn __datafusion_extension_options__<'py>(
2931
&self,
3032
py: Python<'py>,
3133
) -> PyResult<Bound<'py, PyCapsule>> {
3234
let name = cr"datafusion_extension_options".into();
3335

3436
let mut config = FFI_ExtensionOptions::default();
35-
config.add_config(self).map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
37+
config
38+
.add_config(self)
39+
.map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
3640

3741
PyCapsule::new(py, config, Some(name))
3842
}
@@ -64,8 +68,8 @@ impl ExtensionOptions for MyConfig {
6468
Box::new(self.clone())
6569
}
6670

67-
fn set(&mut self, key: &str, value: &str) -> datafusion::common::Result<()> {
68-
datafusion::config::ConfigField::set(self, key, value)
71+
fn set(&mut self, key: &str, value: &str) -> datafusion_common::Result<()> {
72+
datafusion_common::config::ConfigField::set(self, key, value)
6973
}
7074

7175
fn entries(&self) -> Vec<ConfigEntry> {
@@ -103,6 +107,5 @@ impl ConfigField for MyConfig {
103107

104108
_ => config_err!("Config value \"{}\" not found on MyConfig", key),
105109
}
106-
107110
}
108111
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ use pyo3::prelude::*;
1919

2020
use crate::aggregate_udf::MySumUDF;
2121
use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogProviderList};
22+
use crate::config::MyConfig;
2223
use crate::scalar_udf::IsNullUDF;
2324
use crate::table_function::MyTableFunction;
2425
use crate::table_provider::MyTableProvider;
@@ -27,6 +28,7 @@ use crate::window_udf::MyRankUDF;
2728

2829
pub(crate) mod aggregate_udf;
2930
pub(crate) mod catalog_provider;
31+
pub(crate) mod config;
3032
pub(crate) mod scalar_udf;
3133
pub(crate) mod table_function;
3234
pub(crate) mod table_provider;
@@ -46,5 +48,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
4648
m.add_class::<IsNullUDF>()?;
4749
m.add_class::<MySumUDF>()?;
4850
m.add_class::<MyRankUDF>()?;
51+
m.add_class::<MyConfig>()?;
4952
Ok(())
5053
}

0 commit comments

Comments
 (0)