Skip to content

Commit f53e007

Browse files
committed
Debugging a failing test
1 parent 4811621 commit f53e007

10 files changed

Lines changed: 316 additions & 158 deletions

File tree

Cargo.lock

Lines changed: 124 additions & 117 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -70,12 +70,12 @@ codegen-units = 1
7070
# We cannot publish to crates.io with any patches in the below section. Developers
7171
# must remove any entries in this section before creating a release candidate.
7272
[patch.crates-io]
73-
datafusion = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f" }
74-
datafusion-substrait = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f" }
75-
datafusion-proto = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f" }
76-
datafusion-ffi = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f" }
77-
datafusion-catalog = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f" }
78-
datafusion-common = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f" }
79-
datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f" }
80-
datafusion-functions-window = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f" }
81-
datafusion-expr = { git = "https://github.com/apache/datafusion.git", rev = "35749607f585b3bf25b66b7d2289c56c18d03e4f" }
73+
datafusion = { git = "https://github.com/davisp/datafusion.git", branch = "pd/test/ffi-config" }
74+
datafusion-substrait = { git = "https://github.com/davisp/datafusion.git", branch = "pd/test/ffi-config" }
75+
datafusion-proto = { git = "https://github.com/davisp/datafusion.git", branch = "pd/test/ffi-config" }
76+
datafusion-ffi = { git = "https://github.com/davisp/datafusion.git", branch = "pd/test/ffi-config" }
77+
datafusion-catalog = { git = "https://github.com/davisp/datafusion.git", branch = "pd/test/ffi-config" }
78+
datafusion-common = { git = "https://github.com/davisp/datafusion.git", branch = "pd/test/ffi-config" }
79+
datafusion-functions-aggregate = { git = "https://github.com/davisp/datafusion.git", branch = "pd/test/ffi-config" }
80+
datafusion-functions-window = { git = "https://github.com/davisp/datafusion.git", branch = "pd/test/ffi-config" }
81+
datafusion-expr = { git = "https://github.com/davisp/datafusion.git", branch = "pd/test/ffi-config" }

crates/core/Cargo.toml

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -25,39 +25,40 @@ description.workspace = true
2525
homepage.workspace = true
2626
repository.workspace = true
2727
include = [
28-
"src",
29-
"../LICENSE.txt",
30-
"build.rs",
31-
"../pyproject.toml",
32-
"Cargo.toml",
33-
"../Cargo.lock",
28+
"src",
29+
"../LICENSE.txt",
30+
"build.rs",
31+
"../pyproject.toml",
32+
"Cargo.toml",
33+
"../Cargo.lock",
3434
]
3535

3636
[dependencies]
3737
tokio = { workspace = true, features = [
38-
"macros",
39-
"rt",
40-
"rt-multi-thread",
41-
"sync",
38+
"macros",
39+
"rt",
40+
"rt-multi-thread",
41+
"sync",
4242
] }
4343
pyo3 = { workspace = true, features = [
44-
"extension-module",
45-
"abi3",
46-
"abi3-py310",
44+
"extension-module",
45+
"abi3",
46+
"abi3-py310",
4747
] }
4848
pyo3-async-runtimes = { workspace = true, features = ["tokio-runtime"] }
4949
pyo3-log = { workspace = true }
5050
arrow = { workspace = true, features = ["pyarrow"] }
5151
arrow-select = { workspace = true }
5252
datafusion = { workspace = true, features = ["avro", "unicode_expressions"] }
53+
datafusion-common = { workspace = true }
5354
datafusion-substrait = { workspace = true, optional = true }
5455
datafusion-proto = { workspace = true }
5556
datafusion-ffi = { workspace = true }
5657
prost = { workspace = true } # keep in line with `datafusion-substrait`
5758
serde_json = { workspace = true }
5859
uuid = { workspace = true, features = ["v4"] }
5960
mimalloc = { workspace = true, optional = true, features = [
60-
"local_dynamic_tls",
61+
"local_dynamic_tls",
6162
] }
6263
async-trait = { workspace = true }
6364
futures = { workspace = true }

crates/core/src/context.rs

Lines changed: 43 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -187,9 +187,11 @@ impl PySessionConfig {
187187
Self::from(self.config.clone().set_str(key, value))
188188
}
189189

190-
fn with_extension(&self, extension: Bound<PyAny>) -> PyResult<Self> {
191-
let capsule = extension.call_method0("__datafusion_extension_options__")?;
192-
let capsule = capsule.cast::<PyCapsule>().map_err(py_datafusion_err)?;
190+
fn with_extension(&self, mut extension: Bound<PyAny>) -> PyResult<Self> {
191+
if extension.hasattr("__datafusion_extension_options__")? {
192+
extension = extension.call_method0("__datafusion_extension_options__")?;
193+
}
194+
let capsule = extension.cast::<PyCapsule>().map_err(py_datafusion_err)?;
193195

194196
validate_pycapsule(capsule, "datafusion_extension_options")?;
195197

@@ -722,6 +724,44 @@ impl PySessionContext {
722724
Ok(())
723725
}
724726

727+
pub fn register_table_options(
728+
&self,
729+
mut extension: Bound<'_, PyAny>,
730+
) -> PyDataFusionResult<()> {
731+
if extension.hasattr("__datafusion_extension_options__")? {
732+
extension = extension.call_method0("__datafusion_extension_options__")?;
733+
}
734+
735+
let capsule = extension.cast::<PyCapsule>().map_err(py_datafusion_err)?;
736+
validate_pycapsule(capsule, "datafusion_extension_options")?;
737+
738+
let data: NonNull<FFI_ExtensionOptions> = capsule
739+
.pointer_checked(Some(c_str!("datafusion_extension_options")))?
740+
.cast();
741+
let mut extension: FFI_ExtensionOptions = unsafe { data.as_ref().clone() };
742+
743+
let st = self.ctx.state_ref();
744+
let mut lock = st.write();
745+
746+
if let Some(prior_extension) = lock
747+
.table_options()
748+
.extensions
749+
.get::<FFI_ExtensionOptions>()
750+
{
751+
extension
752+
.merge(prior_extension)
753+
.map_err(py_datafusion_err)?;
754+
}
755+
756+
lock.table_options_mut().extensions.insert(extension);
757+
758+
for (key, _) in lock.table_options().extensions.iter() {
759+
eprintln!("{key:?}");
760+
}
761+
762+
Ok(())
763+
}
764+
725765
pub fn register_catalog_provider_list(
726766
&self,
727767
mut provider: Bound<PyAny>,

examples/datafusion-ffi-example/Cargo.toml

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,9 @@ arrow-schema = { workspace = true }
3939
async-trait = { workspace = true }
4040
datafusion-python-util.workspace = true
4141
pyo3 = { workspace = true, features = [
42-
"extension-module",
43-
"abi3",
44-
"abi3-py310",
42+
"extension-module",
43+
"abi3",
44+
"abi3-py310",
4545
] }
4646
pyo3-log = { workspace = true }
4747

@@ -53,9 +53,9 @@ name = "datafusion_ffi_example"
5353
crate-type = ["cdylib", "rlib"]
5454

5555
[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" }
56+
datafusion-catalog = { git = "https://github.com/davisp/datafusion", branch = "pd/test/ffi-config" }
57+
datafusion-common = { git = "https://github.com/davisp/datafusion", branch = "pd/test/ffi-config" }
58+
datafusion-functions-aggregate = { git = "https://github.com/davisp/datafusion", branch = "pd/test/ffi-config" }
59+
datafusion-functions-window = { git = "https://github.com/davisp/datafusion", branch = "pd/test/ffi-config" }
60+
datafusion-expr = { git = "https://github.com/davisp/datafusion", branch = "pd/test/ffi-config" }
61+
datafusion-ffi = { git = "https://github.com/davisp/datafusion", branch = "pd/test/ffi-config" }

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,21 +18,25 @@
1818
from __future__ import annotations
1919

2020
from datafusion import SessionContext
21-
from datafusion_ffi_example import MyTableProviderFactory
21+
from datafusion_ffi_example import MyTableProviderFactory, MyTableOptions
2222

2323

2424
def test_table_provider_factory_ffi() -> None:
2525
ctx = SessionContext()
2626
table = MyTableProviderFactory()
2727

2828
ctx.register_table_factory("MY_FORMAT", table)
29+
ctx.register_table_options(MyTableOptions())
2930

3031
# Create a new external table
3132
ctx.sql("""
3233
CREATE EXTERNAL TABLE
3334
foo
3435
STORED AS my_format
35-
LOCATION '';
36+
LOCATION ''
37+
OPTIONS (
38+
'my_format.some_option' '42'
39+
);
3640
""").collect()
3741

3842
# Query the pre-populated table

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

Lines changed: 93 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,12 @@ use pyo3::types::PyCapsule;
1010
use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
1111

1212
/// My own config options.
13-
#[pyclass(name = "MyConfig", module = "datafusion_ffi_example", subclass)]
13+
#[pyclass(
14+
from_py_object,
15+
name = "MyConfig",
16+
module = "datafusion_ffi_example",
17+
subclass
18+
)]
1419
#[derive(Clone, Debug)]
1520
pub struct MyConfig {
1621
/// Should "foo" be replaced by "bar"?
@@ -109,3 +114,90 @@ impl ConfigField for MyConfig {
109114
}
110115
}
111116
}
117+
118+
/// My own table options.
119+
#[pyclass(
120+
from_py_object,
121+
name = "MyTableOptions",
122+
module = "datafusion_ffi_example",
123+
subclass
124+
)]
125+
#[derive(Clone, Debug)]
126+
pub struct MyTableOptions {
127+
/// A random option name for testing
128+
pub some_option: usize,
129+
}
130+
131+
#[pymethods]
132+
impl MyTableOptions {
133+
#[new]
134+
fn new() -> Self {
135+
Self::default()
136+
}
137+
138+
fn __datafusion_extension_options__<'py>(
139+
&self,
140+
py: Python<'py>,
141+
) -> PyResult<Bound<'py, PyCapsule>> {
142+
let name = cr"datafusion_extension_options".into();
143+
144+
let mut config = FFI_ExtensionOptions::default();
145+
config
146+
.add_config(self)
147+
.map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
148+
149+
PyCapsule::new(py, config, Some(name))
150+
}
151+
}
152+
153+
impl Default for MyTableOptions {
154+
fn default() -> Self {
155+
Self { some_option: 1337 }
156+
}
157+
}
158+
159+
impl ConfigExtension for MyTableOptions {
160+
const PREFIX: &'static str = "my_format";
161+
}
162+
163+
impl ExtensionOptions for MyTableOptions {
164+
fn as_any(&self) -> &dyn Any {
165+
self
166+
}
167+
168+
fn as_any_mut(&mut self) -> &mut dyn Any {
169+
self
170+
}
171+
172+
fn cloned(&self) -> Box<dyn ExtensionOptions> {
173+
Box::new(self.clone())
174+
}
175+
176+
fn set(&mut self, key: &str, value: &str) -> datafusion_common::Result<()> {
177+
datafusion_common::config::ConfigField::set(self, key, value)
178+
}
179+
180+
fn entries(&self) -> Vec<ConfigEntry> {
181+
vec![ConfigEntry {
182+
key: "some_option".to_owned(),
183+
value: Some(format!("{}", self.some_option)),
184+
description: "some description",
185+
}]
186+
}
187+
}
188+
189+
impl ConfigField for MyTableOptions {
190+
fn visit<V: Visit>(&self, v: &mut V, _key: &str, _description: &'static str) {
191+
let key = "some_option";
192+
let desc = "some description";
193+
self.some_option.visit(v, key, desc);
194+
}
195+
196+
fn set(&mut self, key: &str, value: &str) -> Result<(), DataFusionError> {
197+
let (key, rem) = key.split_once('.').unwrap_or((key, ""));
198+
match key {
199+
"some_option" => self.some_option.set(rem, value.as_ref()),
200+
_ => config_err!("Config value \"{}\" not found on MyTableOptions", key),
201+
}
202+
}
203+
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +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;
22+
use crate::config::{MyConfig, MyTableOptions};
2323
use crate::scalar_udf::IsNullUDF;
2424
use crate::table_function::MyTableFunction;
2525
use crate::table_provider::MyTableProvider;
@@ -49,5 +49,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
4949
m.add_class::<MySumUDF>()?;
5050
m.add_class::<MyRankUDF>()?;
5151
m.add_class::<MyConfig>()?;
52+
m.add_class::<MyTableOptions>()?;
5253
Ok(())
5354
}

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use std::sync::Arc;
1919

2020
use async_trait::async_trait;
2121
use datafusion_catalog::{Session, TableProvider, TableProviderFactory};
22-
use datafusion_common::error::Result as DataFusionResult;
22+
use datafusion_common::error::{DataFusionError, Result as DataFusionResult};
2323
use datafusion_expr::CreateExternalTable;
2424
use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory;
2525
use datafusion_python_util::ffi_logical_codec_from_pycapsule;
@@ -42,8 +42,15 @@ impl TableProviderFactory for ExampleTableProviderFactory {
4242
async fn create(
4343
&self,
4444
_state: &dyn Session,
45-
_cmd: &CreateExternalTable,
45+
cmd: &CreateExternalTable,
4646
) -> DataFusionResult<Arc<dyn TableProvider>> {
47+
eprintln!("OPTIONS: {:#?}", cmd.options);
48+
let val = cmd.options.get("my_format.some_option");
49+
if val != Some(&"42".to_owned()) {
50+
return Err(DataFusionError::Internal(format!(
51+
"Expected '42', got '{val:?}'"
52+
)));
53+
}
4754
Ok(catalog_provider::my_table())
4855
}
4956
}

python/datafusion/context.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -852,6 +852,12 @@ def register_table_factory(
852852
"""
853853
self.ctx.register_table_factory(format, factory)
854854

855+
def register_table_options(
856+
self,
857+
factory: Any,
858+
) -> None:
859+
self.ctx.register_table_options(factory)
860+
855861
def catalog_names(self) -> set[str]:
856862
"""Returns the list of catalogs in this context."""
857863
return self.ctx.catalog_names()

0 commit comments

Comments
 (0)