From 9ba147c1697a0cf8a2f4d09e618b86064fbcb2b7 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 31 Jul 2026 15:38:53 -0400 Subject: [PATCH 01/30] Add FFI query planner support AI Disclosure: This code was written in part by an AI agent.: --- Cargo.lock | 78 ++++++++-------- Cargo.toml | 22 +++-- crates/core/Cargo.toml | 1 + crates/core/src/context.rs | 59 +++++++++++- crates/util/src/lib.rs | 33 +++++++ examples/datafusion-ffi-example/Cargo.toml | 1 + .../python/tests/_test_query_planner.py | 39 ++++++++ examples/datafusion-ffi-example/src/lib.rs | 3 + .../src/query_planner.rs | 92 +++++++++++++++++++ python/datafusion/context.py | 43 +++++++++ python/tests/test_context.py | 23 +++++ 11 files changed, 344 insertions(+), 50 deletions(-) create mode 100644 examples/datafusion-ffi-example/python/tests/_test_query_planner.py create mode 100644 examples/datafusion-ffi-example/src/query_planner.rs diff --git a/Cargo.lock b/Cargo.lock index d34862ac7..ed4683fbe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -791,7 +791,7 @@ dependencies = [ [[package]] name = "datafusion" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "arrow-schema", @@ -844,7 +844,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "async-trait", @@ -868,7 +868,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "async-trait", @@ -891,7 +891,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "arrow-ipc", @@ -917,7 +917,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "futures", "log", @@ -927,7 +927,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "async-compression", @@ -963,7 +963,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "arrow-ipc", @@ -986,7 +986,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "arrow-avro", @@ -1004,7 +1004,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "async-trait", @@ -1026,7 +1026,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "async-trait", @@ -1048,7 +1048,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "arrow-schema", @@ -1079,12 +1079,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" [[package]] name = "datafusion-execution" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "arrow-buffer", @@ -1109,7 +1109,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "arrow-schema", @@ -1131,7 +1131,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "datafusion-common", @@ -1142,7 +1142,7 @@ dependencies = [ [[package]] name = "datafusion-ffi" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "arrow-schema", @@ -1188,6 +1188,7 @@ dependencies = [ "datafusion-functions-window", "datafusion-proto", "datafusion-python-util", + "datafusion-session", "pyo3", "pyo3-build-config", "pyo3-log", @@ -1196,7 +1197,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "arrow-buffer", @@ -1227,7 +1228,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "datafusion-common", @@ -1247,7 +1248,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "datafusion-common", @@ -1258,7 +1259,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "arrow-ord", @@ -1282,7 +1283,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "async-trait", @@ -1297,7 +1298,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "datafusion-common", @@ -1313,7 +1314,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1322,7 +1323,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "datafusion-doc", "quote", @@ -1332,7 +1333,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "chrono", @@ -1351,7 +1352,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "datafusion-common", @@ -1373,7 +1374,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "datafusion-common", @@ -1387,7 +1388,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "chrono", @@ -1404,7 +1405,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "datafusion-common", @@ -1423,7 +1424,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "arrow-data", @@ -1459,7 +1460,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "chrono", @@ -1486,7 +1487,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "datafusion-common", @@ -1496,7 +1497,7 @@ dependencies = [ [[package]] name = "datafusion-proto-models" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "datafusion-proto-common", "prost", @@ -1505,7 +1506,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "datafusion-common", @@ -1530,6 +1531,7 @@ dependencies = [ "datafusion-ffi", "datafusion-proto", "datafusion-python-util", + "datafusion-session", "datafusion-spark", "datafusion-substrait", "futures", @@ -1565,7 +1567,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow-schema", "async-trait", @@ -1579,7 +1581,7 @@ dependencies = [ [[package]] name = "datafusion-spark" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "bigdecimal", @@ -1608,7 +1610,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "arrow", "bigdecimal", @@ -1627,7 +1629,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 362159913..01dce04a4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -50,6 +50,7 @@ datafusion-functions-aggregate = { version = "54.1.0" } datafusion-functions-window = { version = "54.1.0" } datafusion-spark = { version = "54.1.0" } datafusion-expr = { version = "54.1.0" } +datafusion-session = { version = "54.1.0" } prost = "0.14.3" serde_json = "1" uuid = { version = "1.23" } @@ -72,13 +73,14 @@ codegen-units = 2 # We cannot publish to crates.io with any patches in the below section. Developers # must remove any entries in this section before creating a release candidate. [patch.crates-io] -datafusion = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" } -datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" } -datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" } -datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" } -datafusion-catalog = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" } -datafusion-common = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" } -datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" } -datafusion-functions-window = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" } -datafusion-spark = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" } -datafusion-expr = { git = "https://github.com/apache/datafusion", rev = "dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" } +datafusion = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } +datafusion-substrait = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } +datafusion-proto = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } +datafusion-ffi = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } +datafusion-catalog = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } +datafusion-common = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } +datafusion-functions-aggregate = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } +datafusion-functions-window = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } +datafusion-spark = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } +datafusion-expr = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } +datafusion-session = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index c5f1e0167..91a1d5f77 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -54,6 +54,7 @@ datafusion-substrait = { workspace = true, optional = true } datafusion-proto = { workspace = true } datafusion-ffi = { workspace = true } datafusion-spark = { workspace = true, features = ["core"] } +datafusion-session = { workspace = true } prost = { workspace = true } # keep in line with `datafusion-substrait` serde_json = { workspace = true } uuid = { workspace = true, features = ["v4"] } diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 7bbeed2f1..ab6bbbbcc 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -24,6 +24,7 @@ use std::sync::Arc; use arrow::array::RecordBatchReader; use arrow::ffi_stream::ArrowArrayStreamReader; use arrow::pyarrow::FromPyArrow; +use async_trait::async_trait; use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion::arrow::pyarrow::PyArrowType; use datafusion::arrow::record_batch::RecordBatch; @@ -36,7 +37,7 @@ use datafusion::datasource::listing::{ }; use datafusion::datasource::{MemTable, TableProvider}; use datafusion::execution::context::{ - DataFilePaths, SQLOptions, SessionConfig, SessionContext, TaskContext, + DataFilePaths, QueryPlanner, SQLOptions, SessionConfig, SessionContext, TaskContext, }; use datafusion::execution::disk_manager::DiskManagerMode; use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool, UnboundedMemoryPool}; @@ -44,6 +45,8 @@ use datafusion::execution::options::{ArrowReadOptions, ReadOptions}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::execution::{FunctionRegistry, TaskContextProvider}; +use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::{ AvroReadOptions, CsvReadOptions, DataFrame, JsonReadOptions, ParquetReadOptions, }; @@ -53,15 +56,18 @@ 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::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_ffi::query_planner::FFI_QueryPlanner; use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory; use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_python_util::{ create_logical_extension_capsule, create_physical_extension_capsule, - ffi_logical_codec_from_pycapsule, get_global_ctx, get_tokio_runtime, + create_query_planner_capsule, ffi_logical_codec_from_pycapsule, + ffi_query_planner_from_pycapsule, get_global_ctx, get_tokio_runtime, physical_codec_from_pycapsule, physical_optimizer_rule_from_pycapsule, spawn_future, wait_for_future, }; +use datafusion_session::Session; use object_store::ObjectStore; use pyo3::IntoPyObjectExt; use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyValueError}; @@ -221,6 +227,25 @@ impl PySessionConfig { } } +#[derive(Debug, Clone)] +struct PythonQueryPlanner { + planner: FFI_QueryPlanner, +} + +#[async_trait] +impl QueryPlanner for PythonQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> datafusion::common::Result> { + let runtime = get_tokio_runtime().handle().clone(); + self.planner + .create_physical_plan_with_session_runtime(logical_plan, session, Some(runtime)) + .await + } +} + /// Runtime options for a SessionContext #[pyclass( from_py_object, @@ -1211,6 +1236,23 @@ impl PySessionContext { Ok(()) } + pub fn with_query_planner(&self, planner: Bound<'_, PyAny>) -> PyDataFusionResult { + let mut planner = ffi_query_planner_from_pycapsule(&planner)?; + planner.logical_codec = self.ffi_logical_codec().as_ref().clone(); + planner.physical_codec = self.ffi_physical_codec().as_ref().clone(); + let planner = Arc::new(PythonQueryPlanner { planner }); + let state = SessionStateBuilder::new_from_existing(self.ctx.state()) + .with_query_planner(planner) + .build(); + let ctx = Arc::new(SessionContext::new_with_state(state)); + + Ok(Self { + ctx, + logical_codec: Arc::clone(&self.logical_codec), + physical_codec: Arc::clone(&self.physical_codec), + }) + } + pub fn table_provider(&self, name: &str, py: Python) -> PyResult { let provider = wait_for_future(py, self.ctx.table_provider(name)) // Outer error: runtime/async failure @@ -1385,6 +1427,19 @@ impl PySessionContext { create_logical_extension_capsule(py, ffi.as_ref()) } + pub fn __datafusion_query_planner__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let planner = Arc::clone(self.ctx.state().query_planner()); + let ffi = FFI_QueryPlanner::new_with_ffi_codecs( + planner, + self.ffi_logical_codec().as_ref().clone(), + self.ffi_physical_codec().as_ref().clone(), + ); + create_query_planner_capsule(py, &ffi) + } + pub fn with_logical_extension_codec<'py>( &self, codec: Bound<'py, PyAny>, diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index 9327d7f2f..7375a034c 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -29,6 +29,7 @@ use datafusion_ffi::execution::FFI_TaskContextProvider; use datafusion_ffi::physical_optimizer::FFI_PhysicalOptimizerRule; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_ffi::query_planner::FFI_QueryPlanner; use datafusion_ffi::table_provider::FFI_TableProvider; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use pyo3::exceptions::{PyImportError, PyTypeError, PyValueError}; @@ -231,6 +232,38 @@ pub fn ffi_logical_codec_from_pycapsule(obj: Bound) -> PyResult( + py: Python<'py>, + planner: &FFI_QueryPlanner, +) -> PyResult> { + PyCapsule::new_with_value(py, planner.clone(), cr"datafusion_query_planner") +} + +pub fn ffi_query_planner_from_pycapsule(obj: &Bound) -> PyResult { + let attr_name = "__datafusion_query_planner__"; + let capsule = if obj.hasattr(attr_name)? { + obj.getattr(attr_name)?.call0()? + } else { + obj.clone() + }; + + let capsule = capsule.cast::()?; + validate_pycapsule(capsule, "datafusion_query_planner")?; + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_query_planner"))? + .cast(); + let planner = unsafe { data.as_ref() }; + let planner_version = unsafe { (planner.version)() }; + let expected_version = datafusion_ffi::version(); + if planner_version != expected_version { + return Err(PyImportError::new_err(format!( + "Incompatible DataFusion query planner version {planner_version}; expected major version {expected_version}." + ))); + } + + Ok(planner.clone()) +} + pub fn create_physical_extension_capsule<'py>( py: Python<'py>, codec: &FFI_PhysicalExtensionCodec, diff --git a/examples/datafusion-ffi-example/Cargo.toml b/examples/datafusion-ffi-example/Cargo.toml index ffc839d56..9132b2845 100644 --- a/examples/datafusion-ffi-example/Cargo.toml +++ b/examples/datafusion-ffi-example/Cargo.toml @@ -34,6 +34,7 @@ datafusion-functions-window = { workspace = true } datafusion-expr = { workspace = true } datafusion-ffi = { workspace = true } datafusion-proto = { workspace = true } +datafusion-session = { workspace = true } arrow = { workspace = true } arrow-array = { workspace = true } diff --git a/examples/datafusion-ffi-example/python/tests/_test_query_planner.py b/examples/datafusion-ffi-example/python/tests/_test_query_planner.py new file mode 100644 index 000000000..91b6e9c4e --- /dev/null +++ b/examples/datafusion-ffi-example/python/tests/_test_query_planner.py @@ -0,0 +1,39 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import pytest +from datafusion import SessionContext +from datafusion_ffi_example import MyQueryPlanner + + +@pytest.mark.parametrize("raw_capsule", [False, True]) +def test_ffi_query_planner_runs_during_planning(raw_capsule: bool): + """A query planner imported from another library creates the physical plan.""" + planner = MyQueryPlanner() + exported_planner = ( + planner.__datafusion_query_planner__() if raw_capsule else planner + ) + ctx = SessionContext().with_query_planner(exported_planner) + + before = planner.plan_calls() + result = ctx.sql("SELECT 1 AS value").collect() + after = planner.plan_calls() + + assert after > before + assert result == [] diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index eccf7b81a..d25e69253 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -23,6 +23,7 @@ use crate::config::MyConfig; use crate::logical_extension_codec::MyLogicalExtensionCodec; use crate::physical_extension_codec::MyPhysicalExtensionCodec; use crate::physical_optimizer::MyPhysicalOptimizerRule; +use crate::query_planner::MyQueryPlanner; use crate::scalar_udf::IsNullUDF; use crate::table_function::MyTableFunction; use crate::table_provider::MyTableProvider; @@ -35,6 +36,7 @@ pub(crate) mod config; pub(crate) mod logical_extension_codec; pub(crate) mod physical_extension_codec; pub(crate) mod physical_optimizer; +pub(crate) mod query_planner; pub(crate) mod scalar_udf; pub(crate) mod table_function; pub(crate) mod table_provider; @@ -58,5 +60,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-example/src/query_planner.rs b/examples/datafusion-ffi-example/src/query_planner.rs new file mode 100644 index 000000000..4637074dd --- /dev/null +++ b/examples/datafusion-ffi-example/src/query_planner.rs @@ -0,0 +1,92 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use arrow::datatypes::{DataType, Field, Schema}; +use async_trait::async_trait; +use datafusion::execution::TaskContextProvider; +use datafusion::execution::context::{QueryPlanner, SessionContext}; +use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion_ffi::query_planner::FFI_QueryPlanner; +use datafusion_python_util::get_tokio_runtime; +use datafusion_session::Session; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +#[derive(Debug)] +struct CountingQueryPlanner { + plan_calls: Arc, +} + +#[async_trait] +impl QueryPlanner for CountingQueryPlanner { + async fn create_physical_plan( + &self, + _logical_plan: &LogicalPlan, + _session: &dyn Session, + ) -> datafusion::common::Result> { + self.plan_calls.fetch_add(1, Ordering::SeqCst); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + Ok(Arc::new(EmptyExec::new(schema))) + } +} + +/// Python-visible query planner used to test planning across a real FFI boundary. +#[pyclass( + from_py_object, + name = "MyQueryPlanner", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Default, Clone)] +pub(crate) struct MyQueryPlanner { + plan_calls: Arc, +} + +#[pymethods] +impl MyQueryPlanner { + #[new] + fn new() -> Self { + Self::default() + } + + fn plan_calls(&self) -> usize { + self.plan_calls.load(Ordering::SeqCst) + } + + fn __datafusion_query_planner__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let planner: Arc = Arc::new(CountingQueryPlanner { + plan_calls: Arc::clone(&self.plan_calls), + }); + let runtime = get_tokio_runtime().handle().clone(); + let ctx_provider = Arc::new(SessionContext::new()) as Arc; + let ffi = FFI_QueryPlanner::new(planner, Some(runtime), &ctx_provider, None, None); + + PyCapsule::new_with_value(py, ffi, cr"datafusion_query_planner") + } +} diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 94b2bb1c6..f92c57b28 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -145,6 +145,16 @@ class PhysicalOptimizerRuleExportable(Protocol): def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 +class QueryPlannerExportable(Protocol): + """Type hint for object that has a __datafusion_query_planner__ PyCapsule. + + The method returns a PyCapsule wrapping an ``FFI_QueryPlanner``, typically + produced by a separate compiled extension. + """ + + def __datafusion_query_planner__(self) -> object: ... # noqa: D105 + + class SessionConfig: """Session configuration options.""" @@ -1759,6 +1769,35 @@ def add_physical_optimizer_rule( """ self.ctx.add_physical_optimizer_rule(rule) + def with_query_planner( + self, planner: QueryPlannerExportable | _PyCapsule + ) -> SessionContext: + """Create a new session context with a custom query planner. + + The planner is imported through its ``__datafusion_query_planner__`` + PyCapsule. The returned context preserves the existing session state and + its logical and physical extension codec settings. + + Args: + planner: Object exposing ``__datafusion_query_planner__`` or a raw + ``datafusion_query_planner`` PyCapsule. + + Returns: + A new context that uses the specified query planner. + + Examples: + >>> from my_extension import DistributedQueryPlanner # doctest: +SKIP + >>> ctx = SessionContext() + >>> planner = DistributedQueryPlanner() # doctest: +SKIP + >>> planner_ctx = ctx.with_query_planner(planner) # doctest: +SKIP + >>> query = planner_ctx.sql("SELECT * FROM remote_table") # doctest: +SKIP + >>> query.collect() # doctest: +SKIP + """ + new_internal = self.ctx.with_query_planner(planner) + new = SessionContext.__new__(SessionContext) + new.ctx = new_internal + return new + def table_provider(self, name: str) -> Table: """Return the :py:class:`~datafusion.catalog.Table` for the given table name. @@ -2182,6 +2221,10 @@ def __datafusion_logical_extension_codec__(self) -> Any: """Access the PyCapsule FFI_LogicalExtensionCodec.""" return self.ctx.__datafusion_logical_extension_codec__() + def __datafusion_query_planner__(self) -> Any: + """Access the ``FFI_QueryPlanner`` PyCapsule for the current planner.""" + return self.ctx.__datafusion_query_planner__() + def with_logical_extension_codec( self, codec: LogicalExtensionCodecExportable | _PyCapsule ) -> SessionContext: diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 7d038c7a5..6e6eaadbe 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -14,6 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import ctypes import datetime as dt import gzip import pathlib @@ -731,6 +732,28 @@ def test_remove_optimizer_rule(ctx): assert ctx.remove_optimizer_rule("nonexistent_rule") is False +def test_with_query_planner_rejects_wrong_capsule(ctx): + with pytest.raises(ValueError, match="datafusion_query_planner"): + ctx.with_query_planner(ctx.__datafusion_task_context_provider__()) + + +def test_with_query_planner_capsule(ctx): + capsule = ctx.__datafusion_query_planner__() + get_name = ctypes.pythonapi.PyCapsule_GetName + get_name.argtypes = [ctypes.py_object] + get_name.restype = ctypes.c_char_p + assert get_name(capsule) == b"datafusion_query_planner" + + ctx.register_record_batches( + "query_planner_test", + [[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]], + ) + planner_context = ctx.with_query_planner(capsule) + assert planner_context.table_exist("query_planner_test") + batches = planner_context.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + def test_table_provider(ctx): batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) ctx.register_record_batches("provider_test", [[batch]]) From 47aea03d5c62cd716462e3a40289b9c76d910447 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 31 Jul 2026 16:45:03 -0400 Subject: [PATCH 02/30] Add three-library FFI planner example AI Disclosure: This code was written in part by an AI agent.: --- .github/workflows/build.yml | 16 +- .github/workflows/test.yml | 2 + Cargo.lock | 15 ++ Cargo.toml | 7 +- crates/core/src/context.rs | 102 ++++++++- docs/source/contributor-guide/ffi.md | 27 +++ examples/README.md | 9 + examples/datafusion-ffi-example/Cargo.toml | 1 - examples/datafusion-ffi-example/README.md | 31 +++ examples/datafusion-ffi-example/src/lib.rs | 3 - .../src/logical_extension_codec.rs | 103 ++++++--- .../src/physical_extension_codec.rs | 92 +++++++- .../src/query_planner.rs | 92 -------- .../Cargo.toml | 49 +++++ .../README.md | 41 ++++ .../build.rs | 20 ++ .../pyproject.toml} | 33 ++- .../_test_three_library_query_planner.py | 120 +++++++++++ .../src/config.rs | 112 ++++++++++ .../src/lib.rs | 32 +++ .../src/planner.rs | 198 ++++++++++++++++++ python/datafusion/context.py | 3 +- 22 files changed, 943 insertions(+), 165 deletions(-) create mode 100644 examples/datafusion-ffi-example/README.md delete mode 100644 examples/datafusion-ffi-example/src/query_planner.rs create mode 100644 examples/datafusion-ffi-query-planner-example/Cargo.toml create mode 100644 examples/datafusion-ffi-query-planner-example/README.md create mode 100644 examples/datafusion-ffi-query-planner-example/build.rs rename examples/{datafusion-ffi-example/python/tests/_test_query_planner.py => datafusion-ffi-query-planner-example/pyproject.toml} (52%) create mode 100644 examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py create mode 100644 examples/datafusion-ffi-query-planner-example/src/config.rs create mode 100644 examples/datafusion-ffi-query-planner-example/src/lib.rs create mode 100644 examples/datafusion-ffi-query-planner-example/src/planner.rs diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index c35801b11..d7af9b663 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -186,7 +186,7 @@ jobs: manylinux: "2_28" # FFI test wheel only needs to be built once per platform; gate to abi3. - - name: Build FFI test library + - name: Build FFI provider test library if: matrix.python-tag == 'abi3' uses: PyO3/maturin-action@v1 with: @@ -196,6 +196,16 @@ jobs: args: --out dist rustup-components: rust-std + - name: Build FFI query planner test library + if: matrix.python-tag == 'abi3' + uses: PyO3/maturin-action@v1 + with: + target: x86_64-unknown-linux-gnu + manylinux: "2_28" + working-directory: examples/datafusion-ffi-query-planner-example + args: --out dist + rustup-components: rust-std + - name: Archive wheels uses: actions/upload-artifact@v7 with: @@ -207,7 +217,9 @@ jobs: uses: actions/upload-artifact@v7 with: name: test-ffi-manylinux-x86_64 - path: examples/datafusion-ffi-example/dist/* + path: | + examples/datafusion-ffi-example/dist/* + examples/datafusion-ffi-query-planner-example/dist/* # ============================================ # Build - Linux ARM64 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 558e751c8..21ecd56f7 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -121,6 +121,8 @@ jobs: run: | cd examples/datafusion-ffi-example uv run --no-project pytest python/tests/_test*.py + cd ../datafusion-ffi-query-planner-example + uv run --no-project pytest python/tests/_test*.py - name: Run tpchgen-cli to create 1 Gb dataset if: matrix.wheel-tag == 'abi3' diff --git a/Cargo.lock b/Cargo.lock index ed4683fbe..e03c75f15 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1188,6 +1188,21 @@ dependencies = [ "datafusion-functions-window", "datafusion-proto", "datafusion-python-util", + "pyo3", + "pyo3-build-config", + "pyo3-log", +] + +[[package]] +name = "datafusion-ffi-query-planner-example" +version = "54.0.0" +dependencies = [ + "async-trait", + "datafusion", + "datafusion-catalog", + "datafusion-common", + "datafusion-ffi", + "datafusion-python-util", "datafusion-session", "pyo3", "pyo3-build-config", diff --git a/Cargo.toml b/Cargo.toml index 01dce04a4..0f5aa8e27 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,7 +27,12 @@ edition = "2024" rust-version = "1.88" [workspace] -members = ["crates/core", "crates/util", "examples/datafusion-ffi-example"] +members = [ + "crates/core", + "crates/util", + "examples/datafusion-ffi-example", + "examples/datafusion-ffi-query-planner-example", +] resolver = "3" [workspace.dependencies] diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index ab6bbbbcc..9853fe9b9 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -227,9 +227,34 @@ impl PySessionConfig { } } +#[derive(Debug)] +struct PlanningTaskContextProvider(Arc); + +impl TaskContextProvider for PlanningTaskContextProvider { + fn task_ctx(&self) -> Arc { + Arc::clone(&self.0) + } +} + #[derive(Debug, Clone)] struct PythonQueryPlanner { planner: FFI_QueryPlanner, + logical_codec: Arc, + physical_codec: Arc, +} + +impl PythonQueryPlanner { + fn with_codecs( + &self, + logical_codec: Arc, + physical_codec: Arc, + ) -> Self { + Self { + planner: self.planner.clone(), + logical_codec, + physical_codec, + } + } } #[async_trait] @@ -240,9 +265,35 @@ impl QueryPlanner for PythonQueryPlanner { session: &dyn Session, ) -> datafusion::common::Result> { let runtime = get_tokio_runtime().handle().clone(); - self.planner + let ctx_provider = Arc::new(PlanningTaskContextProvider(session.task_ctx())); + let (logical_codec, physical_codec) = { + let dyn_ctx_provider: Arc = ctx_provider.clone(); + let logical_codec: Arc = + Arc::clone(&self.logical_codec) as Arc; + let physical_codec: Arc = + Arc::clone(&self.physical_codec) as Arc; + ( + FFI_LogicalExtensionCodec::new( + logical_codec, + Some(runtime.clone()), + &dyn_ctx_provider, + ), + FFI_PhysicalExtensionCodec::new( + physical_codec, + Some(runtime.clone()), + &dyn_ctx_provider, + ), + ) + }; + + let mut planner = self.planner.clone(); + planner.logical_codec = logical_codec; + planner.physical_codec = physical_codec; + let result = planner .create_physical_plan_with_session_runtime(logical_plan, session, Some(runtime)) - .await + .await; + drop(ctx_provider); + result } } @@ -1237,10 +1288,12 @@ impl PySessionContext { } pub fn with_query_planner(&self, planner: Bound<'_, PyAny>) -> PyDataFusionResult { - let mut planner = ffi_query_planner_from_pycapsule(&planner)?; - planner.logical_codec = self.ffi_logical_codec().as_ref().clone(); - planner.physical_codec = self.ffi_physical_codec().as_ref().clone(); - let planner = Arc::new(PythonQueryPlanner { planner }); + let planner = ffi_query_planner_from_pycapsule(&planner)?; + let planner = Arc::new(PythonQueryPlanner { + planner, + logical_codec: Arc::clone(&self.logical_codec), + physical_codec: Arc::clone(&self.physical_codec), + }); let state = SessionStateBuilder::new_from_existing(self.ctx.state()) .with_query_planner(planner) .build(); @@ -1448,10 +1501,13 @@ impl PySessionContext { let inner: Arc = (&inner_ffi).into(); let logical_codec = Arc::new(PythonLogicalCodec::new(inner)); + let physical_codec = Arc::clone(&self.physical_codec); + let ctx = self + .ctx_with_query_planner_codecs(Arc::clone(&logical_codec), Arc::clone(&physical_codec)); Ok(Self { - ctx: Arc::clone(&self.ctx), + ctx, logical_codec, - physical_codec: Arc::clone(&self.physical_codec), + physical_codec, }) } @@ -1470,9 +1526,12 @@ impl PySessionContext { let inner = physical_codec_from_pycapsule(&codec)?; let physical_codec = Arc::new(PythonPhysicalCodec::new(inner)); + let logical_codec = Arc::clone(&self.logical_codec); + let ctx = self + .ctx_with_query_planner_codecs(Arc::clone(&logical_codec), Arc::clone(&physical_codec)); Ok(Self { - ctx: Arc::clone(&self.ctx), - logical_codec: Arc::clone(&self.logical_codec), + ctx, + logical_codec, physical_codec, }) } @@ -1486,8 +1545,10 @@ impl PySessionContext { PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner())) .with_python_udf_inlining(enabled), ); + let ctx = self + .ctx_with_query_planner_codecs(Arc::clone(&logical_codec), Arc::clone(&physical_codec)); Self { - ctx: Arc::clone(&self.ctx), + ctx, logical_codec, physical_codec, } @@ -1495,6 +1556,25 @@ impl PySessionContext { } impl PySessionContext { + fn ctx_with_query_planner_codecs( + &self, + logical_codec: Arc, + physical_codec: Arc, + ) -> Arc { + let state = self.ctx.state(); + let query_planner = state.query_planner(); + let planner_any: &dyn std::any::Any = query_planner.as_ref(); + let Some(planner) = planner_any.downcast_ref::() else { + return Arc::clone(&self.ctx); + }; + + let planner = Arc::new(planner.with_codecs(logical_codec, physical_codec)); + let state = SessionStateBuilder::new_from_existing(self.ctx.state()) + .with_query_planner(planner) + .build(); + Arc::new(SessionContext::new_with_state(state)) + } + async fn _table(&self, name: &str) -> datafusion::common::Result { self.ctx.table(name).await } diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index bf65cad2a..d13eda4d8 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -232,6 +232,31 @@ extension that has been written using this approach and the most thoroughly impl As we continue to expose more of the DataFusion features, we intend to follow this same design pattern. +## Query Planners Across Multiple Libraries + +A query can involve three independent native libraries: `datafusion-python`, a library +that owns table providers or functions, and a library that owns the query planner. The +examples use two separate extension crates so each role has a distinct shared-library +identity: + +- [`datafusion-ffi-example`] owns providers, functions, and their codecs. +- [`datafusion-ffi-query-planner-example`] owns the planner and its configuration. + +The `SessionContext` owns the codecs used for the exchange and supplies them to the +foreign planner. This lets the planner decode provider-owned objects and lets +`datafusion-python` decode the physical plan returned by the planner. The examples use +process-local tokens to demonstrate ownership; production codecs should serialize +durable metadata instead. + +The current Python API has one external logical codec and one external physical codec. +Installing another codec replaces the prior codec rather than composing a registry. +The example therefore has one external codec owner, and the planner uses built-in +physical nodes. Install the provider codecs before the planner where possible. + +The current FFI logical codec supports providers and UDFs but not arbitrary custom +`LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and +local build commands. + ## Alternative Approach Suppose you needed to expose some other features of DataFusion and you could not wait @@ -257,3 +282,5 @@ At the time of this writing, the FFI features are under active development. To s the latest status, we recommend reviewing the code in the [datafusion-ffi] crate. [datafusion-ffi]: https://crates.io/crates/datafusion-ffi +[`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example +[`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example diff --git a/examples/README.md b/examples/README.md index e0e3056d9..7bbb45dcf 100644 --- a/examples/README.md +++ b/examples/README.md @@ -49,6 +49,15 @@ Here is a direct link to the file used in the examples: - [Fan out distinct expressions to a multiprocessing pool](./multiprocessing_pickle_expr.py) - [Distribute expression evaluation across Ray actors](./ray_pickle_expr.py) +### Rust FFI Extensions + +- [Table providers, functions, and codecs](./datafusion-ffi-example/) +- [Independent query planner and planner configuration](./datafusion-ffi-query-planner-example/) + +These two crates form a three-library interoperability example with +`datafusion-python`. They are separate shared libraries so the tests exercise real FFI +type and codec boundaries rather than same-library Rust downcasts. + ### Substrait Support - [Serialize query plans using Substrait](./substrait.py) diff --git a/examples/datafusion-ffi-example/Cargo.toml b/examples/datafusion-ffi-example/Cargo.toml index 9132b2845..ffc839d56 100644 --- a/examples/datafusion-ffi-example/Cargo.toml +++ b/examples/datafusion-ffi-example/Cargo.toml @@ -34,7 +34,6 @@ datafusion-functions-window = { workspace = true } datafusion-expr = { workspace = true } datafusion-ffi = { workspace = true } datafusion-proto = { workspace = true } -datafusion-session = { workspace = true } arrow = { workspace = true } arrow-array = { workspace = true } diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md new file mode 100644 index 000000000..4a35e75da --- /dev/null +++ b/examples/datafusion-ffi-example/README.md @@ -0,0 +1,31 @@ +# DataFusion Python FFI provider example + +This crate is the **provider library** in the three-library query-planning example. It exports table providers, functions, and the logical and physical codecs needed to serialize objects owned by this library. The companion planner is in [`../datafusion-ffi-query-planner-example`](../datafusion-ffi-query-planner-example/). + +The example intentionally uses separate `cdylib` crates for these roles: + +1. **A — `datafusion-python`:** owns the `SessionContext` and executes the result. +2. **B — this crate:** owns table providers, functions, and provider execution plans. +3. **C — the planner crate:** receives the logical plan and returns a physical plan. + +Separate shared libraries guarantee distinct DataFusion library markers. This catches type-identity mistakes that a planner and provider compiled into one shared library would hide. + +## Codec behavior + +`MyLogicalExtensionCodec` serializes this example's in-memory table providers, and `MyPhysicalExtensionCodec` serializes provider-owned memory scans and opaque FFI wrappers around them. Both use documented, process-local, one-shot token registries. The registries make ownership and callback routing visible without pretending to be a portable format. They assume trusted in-process payloads and consume each token during decoding. A production provider should instead encode durable metadata from which its provider and plans can be reconstructed. + +The example codecs do not inspect the callback `TaskContext`. A production codec that depends on session configuration or registered functions must ensure its exported FFI codec is bound to, and retains, the appropriate host `TaskContextProvider`. + +The current Python API installs one external logical codec and one external physical codec. It does not yet compose codecs from several independent plugin owners. This example therefore makes the provider library the sole external codec owner; the planner uses built-in physical nodes and receives the provider codecs from the host. + +Register both provider codecs before installing the planner: + +```python +ctx = ctx.with_logical_extension_codec(provider_logical_codec) +ctx = ctx.with_physical_extension_codec(provider_physical_codec) +ctx = ctx.with_query_planner(planner) +``` + +Derived contexts also rebind an installed planner when codecs change, but planner-last order is recommended because it states the ownership flow clearly. + +Arbitrary custom `LogicalPlan::Extension` nodes are not supported by the current DataFusion FFI logical codec. This example covers foreign table providers, UDFs, and physical execution plans only. diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index d25e69253..eccf7b81a 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -23,7 +23,6 @@ use crate::config::MyConfig; use crate::logical_extension_codec::MyLogicalExtensionCodec; use crate::physical_extension_codec::MyPhysicalExtensionCodec; use crate::physical_optimizer::MyPhysicalOptimizerRule; -use crate::query_planner::MyQueryPlanner; use crate::scalar_udf::IsNullUDF; use crate::table_function::MyTableFunction; use crate::table_provider::MyTableProvider; @@ -36,7 +35,6 @@ pub(crate) mod config; pub(crate) mod logical_extension_codec; pub(crate) mod physical_extension_codec; pub(crate) mod physical_optimizer; -pub(crate) mod query_planner; pub(crate) mod scalar_udf; pub(crate) mod table_function; pub(crate) mod table_provider; @@ -60,6 +58,5 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-example/src/logical_extension_codec.rs b/examples/datafusion-ffi-example/src/logical_extension_codec.rs index 8c3976d37..0474a8d39 100644 --- a/examples/datafusion-ffi-example/src/logical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/logical_extension_codec.rs @@ -15,11 +15,14 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::collections::HashMap; +use std::fmt; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; use arrow::datatypes::SchemaRef; -use datafusion::common::{Result, TableReference}; +use datafusion::catalog::MemTable; +use datafusion::common::{DataFusionError, Result, TableReference}; use datafusion::datasource::TableProvider; use datafusion::execution::{TaskContext, TaskContextProvider}; use datafusion::logical_expr::{Extension, LogicalPlan, ScalarUDF}; @@ -30,25 +33,49 @@ use datafusion_python_util::get_tokio_runtime; use pyo3::prelude::*; use pyo3::types::PyCapsule; -/// Tracks how often each `try_*_udf` entry point fires. Surface for -/// Python tests to assert the session routed UDF -/// encode/decode through this user-supplied codec rather than the -/// upstream default. +const TABLE_PROVIDER_TOKEN: &[u8] = b"DFPYEXTP"; +static NEXT_TABLE_PROVIDER_ID: AtomicU64 = AtomicU64::new(1); +static TABLE_PROVIDERS: OnceLock>>> = OnceLock::new(); + +fn table_providers() -> &'static Mutex>> { + TABLE_PROVIDERS.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn token_id(buf: &[u8], prefix: &[u8]) -> Option { + let id: [u8; 8] = buf.strip_prefix(prefix)?.try_into().ok()?; + Some(u64::from_le_bytes(id)) +} + #[derive(Debug, Default)] pub(crate) struct CallCounters { pub encode_udf: AtomicUsize, pub decode_udf: AtomicUsize, + pub encode_table_provider: AtomicUsize, + pub decode_table_provider: AtomicUsize, } -/// Minimal user-supplied `LogicalExtensionCodec` for integration tests. -/// Delegates everything to `DefaultLogicalExtensionCodec` and bumps -/// counters on the UDF entry points so tests can prove the wrapper -/// installed via `SessionContext.with_logical_extension_codec(...)` -/// actually gets consulted. -#[derive(Debug)] +/// Example codec for objects owned by this extension library. +/// +/// The table-provider token registry is intentionally process-local. It is a compact +/// example of preserving Rust type identity across three loaded libraries, not a +/// network serialization format. Production libraries should encode reconstructible +/// provider metadata rather than retaining objects in a global registry. struct CountingLogicalExtensionCodec { inner: DefaultLogicalExtensionCodec, counters: Arc, + // The FFI task-context handle is weak. Retain its provider for as long as + // this codec can be called, even if Python drops the exporter object. + _ctx_provider: Arc, +} + +impl fmt::Debug for CountingLogicalExtensionCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CountingLogicalExtensionCodec") + .field("inner", &self.inner) + .field("counters", &self.counters) + .finish_non_exhaustive() + } } impl LogicalExtensionCodec for CountingLogicalExtensionCodec { @@ -72,6 +99,20 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec { schema: SchemaRef, ctx: &TaskContext, ) -> Result> { + if let Some(id) = token_id(buf, TABLE_PROVIDER_TOKEN) { + self.counters + .decode_table_provider + .fetch_add(1, Ordering::SeqCst); + return table_providers() + .lock() + .map_err(|err| DataFusionError::Internal(err.to_string()))? + .remove(&id) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "Unknown datafusion-ffi-example table provider token {id}" + )) + }); + } self.inner .try_decode_table_provider(buf, table_ref, schema, ctx) } @@ -82,6 +123,19 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec { node: Arc, buf: &mut Vec, ) -> Result<()> { + if node.downcast_ref::().is_some() { + self.counters + .encode_table_provider + .fetch_add(1, Ordering::SeqCst); + let id = NEXT_TABLE_PROVIDER_ID.fetch_add(1, Ordering::SeqCst); + table_providers() + .lock() + .map_err(|err| DataFusionError::Internal(err.to_string()))? + .insert(id, node); + buf.extend_from_slice(TABLE_PROVIDER_TOKEN); + buf.extend_from_slice(&id.to_le_bytes()); + return Ok(()); + } self.inner.try_encode_table_provider(table_ref, node, buf) } @@ -105,6 +159,7 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec { #[derive(Clone)] pub(crate) struct MyLogicalExtensionCodec { counters: Arc, + ctx_provider: Arc, } #[pymethods] @@ -113,26 +168,26 @@ impl MyLogicalExtensionCodec { fn new() -> Self { Self { counters: Arc::new(CallCounters::default()), + ctx_provider: Arc::new(SessionContext::new()), } } - /// Number of `try_encode_udf` invocations observed since - /// construction. fn encode_udf_calls(&self) -> usize { self.counters.encode_udf.load(Ordering::SeqCst) } - /// Number of `try_decode_udf` invocations observed. fn decode_udf_calls(&self) -> usize { self.counters.decode_udf.load(Ordering::SeqCst) } - /// Capsule entry point consumed by - /// `datafusion_python_util::ffi_logical_codec_from_pycapsule`. - /// datafusion-python invokes this with no arguments when the user - /// calls `ctx.with_logical_extension_codec(my_codec)`. The codec - /// owns its own bare `SessionContext` as a TaskContextProvider — - /// good enough for tests that only exercise UDF encode/decode. + fn table_provider_encode_calls(&self) -> usize { + self.counters.encode_table_provider.load(Ordering::SeqCst) + } + + fn table_provider_decode_calls(&self) -> usize { + self.counters.decode_table_provider.load(Ordering::SeqCst) + } + fn __datafusion_logical_extension_codec__<'py>( &self, py: Python<'py>, @@ -140,11 +195,11 @@ impl MyLogicalExtensionCodec { let inner: Arc = Arc::new(CountingLogicalExtensionCodec { inner: DefaultLogicalExtensionCodec {}, counters: Arc::clone(&self.counters), + _ctx_provider: Arc::clone(&self.ctx_provider), }); let runtime = get_tokio_runtime().handle().clone(); - let bare_session: Arc = Arc::new(SessionContext::new()); - let ctx_provider = bare_session as Arc; + let ctx_provider: Arc = self.ctx_provider.clone(); let ffi = FFI_LogicalExtensionCodec::new(inner, Some(runtime), &ctx_provider); PyCapsule::new_with_value(py, ffi, cr"datafusion_logical_extension_codec") diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs index 35ef77f6b..d1b9ed63b 100644 --- a/examples/datafusion-ffi-example/src/physical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -15,14 +15,18 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::collections::HashMap; +use std::fmt; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; -use datafusion::common::Result; +use datafusion::common::{DataFusionError, Result}; +use datafusion::datasource::source::DataSourceExec; use datafusion::execution::{TaskContext, TaskContextProvider}; use datafusion::logical_expr::ScalarUDF; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionContext; +use datafusion_ffi::execution_plan::ForeignExecutionPlan; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, @@ -31,20 +35,48 @@ use datafusion_python_util::get_tokio_runtime; use pyo3::prelude::*; use pyo3::types::PyCapsule; +const EXECUTION_PLAN_TOKEN: &[u8] = b"DFPYEXEP"; +static NEXT_EXECUTION_PLAN_ID: AtomicU64 = AtomicU64::new(1); +static EXECUTION_PLANS: OnceLock>>> = OnceLock::new(); + +fn execution_plans() -> &'static Mutex>> { + EXECUTION_PLANS.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn token_id(buf: &[u8]) -> Option { + let id: [u8; 8] = buf.strip_prefix(EXECUTION_PLAN_TOKEN)?.try_into().ok()?; + Some(u64::from_le_bytes(id)) +} + #[derive(Debug, Default)] pub(crate) struct PhysicalCallCounters { pub encode_udf: AtomicUsize, pub decode_udf: AtomicUsize, + pub encode_execution_plan: AtomicUsize, + pub decode_execution_plan: AtomicUsize, } -/// Mirror of [`super::logical_extension_codec::CountingLogicalExtensionCodec`] -/// for the physical layer. Delegates to `DefaultPhysicalExtensionCodec` -/// and bumps counters on UDF encode/decode so tests can prove the -/// session routed through a user-supplied physical codec. -#[derive(Debug)] +/// Physical companion to the logical example codec. +/// +/// Provider-owned memory scan plans use a same-process token registry so the +/// owning cdylib can restore their concrete Rust type after the plan travels +/// through the independent query-planner and datafusion-python libraries. struct CountingPhysicalExtensionCodec { inner: DefaultPhysicalExtensionCodec, counters: Arc, + // The FFI task-context handle is weak. Keep its provider alive with the + // codec rather than relying on the lifetime of the Python exporter. + _ctx_provider: Arc, +} + +impl fmt::Debug for CountingPhysicalExtensionCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CountingPhysicalExtensionCodec") + .field("inner", &self.inner) + .field("counters", &self.counters) + .finish_non_exhaustive() + } } impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { @@ -55,6 +87,20 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { ctx: &TaskContext, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { + if let Some(id) = token_id(buf) { + self.counters + .decode_execution_plan + .fetch_add(1, Ordering::SeqCst); + return execution_plans() + .lock() + .map_err(|err| DataFusionError::Internal(err.to_string()))? + .remove(&id) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "Unknown datafusion-ffi-example execution plan token {id}" + )) + }); + } self.inner.try_decode(buf, inputs, ctx, proto_converter) } @@ -64,6 +110,22 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { buf: &mut Vec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { + // The provider owns DataSourceExec. A ForeignExecutionPlan can wrap a + // host-added execution decorator around that scan; retaining the opaque + // wrapper preserves its original library identity without downcasting it. + if node.is::() || node.is::() { + self.counters + .encode_execution_plan + .fetch_add(1, Ordering::SeqCst); + let id = NEXT_EXECUTION_PLAN_ID.fetch_add(1, Ordering::SeqCst); + execution_plans() + .lock() + .map_err(|err| DataFusionError::Internal(err.to_string()))? + .insert(id, node); + buf.extend_from_slice(EXECUTION_PLAN_TOKEN); + buf.extend_from_slice(&id.to_le_bytes()); + return Ok(()); + } self.inner.try_encode(node, buf, proto_converter) } @@ -87,6 +149,7 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { #[derive(Clone)] pub(crate) struct MyPhysicalExtensionCodec { counters: Arc, + ctx_provider: Arc, } #[pymethods] @@ -95,6 +158,7 @@ impl MyPhysicalExtensionCodec { fn new() -> Self { Self { counters: Arc::new(PhysicalCallCounters::default()), + ctx_provider: Arc::new(SessionContext::new()), } } @@ -106,6 +170,14 @@ impl MyPhysicalExtensionCodec { self.counters.decode_udf.load(Ordering::SeqCst) } + fn execution_plan_encode_calls(&self) -> usize { + self.counters.encode_execution_plan.load(Ordering::SeqCst) + } + + fn execution_plan_decode_calls(&self) -> usize { + self.counters.decode_execution_plan.load(Ordering::SeqCst) + } + fn __datafusion_physical_extension_codec__<'py>( &self, py: Python<'py>, @@ -114,11 +186,11 @@ impl MyPhysicalExtensionCodec { Arc::new(CountingPhysicalExtensionCodec { inner: DefaultPhysicalExtensionCodec {}, counters: Arc::clone(&self.counters), + _ctx_provider: Arc::clone(&self.ctx_provider), }); let runtime = get_tokio_runtime().handle().clone(); - let bare_session: Arc = Arc::new(SessionContext::new()); - let ctx_provider = bare_session as Arc; + let ctx_provider: Arc = self.ctx_provider.clone(); let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), &ctx_provider); PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") diff --git a/examples/datafusion-ffi-example/src/query_planner.rs b/examples/datafusion-ffi-example/src/query_planner.rs deleted file mode 100644 index 4637074dd..000000000 --- a/examples/datafusion-ffi-example/src/query_planner.rs +++ /dev/null @@ -1,92 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; - -use arrow::datatypes::{DataType, Field, Schema}; -use async_trait::async_trait; -use datafusion::execution::TaskContextProvider; -use datafusion::execution::context::{QueryPlanner, SessionContext}; -use datafusion::logical_expr::LogicalPlan; -use datafusion::physical_plan::ExecutionPlan; -use datafusion::physical_plan::empty::EmptyExec; -use datafusion_ffi::query_planner::FFI_QueryPlanner; -use datafusion_python_util::get_tokio_runtime; -use datafusion_session::Session; -use pyo3::prelude::*; -use pyo3::types::PyCapsule; - -#[derive(Debug)] -struct CountingQueryPlanner { - plan_calls: Arc, -} - -#[async_trait] -impl QueryPlanner for CountingQueryPlanner { - async fn create_physical_plan( - &self, - _logical_plan: &LogicalPlan, - _session: &dyn Session, - ) -> datafusion::common::Result> { - self.plan_calls.fetch_add(1, Ordering::SeqCst); - let schema = Arc::new(Schema::new(vec![Field::new( - "value", - DataType::Int64, - false, - )])); - Ok(Arc::new(EmptyExec::new(schema))) - } -} - -/// Python-visible query planner used to test planning across a real FFI boundary. -#[pyclass( - from_py_object, - name = "MyQueryPlanner", - module = "datafusion_ffi_example", - subclass -)] -#[derive(Debug, Default, Clone)] -pub(crate) struct MyQueryPlanner { - plan_calls: Arc, -} - -#[pymethods] -impl MyQueryPlanner { - #[new] - fn new() -> Self { - Self::default() - } - - fn plan_calls(&self) -> usize { - self.plan_calls.load(Ordering::SeqCst) - } - - fn __datafusion_query_planner__<'py>( - &self, - py: Python<'py>, - ) -> PyResult> { - let planner: Arc = Arc::new(CountingQueryPlanner { - plan_calls: Arc::clone(&self.plan_calls), - }); - let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider = Arc::new(SessionContext::new()) as Arc; - let ffi = FFI_QueryPlanner::new(planner, Some(runtime), &ctx_provider, None, None); - - PyCapsule::new_with_value(py, ffi, cr"datafusion_query_planner") - } -} diff --git a/examples/datafusion-ffi-query-planner-example/Cargo.toml b/examples/datafusion-ffi-query-planner-example/Cargo.toml new file mode 100644 index 000000000..4d02c69f1 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/Cargo.toml @@ -0,0 +1,49 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "datafusion-ffi-query-planner-example" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +description.workspace = true +homepage.workspace = true +repository.workspace = true +publish = false + +[dependencies] +datafusion = { workspace = true } +datafusion-catalog = { workspace = true, default-features = false } +datafusion-common = { workspace = true, default-features = false } +datafusion-ffi = { workspace = true } +datafusion-session = { workspace = true } +async-trait = { workspace = true } +datafusion-python-util.workspace = true +pyo3 = { workspace = true, features = [ + "extension-module", + "abi3", + "abi3-py310", +] } +pyo3-log = { workspace = true } + +[build-dependencies] +pyo3-build-config = { workspace = true } + +[lib] +name = "datafusion_ffi_query_planner_example" +crate-type = ["cdylib", "rlib"] diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md new file mode 100644 index 000000000..9756f246c --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -0,0 +1,41 @@ +# DataFusion Python FFI query planner example + +This crate is an independent query-planner Python extension. Together with [`../datafusion-ffi-example`](../datafusion-ffi-example/) it demonstrates a real three-library plan exchange: + +- **A — `datafusion-python`:** owns the session and final execution. +- **B — `datafusion-ffi-example`:** owns a table provider, UDF, and provider codecs. +- **C — this crate:** owns the query planner and its custom configuration. + +Two extension crates are used rather than placing the planner in the provider crate. Loading distinct `cdylib` images gives each library a distinct DataFusion marker and proves that foreign sessions, providers, and plans survive the actual ABI boundary. + +## Running the example + +From the repository root, build and install all three extensions, then run the +integration tests: + +```bash +maturin develop --uv +uv run maturin develop --manifest-path examples/datafusion-ffi-example/Cargo.toml +uv run maturin develop \ + --manifest-path examples/datafusion-ffi-query-planner-example/Cargo.toml +uv run pytest \ + examples/datafusion-ffi-query-planner-example/python/tests/_test*.py +``` + +The integration test follows this setup: + +```python +config = SessionConfig().with_extension(PlannerConfig(max_rows=3)) +ctx = SessionContext(config) +ctx = ctx.with_logical_extension_codec(provider_logical_codec) +ctx = ctx.with_physical_extension_codec(provider_physical_codec) +ctx.register_table("numbers", provider) +ctx.register_udf(provider_udf) +ctx = ctx.with_query_planner(MyQueryPlanner()) +``` + +`PlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit. + +The provider's codec pair is injected into the planner for each planning call and is also used to decode the returned physical plan in `datafusion-python`. The API currently supports one external codec owner rather than a registry of independently composed codecs, so this planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; derived contexts rebind codecs after planner installation, but planner-last order is easier to audit. + +The pinned FFI logical codec cannot encode arbitrary custom `LogicalPlan::Extension` nodes. The example therefore demonstrates table-provider, UDF, and physical-plan interoperability without claiming custom logical extension support. diff --git a/examples/datafusion-ffi-query-planner-example/build.rs b/examples/datafusion-ffi-query-planner-example/build.rs new file mode 100644 index 000000000..4878d8b0e --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/build.rs @@ -0,0 +1,20 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +fn main() { + pyo3_build_config::add_extension_module_link_args(); +} diff --git a/examples/datafusion-ffi-example/python/tests/_test_query_planner.py b/examples/datafusion-ffi-query-planner-example/pyproject.toml similarity index 52% rename from examples/datafusion-ffi-example/python/tests/_test_query_planner.py rename to examples/datafusion-ffi-query-planner-example/pyproject.toml index 91b6e9c4e..9e34b4cd4 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/pyproject.toml @@ -15,25 +15,18 @@ # specific language governing permissions and limitations # under the License. -from __future__ import annotations +[build-system] +requires = ["maturin>=1.6,<2.0"] +build-backend = "maturin" -import pytest -from datafusion import SessionContext -from datafusion_ffi_example import MyQueryPlanner +[project] +name = "datafusion_ffi_query_planner_example" +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", +] +dynamic = ["version"] - -@pytest.mark.parametrize("raw_capsule", [False, True]) -def test_ffi_query_planner_runs_during_planning(raw_capsule: bool): - """A query planner imported from another library creates the physical plan.""" - planner = MyQueryPlanner() - exported_planner = ( - planner.__datafusion_query_planner__() if raw_capsule else planner - ) - ctx = SessionContext().with_query_planner(exported_planner) - - before = planner.plan_calls() - result = ctx.sql("SELECT 1 AS value").collect() - after = planner.plan_calls() - - assert after > before - assert result == [] +[tool.maturin] +features = ["pyo3/extension-module"] diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py new file mode 100644 index 000000000..8279879d1 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -0,0 +1,120 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import gc + +import pytest +from datafusion import SessionConfig, SessionContext, udf +from datafusion_ffi_example import ( + IsNullUDF, + MyLogicalExtensionCodec, + MyPhysicalExtensionCodec, + MyTableProvider, +) +from datafusion_ffi_query_planner_example import MyQueryPlanner, PlannerConfig + + +def configured_context(max_rows: int): + config = SessionConfig().with_extension(PlannerConfig(max_rows=max_rows)) + logical_codec = MyLogicalExtensionCodec() + physical_codec = MyPhysicalExtensionCodec() + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(physical_codec) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + return ctx, logical_codec, physical_codec + + +@pytest.mark.parametrize("raw_capsule", [False, True]) +def test_three_library_query_planner(raw_capsule: bool): + """Host, provider, and planner exchange a real non-empty plan over FFI.""" + ctx, logical_codec, physical_codec = configured_context(max_rows=3) + planner = MyQueryPlanner() + exported_planner = ( + planner.__datafusion_query_planner__() if raw_capsule else planner + ) + ctx = ctx.with_query_planner(exported_planner) + + batches = ctx.sql( + 'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"' + ).collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert batches[0].column(1).to_pylist() == [False, False, False] + assert planner.last_max_rows() == 3 + + ctx.sql("SET ffi_query_planner.max_rows = 2").collect() + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner.last_max_rows() == 2 + + assert planner.plan_calls() >= 2 + assert planner.foreign_session_observed() + assert planner.foreign_provider_observed() + assert planner.foreign_plan_observed() + assert logical_codec.table_provider_encode_calls() > 0 + assert logical_codec.table_provider_decode_calls() > 0 + assert physical_codec.execution_plan_encode_calls() > 0 + assert physical_codec.execution_plan_decode_calls() > 0 + + +def test_installed_codecs_outlive_python_exporters(): + ctx, logical_codec, physical_codec = configured_context(max_rows=2) + del logical_codec, physical_codec + gc.collect() + + ctx = ctx.with_query_planner(MyQueryPlanner()) + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + + +def test_provider_codecs_can_be_installed_after_planner(): + config = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + planner = MyQueryPlanner() + logical_codec = MyLogicalExtensionCodec() + physical_codec = MyPhysicalExtensionCodec() + ctx = SessionContext(config).with_query_planner(planner) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(physical_codec) + ctx.register_table("numbers", MyTableProvider(1, 4, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner.last_max_rows() == 2 + assert logical_codec.table_provider_decode_calls() > 0 + assert physical_codec.execution_plan_decode_calls() > 0 + + +def test_query_planner_requires_provider_codec(): + config = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + ctx = SessionContext(config) + ctx.register_table("numbers", MyTableProvider(1, 3, 1)) + ctx = ctx.with_query_planner(MyQueryPlanner()) + + with pytest.raises(Exception, match=r"LogicalExtensionCodec|TableProvider"): + ctx.sql('SELECT "A" FROM numbers').collect() + + +@pytest.mark.parametrize("max_rows", ["0", "oops"]) +def test_query_planner_rejects_invalid_config(max_rows: str): + ctx, _logical_codec, _physical_codec = configured_context(max_rows=2) + ctx = ctx.with_query_planner(MyQueryPlanner()) + + with pytest.raises(Exception, match=r"max_rows|Invalid value"): + ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect() diff --git a/examples/datafusion-ffi-query-planner-example/src/config.rs b/examples/datafusion-ffi-query-planner-example/src/config.rs new file mode 100644 index 000000000..801cee9a0 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/config.rs @@ -0,0 +1,112 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::any::Any; + +use datafusion_common::config::{ + ConfigEntry, ConfigExtension, ConfigField, ExtensionOptions, Visit, +}; +use datafusion_common::{DataFusionError, config_err}; +use datafusion_ffi::config::extension_options::FFI_ExtensionOptions; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +#[pyclass( + from_py_object, + name = "PlannerConfig", + module = "datafusion_ffi_query_planner_example", + subclass +)] +#[derive(Clone, Debug)] +pub(crate) struct PlannerConfig { + pub max_rows: usize, +} + +#[pymethods] +impl PlannerConfig { + #[new] + #[pyo3(signature = (max_rows=10))] + fn new(max_rows: usize) -> Self { + Self { max_rows } + } + + fn __datafusion_extension_options__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let mut config = FFI_ExtensionOptions::default(); + config + .add_config(self) + .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; + PyCapsule::new_with_value(py, config, cr"datafusion_extension_options") + } +} + +impl Default for PlannerConfig { + fn default() -> Self { + Self { max_rows: 10 } + } +} + +impl ConfigExtension for PlannerConfig { + const PREFIX: &'static str = "ffi_query_planner"; +} + +impl ExtensionOptions for PlannerConfig { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn cloned(&self) -> Box { + Box::new(self.clone()) + } + + fn set(&mut self, key: &str, value: &str) -> datafusion_common::Result<()> { + ConfigField::set(self, key, value) + } + + fn entries(&self) -> Vec { + vec![ConfigEntry { + key: "max_rows".to_owned(), + value: Some(self.max_rows.to_string()), + description: "Maximum rows returned by the example query planner", + }] + } +} + +impl ConfigField for PlannerConfig { + fn visit(&self, visitor: &mut V, _key: &str, _description: &'static str) { + self.max_rows.visit( + visitor, + "max_rows", + "Maximum rows returned by the example query planner", + ); + } + + fn set(&mut self, key: &str, value: &str) -> Result<(), DataFusionError> { + let (key, rem) = key.split_once('.').unwrap_or((key, "")); + match key { + "max_rows" => self.max_rows.set(rem, value), + _ => config_err!("Config value '{key}' not found on PlannerConfig"), + } + } +} diff --git a/examples/datafusion-ffi-query-planner-example/src/lib.rs b/examples/datafusion-ffi-query-planner-example/src/lib.rs new file mode 100644 index 000000000..7635c2992 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/lib.rs @@ -0,0 +1,32 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use pyo3::prelude::*; + +use crate::config::PlannerConfig; +use crate::planner::MyQueryPlanner; + +mod config; +mod planner; + +#[pymodule] +fn datafusion_ffi_query_planner_example(m: &Bound<'_, PyModule>) -> PyResult<()> { + pyo3_log::init(); + m.add_class::()?; + m.add_class::()?; + Ok(()) +} diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs new file mode 100644 index 000000000..65f4ebec3 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -0,0 +1,198 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +use async_trait::async_trait; +use datafusion::execution::TaskContextProvider; +use datafusion::execution::context::SessionContext; +use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::limit::GlobalLimitExec; +use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}; +use datafusion_catalog::default_table_source::source_as_provider; +use datafusion_ffi::config::ExtensionOptionsFFIProvider; +use datafusion_ffi::execution_plan::ForeignExecutionPlan; +use datafusion_ffi::query_planner::FFI_QueryPlanner; +use datafusion_ffi::session::ForeignSession; +use datafusion_ffi::table_provider::ForeignTableProvider; +use datafusion_python_util::get_tokio_runtime; +use datafusion_session::{QueryPlanner, Session}; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +use crate::config::PlannerConfig; + +#[derive(Debug, Default)] +struct PlannerObservations { + plan_calls: AtomicUsize, + last_max_rows: AtomicUsize, + foreign_session: AtomicBool, + foreign_provider: AtomicBool, + foreign_plan: AtomicBool, +} + +fn logical_plan_has_foreign_provider(plan: &LogicalPlan) -> bool { + if let LogicalPlan::TableScan(scan) = plan + && let Ok(provider) = source_as_provider(&scan.source) + && provider.downcast_ref::().is_some() + { + return true; + } + plan.inputs() + .iter() + .any(|input| logical_plan_has_foreign_provider(input)) +} + +fn physical_plan_has_foreign_plan(plan: &Arc) -> bool { + plan.is::() + || plan + .children() + .iter() + .any(|child| physical_plan_has_foreign_plan(child)) +} + +fn planner_config(session: &dyn Session) -> datafusion::common::Result { + let options = session.config_options(); + + // Read the flattened entry first. Some DataFusion revisions add an extra + // `datafusion_ffi` namespace while reconstructing a ForeignSession. Parsing + // it directly also ensures malformed values are reported instead of being + // replaced silently by PlannerConfig::default(). + if let Some(entry) = options + .entries() + .into_iter() + .find(|entry| entry.key.ends_with("ffi_query_planner.max_rows")) + { + let value = entry.value.ok_or_else(|| { + datafusion::common::DataFusionError::Configuration(format!( + "{} must have a value", + entry.key + )) + })?; + let max_rows = value.parse::().map_err(|err| { + datafusion::common::DataFusionError::Configuration(format!( + "Invalid value '{value}' for {}: {err}", + entry.key + )) + })?; + if max_rows == 0 { + return Err(datafusion::common::DataFusionError::Configuration( + "ffi_query_planner.max_rows must be greater than zero".to_owned(), + )); + } + return Ok(PlannerConfig { max_rows }); + } + + Ok(options + .local_or_ffi_extension::() + .unwrap_or_default()) +} + +#[derive(Debug)] +struct DistributedQueryPlanner { + observations: Arc, +} + +#[async_trait] +impl QueryPlanner for DistributedQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> datafusion::common::Result> { + self.observations.plan_calls.fetch_add(1, Ordering::SeqCst); + self.observations + .foreign_session + .store(session.as_any().is::(), Ordering::SeqCst); + self.observations.foreign_provider.store( + logical_plan_has_foreign_provider(logical_plan), + Ordering::SeqCst, + ); + + let config = planner_config(session)?; + self.observations + .last_max_rows + .store(config.max_rows, Ordering::SeqCst); + + let plan = DefaultPhysicalPlanner::default() + .create_physical_plan(logical_plan, session) + .await?; + self.observations + .foreign_plan + .store(physical_plan_has_foreign_plan(&plan), Ordering::SeqCst); + + Ok(Arc::new(GlobalLimitExec::new( + plan, + 0, + Some(config.max_rows), + ))) + } +} + +#[pyclass( + from_py_object, + name = "MyQueryPlanner", + module = "datafusion_ffi_query_planner_example", + subclass +)] +#[derive(Debug, Default, Clone)] +pub(crate) struct MyQueryPlanner { + observations: Arc, +} + +#[pymethods] +impl MyQueryPlanner { + #[new] + fn new() -> Self { + Self::default() + } + + fn plan_calls(&self) -> usize { + self.observations.plan_calls.load(Ordering::SeqCst) + } + + fn last_max_rows(&self) -> usize { + self.observations.last_max_rows.load(Ordering::SeqCst) + } + + fn foreign_session_observed(&self) -> bool { + self.observations.foreign_session.load(Ordering::SeqCst) + } + + fn foreign_provider_observed(&self) -> bool { + self.observations.foreign_provider.load(Ordering::SeqCst) + } + + fn foreign_plan_observed(&self) -> bool { + self.observations.foreign_plan.load(Ordering::SeqCst) + } + + fn __datafusion_query_planner__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let planner: Arc = Arc::new(DistributedQueryPlanner { + observations: Arc::clone(&self.observations), + }); + let runtime = get_tokio_runtime().handle().clone(); + let ctx_provider = Arc::new(SessionContext::new()) as Arc; + let ffi = FFI_QueryPlanner::new(planner, Some(runtime), &ctx_provider, None, None); + PyCapsule::new_with_value(py, ffi, cr"datafusion_query_planner") + } +} diff --git a/python/datafusion/context.py b/python/datafusion/context.py index f92c57b28..c7b73c5da 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1776,7 +1776,8 @@ def with_query_planner( The planner is imported through its ``__datafusion_query_planner__`` PyCapsule. The returned context preserves the existing session state and - its logical and physical extension codec settings. + its logical and physical extension codec settings. Codec changes made on + a derived context are rebound to the planner before planning. Args: planner: Object exposing ``__datafusion_query_planner__`` or a raw From 05ab5da6c6ee895839243c9be15701ef1913bba0 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 09:36:50 -0400 Subject: [PATCH 03/30] Update FFI query planner integration AI Disclosure: This code was written in part by an AI agent.: --- Cargo.lock | 180 +++++++++--------- Cargo.toml | 22 +-- crates/core/src/context.rs | 143 +++++++------- .../Cargo.toml | 1 + .../README.md | 2 +- .../src/planner.rs | 10 +- 6 files changed, 183 insertions(+), 175 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e03c75f15..fdd13713f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -99,9 +99,9 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b952ca5a8046ad741b60f142d6eca4aeebcad615694202bc64c5341f23e32c5b" +checksum = "61d285d16bce7d0be61912f7928342b673067b6b7d7ef6cc179258ba7de1fecf" dependencies = [ "arrow-arith", "arrow-array", @@ -121,9 +121,9 @@ dependencies = [ [[package]] name = "arrow-arith" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64a13b8d3008c4e9063c597a08f46446fe3fd5789277127672d6c0bdbb43b1ff" +checksum = "757ef1836251e88222542a7da2623bc1c9cb9e20afefa6db2c41e79991cd91d4" dependencies = [ "arrow-array", "arrow-buffer", @@ -135,9 +135,9 @@ dependencies = [ [[package]] name = "arrow-array" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9486151b2f0785bafc6fa04fc5c99fcb4495455662e58787ea32eaaed33c4192" +checksum = "bc9a4a4b2b5ecd0e04df03471661cb61f28bed3c7fd50994715129b01b2edb97" dependencies = [ "ahash", "arrow-buffer", @@ -147,6 +147,7 @@ dependencies = [ "chrono-tz", "half", "hashbrown 0.17.1", + "libc", "num-complex", "num-integer", "num-traits", @@ -154,9 +155,9 @@ dependencies = [ [[package]] name = "arrow-avro" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e4f9b23a0d7b613acb59fa20bdbe0f80ffdae6411498378340b3915e45f5b84" +checksum = "9fb45cd6bd2b25c0965793b83200eaca82214273a8030fbbc2d783e4c7c65a61" dependencies = [ "arrow-array", "arrow-buffer", @@ -178,21 +179,21 @@ dependencies = [ [[package]] name = "arrow-buffer" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4776577a87794bfdf0b4e90e2ea12454fa7738ea2823c4be5b9d1851da7b434" +checksum = "c12b576ef18c1deb80925a248b25ad84f419198d791b8e293fc6aaa60441fe90" dependencies = [ "bytes", "half", - "num-bigint", + "num-bigint 0.5.1", "num-traits", ] [[package]] name = "arrow-cast" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9ad451ce4f98710828a455b96991b8f031deb2e67f5fcad6773f017e4a69c3a" +checksum = "68338a9096a5dc9bc11927c58c43a8526d96bf6abd2012ef6c0c9f505991cc79" dependencies = [ "arrow-array", "arrow-buffer", @@ -201,7 +202,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64 0.22.1", + "base64 0.23.0", "chrono", "comfy-table", "half", @@ -212,9 +213,9 @@ dependencies = [ [[package]] name = "arrow-csv" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8aa7bf96d6141a7bcca2eed57c7c9767d2a2175281857b8a7b68308992864784" +checksum = "25011b52b346407d497ef0030e12b45e4f2d0cc279efc09c4f3d09106db30e36" dependencies = [ "arrow-array", "arrow-cast", @@ -227,9 +228,9 @@ dependencies = [ [[package]] name = "arrow-data" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b38fe43e2e8704360f1464e6e8cc4fc381ef02cc4fb0192afa8df1aaa0115c66" +checksum = "723fe4aeed7604e00b9883a465af4ff0a0e6c44c03e41a68c3d1cbc403e0e44d" dependencies = [ "arrow-buffer", "arrow-schema", @@ -240,9 +241,9 @@ dependencies = [ [[package]] name = "arrow-ipc" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29dac499fcbc6ba74ee0324057821d381929a48526a3966bd9dffb44aa06d98c" +checksum = "149437b14371f5b9ec60f5ddc751483ae99d7a7072653c0075e5e469156eea7b" dependencies = [ "arrow-array", "arrow-buffer", @@ -256,9 +257,9 @@ dependencies = [ [[package]] name = "arrow-json" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fe05e916ddc50f4c7a363cd69c0ef5894fcee063517e9a0b8582f0c56746af6" +checksum = "f18b9123ccfec418a663f821c9a034af339711678c11ffe00d3ec07da5ff9f7e" dependencies = [ "arrow-array", "arrow-buffer", @@ -281,9 +282,9 @@ dependencies = [ [[package]] name = "arrow-ord" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e13dbdc2a9c053c10c7baa6e30faee04a180aa7ce88e471835850ce37abd20b" +checksum = "e6c08dff0686cf23ca4f562803f191ccbeb726dbae6309cd4b4aaf65e0f2c979" dependencies = [ "arrow-array", "arrow-buffer", @@ -294,9 +295,9 @@ dependencies = [ [[package]] name = "arrow-pyarrow" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf8d967bdece4fa5a0199706730175b3df3448b87e350250a86eb6c22639e445" +checksum = "c196ecc25b3a8dcbc1d842f2619cee653dcfa2fb8b56a291bc0481c3cf5c3821" dependencies = [ "arrow-array", "arrow-data", @@ -306,9 +307,9 @@ dependencies = [ [[package]] name = "arrow-row" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4d5a1f8c733d15260b305683472ee8ad89c62cbd706703ca873b90d051b41592" +checksum = "bbec439386df71ad570e6758a946111322b9e9dc8db83b5527321f0b4c9119c2" dependencies = [ "arrow-array", "arrow-buffer", @@ -319,9 +320,9 @@ dependencies = [ [[package]] name = "arrow-schema" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9e4969dc350d571766247143ab36a5187d095d3d3690970408bc630d47c69e5" +checksum = "e6fed2ca0d1eade57e811cbe73b98ad50cc08a1183e13b2d2aa43a7df593f40e" dependencies = [ "bitflags", "serde_core", @@ -330,9 +331,9 @@ dependencies = [ [[package]] name = "arrow-select" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "402770dba90865359d98d1ef92ef16e23d75c0cca9c2c880c8a05468b7743bf9" +checksum = "466b19cf75130b891dc1b23a84b343c714c62c64c9c62e365c76aa0ff90a53fb" dependencies = [ "ahash", "arrow-array", @@ -344,9 +345,9 @@ dependencies = [ [[package]] name = "arrow-string" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2b0afbb8b9016700938291123df30838b89decc3213dba00852021988b170d3" +checksum = "c838a25bb3691e919e0f617616ac51a4ff8517a952e29ca133cf0c22b2ce65b1" dependencies = [ "arrow-array", "arrow-buffer", @@ -440,7 +441,7 @@ checksum = "4d6867f1565b3aad85681f1015055b087fcfd840d6aeee6eee7f2da317603695" dependencies = [ "autocfg", "libm", - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", ] @@ -791,7 +792,7 @@ dependencies = [ [[package]] name = "datafusion" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-schema", @@ -844,7 +845,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "async-trait", @@ -868,7 +869,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "async-trait", @@ -891,7 +892,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-ipc", @@ -917,7 +918,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "futures", "log", @@ -927,7 +928,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "async-compression", @@ -963,7 +964,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-ipc", @@ -986,7 +987,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-avro", @@ -1004,7 +1005,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "async-trait", @@ -1016,6 +1017,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", @@ -1026,7 +1028,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "async-trait", @@ -1038,6 +1040,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", @@ -1048,7 +1051,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-schema", @@ -1065,6 +1068,7 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-pruning", "datafusion-session", "futures", @@ -1079,12 +1083,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" [[package]] name = "datafusion-execution" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-buffer", @@ -1109,7 +1113,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-schema", @@ -1131,7 +1135,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1142,7 +1146,7 @@ dependencies = [ [[package]] name = "datafusion-ffi" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-schema", @@ -1202,6 +1206,7 @@ dependencies = [ "datafusion-catalog", "datafusion-common", "datafusion-ffi", + "datafusion-proto", "datafusion-python-util", "datafusion-session", "pyo3", @@ -1212,7 +1217,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-buffer", @@ -1243,7 +1248,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1263,7 +1268,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1274,7 +1279,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-ord", @@ -1298,7 +1303,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "async-trait", @@ -1313,7 +1318,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1329,7 +1334,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1338,7 +1343,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "datafusion-doc", "quote", @@ -1348,7 +1353,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "chrono", @@ -1367,7 +1372,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1389,7 +1394,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1403,7 +1408,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "chrono", @@ -1420,7 +1425,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1439,7 +1444,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-data", @@ -1475,7 +1480,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "chrono", @@ -1502,7 +1507,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1512,7 +1517,7 @@ dependencies = [ [[package]] name = "datafusion-proto-models" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "datafusion-proto-common", "prost", @@ -1521,7 +1526,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1582,7 +1587,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow-schema", "async-trait", @@ -1596,7 +1601,7 @@ dependencies = [ [[package]] name = "datafusion-spark" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "bigdecimal", @@ -1625,7 +1630,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "bigdecimal", @@ -1644,7 +1649,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "54.1.0" -source = "git+https://github.com/timsaucer/datafusion?rev=b18d41d81089ab95d824b658cd5232f9c1a695ed#b18d41d81089ab95d824b658cd5232f9c1a695ed" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "async-recursion", "async-trait", @@ -2463,9 +2468,9 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] name = "lz4_flex" -version = "0.13.1" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226" dependencies = [ "twox-hash", ] @@ -2542,6 +2547,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -2656,9 +2671,9 @@ dependencies = [ [[package]] name = "parquet" -version = "59.1.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5302d4da74d6596a1f11f9928767995b53bca657cbeea1e4e8c5074f8a1157dd" +checksum = "7065842956a20c2a536924ce8e4d9955f7422451511b9eb7500d7bfe5077e59c" dependencies = [ "ahash", "arrow-array", @@ -2667,7 +2682,7 @@ dependencies = [ "arrow-ipc", "arrow-schema", "arrow-select", - "base64 0.22.1", + "base64 0.23.0", "brotli", "bytes", "chrono", @@ -2676,11 +2691,10 @@ dependencies = [ "half", "hashbrown 0.17.1", "lz4_flex", - "num-bigint", + "num-bigint 0.5.1", "num-integer", "num-traits", "object_store", - "paste", "seq-macro", "simdutf8", "snap", @@ -2689,12 +2703,6 @@ dependencies = [ "zstd", ] -[[package]] -name = "paste" -version = "1.0.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" - [[package]] name = "pbjson" version = "0.8.0" diff --git a/Cargo.toml b/Cargo.toml index 0f5aa8e27..809ef7a98 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -78,14 +78,14 @@ codegen-units = 2 # We cannot publish to crates.io with any patches in the below section. Developers # must remove any entries in this section before creating a release candidate. [patch.crates-io] -datafusion = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } -datafusion-substrait = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } -datafusion-proto = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } -datafusion-ffi = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } -datafusion-catalog = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } -datafusion-common = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } -datafusion-functions-aggregate = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } -datafusion-functions-window = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } -datafusion-spark = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } -datafusion-expr = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } -datafusion-session = { git = "https://github.com/timsaucer/datafusion", rev = "b18d41d81089ab95d824b658cd5232f9c1a695ed" } +datafusion = { git = "https://github.com/apache/datafusion", rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } +datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } +datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } +datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } +datafusion-catalog = { git = "https://github.com/apache/datafusion", rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } +datafusion-common = { git = "https://github.com/apache/datafusion", rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } +datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion", rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } +datafusion-functions-window = { git = "https://github.com/apache/datafusion", rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } +datafusion-spark = { git = "https://github.com/apache/datafusion", rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } +datafusion-expr = { git = "https://github.com/apache/datafusion", rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } +datafusion-session = { git = "https://github.com/apache/datafusion", rev = "e08aed1e5de41dcf81d529140dae07723b942a5e" } diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 9853fe9b9..a1dc0169f 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -227,73 +227,39 @@ impl PySessionConfig { } } -#[derive(Debug)] -struct PlanningTaskContextProvider(Arc); - -impl TaskContextProvider for PlanningTaskContextProvider { - fn task_ctx(&self) -> Arc { - Arc::clone(&self.0) - } -} - +/// Adapts an FFI planner to the Tokio runtime owned by datafusion-python. +/// +/// Upstream's `ForeignQueryPlanner` cannot recover the runtime handle from the +/// `QueryPlanner` trait, so embedders that own the runtime must call +/// `create_physical_plan_with_session_runtime` directly. #[derive(Debug, Clone)] -struct PythonQueryPlanner { +struct RuntimeAwareQueryPlanner { planner: FFI_QueryPlanner, - logical_codec: Arc, - physical_codec: Arc, } -impl PythonQueryPlanner { - fn with_codecs( +impl RuntimeAwareQueryPlanner { + fn with_ffi_codecs( &self, - logical_codec: Arc, - physical_codec: Arc, + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, ) -> Self { - Self { - planner: self.planner.clone(), - logical_codec, - physical_codec, - } + let planner: Arc = (&self.planner).into(); + let planner = FFI_QueryPlanner::new_with_ffi_codecs(planner, logical_codec, physical_codec); + Self { planner } } } #[async_trait] -impl QueryPlanner for PythonQueryPlanner { +impl QueryPlanner for RuntimeAwareQueryPlanner { async fn create_physical_plan( &self, logical_plan: &LogicalPlan, session: &dyn Session, ) -> datafusion::common::Result> { let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider = Arc::new(PlanningTaskContextProvider(session.task_ctx())); - let (logical_codec, physical_codec) = { - let dyn_ctx_provider: Arc = ctx_provider.clone(); - let logical_codec: Arc = - Arc::clone(&self.logical_codec) as Arc; - let physical_codec: Arc = - Arc::clone(&self.physical_codec) as Arc; - ( - FFI_LogicalExtensionCodec::new( - logical_codec, - Some(runtime.clone()), - &dyn_ctx_provider, - ), - FFI_PhysicalExtensionCodec::new( - physical_codec, - Some(runtime.clone()), - &dyn_ctx_provider, - ), - ) - }; - - let mut planner = self.planner.clone(); - planner.logical_codec = logical_codec; - planner.physical_codec = physical_codec; - let result = planner + self.planner .create_physical_plan_with_session_runtime(logical_plan, session, Some(runtime)) - .await; - drop(ctx_provider); - result + .await } } @@ -1289,15 +1255,22 @@ impl PySessionContext { pub fn with_query_planner(&self, planner: Bound<'_, PyAny>) -> PyDataFusionResult { let planner = ffi_query_planner_from_pycapsule(&planner)?; - let planner = Arc::new(PythonQueryPlanner { + + // Build the codecs against the derived context, then update that same + // context in place. FFI codecs keep a weak task-context provider, so + // rebuilding the context after creating them would leave a stale link. + let ctx = Arc::new(SessionContext::new_with_state(self.ctx.state())); + let planner: Arc = (&planner).into(); + let planner = FFI_QueryPlanner::new_with_ffi_codecs( planner, - logical_codec: Arc::clone(&self.logical_codec), - physical_codec: Arc::clone(&self.physical_codec), - }); - let state = SessionStateBuilder::new_from_existing(self.ctx.state()) + Self::ffi_logical_codec_for(&ctx, &self.logical_codec), + Self::ffi_physical_codec_for(&ctx, &self.physical_codec), + ); + let planner = Arc::new(RuntimeAwareQueryPlanner { planner }); + let state = SessionStateBuilder::new_from_existing(ctx.state()) .with_query_planner(planner) .build(); - let ctx = Arc::new(SessionContext::new_with_state(state)); + *ctx.state_ref().write() = state; Ok(Self { ctx, @@ -1564,15 +1537,24 @@ impl PySessionContext { let state = self.ctx.state(); let query_planner = state.query_planner(); let planner_any: &dyn std::any::Any = query_planner.as_ref(); - let Some(planner) = planner_any.downcast_ref::() else { + let Some(planner) = planner_any + .downcast_ref::() + .cloned() + else { return Arc::clone(&self.ctx); }; - let planner = Arc::new(planner.with_codecs(logical_codec, physical_codec)); - let state = SessionStateBuilder::new_from_existing(self.ctx.state()) + // Preserve the context identity captured by the replacement codecs. + let ctx = Arc::new(SessionContext::new_with_state(state)); + let planner = Arc::new(planner.with_ffi_codecs( + Self::ffi_logical_codec_for(&ctx, &logical_codec), + Self::ffi_physical_codec_for(&ctx, &physical_codec), + )); + let state = SessionStateBuilder::new_from_existing(ctx.state()) .with_query_planner(planner) .build(); - Arc::new(SessionContext::new_with_state(state)) + *ctx.state_ref().write() = state; + ctx } async fn _table(&self, name: &str) -> datafusion::common::Result { @@ -1636,29 +1618,38 @@ impl PySessionContext { /// Used at every site that exports the codec across an FFI boundary /// (capsule getters, Rust wrappers for Python-defined providers, etc.). pub(crate) fn ffi_logical_codec(&self) -> Arc { - let inner: Arc = - Arc::clone(&self.logical_codec) as Arc; + Arc::new(Self::ffi_logical_codec_for(&self.ctx, &self.logical_codec)) + } + + fn ffi_logical_codec_for( + ctx: &Arc, + codec: &Arc, + ) -> FFI_LogicalExtensionCodec { + let codec: Arc = + Arc::clone(codec) as Arc; let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider = Arc::clone(&self.ctx) as Arc; - Arc::new(FFI_LogicalExtensionCodec::new( - inner, - Some(runtime), - &ctx_provider, - )) + let ctx_provider = Arc::clone(ctx) as Arc; + FFI_LogicalExtensionCodec::new(codec, Some(runtime), &ctx_provider) } /// Build an FFI-wrapped clone of the session's physical codec on demand. pub(crate) fn ffi_physical_codec(&self) -> Arc { - let inner: Arc = - Arc::clone(&self.physical_codec) as Arc; - let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider = Arc::clone(&self.ctx) as Arc; - Arc::new(FFI_PhysicalExtensionCodec::new( - inner, - Some(runtime), - &ctx_provider, + Arc::new(Self::ffi_physical_codec_for( + &self.ctx, + &self.physical_codec, )) } + + fn ffi_physical_codec_for( + ctx: &Arc, + codec: &Arc, + ) -> FFI_PhysicalExtensionCodec { + let codec: Arc = + Arc::clone(codec) as Arc; + let runtime = get_tokio_runtime().handle().clone(); + let ctx_provider = Arc::clone(ctx) as Arc; + FFI_PhysicalExtensionCodec::new(codec, Some(runtime), &ctx_provider) + } } pub fn parse_file_compression_type( diff --git a/examples/datafusion-ffi-query-planner-example/Cargo.toml b/examples/datafusion-ffi-query-planner-example/Cargo.toml index 4d02c69f1..263f034b8 100644 --- a/examples/datafusion-ffi-query-planner-example/Cargo.toml +++ b/examples/datafusion-ffi-query-planner-example/Cargo.toml @@ -31,6 +31,7 @@ datafusion = { workspace = true } datafusion-catalog = { workspace = true, default-features = false } datafusion-common = { workspace = true, default-features = false } datafusion-ffi = { workspace = true } +datafusion-proto = { workspace = true } datafusion-session = { workspace = true } async-trait = { workspace = true } datafusion-python-util.workspace = true diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 9756f246c..02bfa75fe 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -36,6 +36,6 @@ ctx = ctx.with_query_planner(MyQueryPlanner()) `PlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit. -The provider's codec pair is injected into the planner for each planning call and is also used to decode the returned physical plan in `datafusion-python`. The API currently supports one external codec owner rather than a registry of independently composed codecs, so this planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; derived contexts rebind codecs after planner installation, but planner-last order is easier to audit. +The provider's codec pair is attached to the planner when the derived context is created and is also used to decode the returned physical plan in `datafusion-python`. The API currently supports one external codec owner rather than a registry of independently composed codecs, so this planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; derived contexts rebind codecs after planner installation, but planner-last order is easier to audit. The pinned FFI logical codec cannot encode arbitrary custom `LogicalPlan::Extension` nodes. The example therefore demonstrates table-provider, UDF, and physical-plan interoperability without claiming custom logical extension support. diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index 65f4ebec3..cb767ffa5 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -31,6 +31,8 @@ use datafusion_ffi::execution_plan::ForeignExecutionPlan; use datafusion_ffi::query_planner::FFI_QueryPlanner; use datafusion_ffi::session::ForeignSession; use datafusion_ffi::table_provider::ForeignTableProvider; +use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; +use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; use datafusion_python_util::get_tokio_runtime; use datafusion_session::{QueryPlanner, Session}; use pyo3::prelude::*; @@ -192,7 +194,13 @@ impl MyQueryPlanner { }); let runtime = get_tokio_runtime().handle().clone(); let ctx_provider = Arc::new(SessionContext::new()) as Arc; - let ffi = FFI_QueryPlanner::new(planner, Some(runtime), &ctx_provider, None, None); + let ffi = FFI_QueryPlanner::new( + planner, + Some(runtime), + &ctx_provider, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::new(DefaultPhysicalExtensionCodec {}), + ); PyCapsule::new_with_value(py, ffi, cr"datafusion_query_planner") } } From 27147f3b42c99338ee3a5966baa819ecd429e146 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 11:05:35 -0400 Subject: [PATCH 04/30] add rat --- examples/datafusion-ffi-example/README.md | 19 +++++++++++++++++++ .../README.md | 19 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md index 4a35e75da..c897c067a 100644 --- a/examples/datafusion-ffi-example/README.md +++ b/examples/datafusion-ffi-example/README.md @@ -1,3 +1,22 @@ + + # DataFusion Python FFI provider example This crate is the **provider library** in the three-library query-planning example. It exports table providers, functions, and the logical and physical codecs needed to serialize objects owned by this library. The companion planner is in [`../datafusion-ffi-query-planner-example`](../datafusion-ffi-query-planner-example/). diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 02bfa75fe..53e2bddc5 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -1,3 +1,22 @@ + + # DataFusion Python FFI query planner example This crate is an independent query-planner Python extension. Together with [`../datafusion-ffi-example`](../datafusion-ffi-example/) it demonstrates a real three-library plan exchange: From 7982ee5b2fa57f658631f7d1d62acba2c6c1b7a8 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sat, 8 Aug 2026 12:07:21 -0400 Subject: [PATCH 05/30] fix: install FFI test wheels from nested artifact paths The FFI test wheel artifact now bundles two projects, so upload-artifact preserves a `/dist/` prefix instead of placing the wheels at the artifact root. The install step globbed `wheels/*.whl`, which no longer matched them, so the FFI wheels were silently skipped and the FFI unit tests failed with `ModuleNotFoundError: No module named 'datafusion_ffi_example'`. Install the recursive `find` results instead of re-globbing. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/test.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 21ecd56f7..047b35039 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -93,11 +93,15 @@ jobs: uv venv --python "${{ steps.setup-python.outputs.python-path }}" VENV_PY="$PWD/.venv/bin/python" uv sync --python "$VENV_PY" --dev --no-install-package datafusion + # Search recursively: the FFI artifact bundles more than one + # project, so upload-artifact keeps a `/dist/` prefix + # and the wheels are not all at the top of wheels/. WHEELS=$(find wheels/ -name "*.whl") if [ -n "$WHEELS" ]; then echo "Installing wheels:" echo "$WHEELS" - uv pip install --python "$VENV_PY" wheels/*.whl + # shellcheck disable=SC2086 # intentional split on newlines + uv pip install --python "$VENV_PY" $WHEELS else echo "ERROR: No wheels found!" exit 1 From 81f68ed313251c7caafb47b637ce133465b4e572 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 25 Aug 2026 08:54:55 -0400 Subject: [PATCH 06/30] refactor: address review of FFI query planner support Collapse the two duplicated planner-install blocks into a single `ctx_with_rebound_planner`. A derived context shares the existing `SessionContext` when there is no foreign planner to rebind, and forks only when one is installed, since the FFI codecs capture the context they are built against. Document what that fork shares. Catalogs, tables, and the runtime environment stay shared; registered functions, configuration, and the optimizer rule lists are snapshotted. The caveat lands on all four derivation methods and on a new contributor-guide subsection, with tests covering both halves. Explain why `RuntimeAwareQueryPlanner` exists at all. Upstream's `ForeignQueryPlanner` is the consumer-side adapter that lets an `FFI_QueryPlanner` satisfy the `QueryPlanner` trait, which is what makes a planner from another shared library installable in a `SessionState`. Its trait method receives only a `&LogicalPlan` and a `&dyn Session`, so it has nowhere to obtain a runtime handle and passes `None`. Throughout datafusion-ffi each library attaches its own runtime to the objects it exports, so a producer-side wrapper can enter that runtime before running its own library's code. A provider owned by another library keeps its owner's runtime even when it travels through our catalog, because `FFI_TableProvider::new_with_ffi_codec` unwraps a `ForeignTableProvider` back to the original handle and discards the runtime passed alongside it. `session_runtime` is that same rule applied to the session: `FFI_SessionRef` is our object and every callback on it runs our code. It matters for what those callbacks hand back. A plan produced by our own planner returns as `FFI_ExecutionPlan::new(plan, runtime)`, and `execute` enters that runtime before calling into the plan; the same holds for our physical optimizer rules and for tables we own rather than re-export. The delegation case this type exists for is exactly that shape. A foreign planner falling back to our planner through `__datafusion_query_planner__` receives a plan whose execution needs our runtime, and datafusion-python owns that runtime as a process global while the Python thread calling in carries no ambient one. The same reasoning is why `__datafusion_query_planner__` re-exports through the adapter rather than unwrapping to the inner handle. A consumer reaching us through `ForeignQueryPlanner` calls with `None`, so the adapter is what restores our handle on the way back out. Unwrapping would save a planning-time round trip and silently drop it. In the planner example, match the two real spellings of the row-limit config key exactly instead of by suffix, and validate after both lookup paths so the fallback cannot accept `max_rows = 0`. The key appears twice because rebuilding a `ConfigOptions` across the FFI boundary parks every foreign extension inside a single `FFI_ExtensionOptions`, itself namespaced under `datafusion_ffi`. Also declare `requires-python = ">=3.10"` on the provider example to match the `abi3-py310` feature it builds against, and link both example READMEs to the contributor guide rather than restating its caveats. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 137 +++++++++++------- docs/source/contributor-guide/ffi.md | 42 ++++++ examples/datafusion-ffi-example/README.md | 8 +- .../datafusion-ffi-example/pyproject.toml | 3 +- .../README.md | 4 +- .../_test_three_library_query_planner.py | 26 ++++ .../src/planner.rs | 66 +++++---- followup-issue.md | 40 +++++ pr-body-addition.md | 10 ++ python/datafusion/context.py | 46 +++++- python/tests/test_context.py | 52 +++++++ 11 files changed, 341 insertions(+), 93 deletions(-) create mode 100644 followup-issue.md create mode 100644 pr-body-addition.md diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index a1dc0169f..59e8b9707 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -229,26 +229,39 @@ impl PySessionConfig { /// Adapts an FFI planner to the Tokio runtime owned by datafusion-python. /// -/// Upstream's `ForeignQueryPlanner` cannot recover the runtime handle from the -/// `QueryPlanner` trait, so embedders that own the runtime must call -/// `create_physical_plan_with_session_runtime` directly. -#[derive(Debug, Clone)] +/// Upstream's `ForeignQueryPlanner` is the consumer-side adapter that lets an +/// `FFI_QueryPlanner` satisfy the `QueryPlanner` trait, which is what makes a +/// planner from another shared library installable in a `SessionState`. Its +/// trait method receives only a `&LogicalPlan` and a `&dyn Session`, so it has +/// nowhere to obtain a runtime handle and calls +/// `create_physical_plan_with_session_runtime` with `None`. +/// +/// Throughout datafusion-ffi each library attaches its *own* runtime to the +/// objects it exports, so a producer-side wrapper can `Handle::enter` before +/// running that library's code. A provider owned by another library keeps its +/// owner's runtime even when it travels through our catalog, because +/// `FFI_TableProvider::new_with_ffi_codec` unwraps a `ForeignTableProvider` +/// back to the original handle and discards the runtime passed alongside it. +/// `session_runtime` is that same rule applied to the session: `FFI_SessionRef` +/// is our object, and every callback on it runs our code. +/// +/// It matters for what those callbacks hand back. A plan produced by our own +/// planner returns as `FFI_ExecutionPlan::new(plan, runtime)`, and `execute` +/// enters that runtime before calling into the plan. The same holds for our +/// physical optimizer rules and for tables we own rather than re-export. The +/// delegation case this type exists for is exactly that shape: a foreign +/// planner that falls back to our planner through +/// `__datafusion_query_planner__` receives a plan whose execution needs our +/// runtime, and datafusion-python owns that runtime as a process global while +/// the Python thread calling in carries no ambient one. +/// +/// This adapter is where the handle gets attached. It wraps the foreign handle +/// and calls `create_physical_plan_with_session_runtime` with `Some(handle)`. +#[derive(Debug)] struct RuntimeAwareQueryPlanner { planner: FFI_QueryPlanner, } -impl RuntimeAwareQueryPlanner { - fn with_ffi_codecs( - &self, - logical_codec: FFI_LogicalExtensionCodec, - physical_codec: FFI_PhysicalExtensionCodec, - ) -> Self { - let planner: Arc = (&self.planner).into(); - let planner = FFI_QueryPlanner::new_with_ffi_codecs(planner, logical_codec, physical_codec); - Self { planner } - } -} - #[async_trait] impl QueryPlanner for RuntimeAwareQueryPlanner { async fn create_physical_plan( @@ -1255,22 +1268,8 @@ impl PySessionContext { pub fn with_query_planner(&self, planner: Bound<'_, PyAny>) -> PyDataFusionResult { let planner = ffi_query_planner_from_pycapsule(&planner)?; - - // Build the codecs against the derived context, then update that same - // context in place. FFI codecs keep a weak task-context provider, so - // rebuilding the context after creating them would leave a stale link. - let ctx = Arc::new(SessionContext::new_with_state(self.ctx.state())); - let planner: Arc = (&planner).into(); - let planner = FFI_QueryPlanner::new_with_ffi_codecs( - planner, - Self::ffi_logical_codec_for(&ctx, &self.logical_codec), - Self::ffi_physical_codec_for(&ctx, &self.physical_codec), - ); - let planner = Arc::new(RuntimeAwareQueryPlanner { planner }); - let state = SessionStateBuilder::new_from_existing(ctx.state()) - .with_query_planner(planner) - .build(); - *ctx.state_ref().write() = state; + let ctx = + self.ctx_with_rebound_planner(&self.logical_codec, &self.physical_codec, Some(planner)); Ok(Self { ctx, @@ -1457,6 +1456,14 @@ impl PySessionContext { &self, py: Python<'py>, ) -> PyResult> { + // When the installed planner is already foreign, this exports a handle + // that wraps `RuntimeAwareQueryPlanner`, which in turn wraps the + // original foreign handle. That extra layer looks redundant but is not: + // a consumer reaching us through `ForeignQueryPlanner` always calls + // `create_physical_plan_with_session_runtime` with `None`, so the + // adapter is what puts our Tokio handle back on the session before the + // call continues outward. Unwrapping to the inner handle here would + // save one planning-time round trip and silently drop that runtime. let planner = Arc::clone(self.ctx.state().query_planner()); let ffi = FFI_QueryPlanner::new_with_ffi_codecs( planner, @@ -1475,8 +1482,7 @@ impl PySessionContext { let logical_codec = Arc::new(PythonLogicalCodec::new(inner)); let physical_codec = Arc::clone(&self.physical_codec); - let ctx = self - .ctx_with_query_planner_codecs(Arc::clone(&logical_codec), Arc::clone(&physical_codec)); + let ctx = self.ctx_with_rebound_planner(&logical_codec, &physical_codec, None); Ok(Self { ctx, logical_codec, @@ -1500,8 +1506,7 @@ impl PySessionContext { let physical_codec = Arc::new(PythonPhysicalCodec::new(inner)); let logical_codec = Arc::clone(&self.logical_codec); - let ctx = self - .ctx_with_query_planner_codecs(Arc::clone(&logical_codec), Arc::clone(&physical_codec)); + let ctx = self.ctx_with_rebound_planner(&logical_codec, &physical_codec, None); Ok(Self { ctx, logical_codec, @@ -1518,8 +1523,7 @@ impl PySessionContext { PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner())) .with_python_udf_inlining(enabled), ); - let ctx = self - .ctx_with_query_planner_codecs(Arc::clone(&logical_codec), Arc::clone(&physical_codec)); + let ctx = self.ctx_with_rebound_planner(&logical_codec, &physical_codec, None); Self { ctx, logical_codec, @@ -1529,31 +1533,60 @@ impl PySessionContext { } impl PySessionContext { - fn ctx_with_query_planner_codecs( + /// Return the [`SessionContext`] a derived `PySessionContext` should use, + /// binding a foreign query planner to `logical_codec` and `physical_codec`. + /// + /// Pass `Some(planner)` to install one, or `None` to rebind whichever + /// planner the session already holds. + /// + /// With no foreign planner in play there is nothing to rebind, so the + /// existing context is shared and swapping codecs alone does not change + /// what a derived context observes. A foreign planner does have to be + /// rebound, and the FFI codecs capture the context they are built against, + /// so that case forks the session: the codecs are built from the fork and + /// then the fork's state is overwritten in place, because rebuilding the + /// context afterwards would leave the codecs pointing at a session that is + /// no longer used for planning. + /// + /// A fork is not a deep copy. `SessionState` keeps its catalog list behind + /// an `Arc`, so catalogs and tables stay shared with the original session, + /// while registered functions and the configuration are snapshotted at the + /// time of the call. The session id is deliberately carried over. + fn ctx_with_rebound_planner( &self, - logical_codec: Arc, - physical_codec: Arc, + logical_codec: &Arc, + physical_codec: &Arc, + planner: Option, ) -> Arc { let state = self.ctx.state(); - let query_planner = state.query_planner(); - let planner_any: &dyn std::any::Any = query_planner.as_ref(); - let Some(planner) = planner_any - .downcast_ref::() - .cloned() - else { + + // When the caller is only replacing codecs, recover the foreign planner + // already installed so it can be rebound below. + let planner = planner.or_else(|| { + let installed: &dyn std::any::Any = state.query_planner().as_ref(); + installed + .downcast_ref::() + .map(|planner| planner.planner.clone()) + }); + + let Some(planner) = planner else { return Arc::clone(&self.ctx); }; - // Preserve the context identity captured by the replacement codecs. let ctx = Arc::new(SessionContext::new_with_state(state)); - let planner = Arc::new(planner.with_ffi_codecs( - Self::ffi_logical_codec_for(&ctx, &logical_codec), - Self::ffi_physical_codec_for(&ctx, &physical_codec), - )); + let inner: Arc = (&planner).into(); + let planner = Arc::new(RuntimeAwareQueryPlanner { + planner: FFI_QueryPlanner::new_with_ffi_codecs( + inner, + Self::ffi_logical_codec_for(&ctx, logical_codec), + Self::ffi_physical_codec_for(&ctx, physical_codec), + ), + }); let state = SessionStateBuilder::new_from_existing(ctx.state()) .with_query_planner(planner) .build(); *ctx.state_ref().write() = state; + ctx } diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index d13eda4d8..ded7b4295 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -257,6 +257,48 @@ The current FFI logical codec supports providers and UDFs but not arbitrary cust `LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and local build commands. +### What a derived context shares + +`with_query_planner`, `with_logical_extension_codec`, `with_physical_extension_codec`, +and `with_python_udf_inlining` all return a new `SessionContext` rather than mutating +the receiver. How much the two contexts then share depends on whether a foreign query +planner is involved. + +Without one, the derived context wraps the *same* underlying session, so a registration +on either side is visible to both. + +`with_query_planner` is different, and so is any codec change made on a session that +already has a foreign planner installed. A foreign planner holds the FFI codecs it was +built with, so changing the codecs means rebuilding the planner against the context +that will actually run the query. That forks the session state, and the two halves of +the fork behave differently: + +- **Shared.** Catalogs, schemas, and tables. `SessionState` holds its catalog list + behind an `Arc`, so a table registered on either context is visible to both. The + runtime environment is shared for the same reason. +- **Copied.** Registered scalar, aggregate, and window functions, table functions, the + session configuration, and the analyzer and optimizer rule lists. These are + snapshotted when the derived context is created, so a UDF registered on the original + context afterwards is not visible to the derived one, and a `SET` applied to one does + not reach the other. + +The session id is carried over to the fork, so both contexts report the same id. + +Register functions before deriving, or register them directly on the derived context: + +```python +ctx = SessionContext(config) +ctx = ctx.with_logical_extension_codec(provider_logical_codec) +ctx = ctx.with_physical_extension_codec(provider_physical_codec) +ctx = ctx.with_query_planner(planner) +ctx.register_udf(my_udf) # registered on the context that will run the query +``` + +A session holds exactly one query planner. Calling `with_query_planner` again replaces +the installed planner instead of layering another one. To chain planners, have the new +planner wrap the capsule returned by `SessionContext.__datafusion_query_planner__()` +and delegate to it explicitly. + ## Alternative Approach Suppose you needed to expose some other features of DataFusion and you could not wait diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md index c897c067a..f2945251f 100644 --- a/examples/datafusion-ffi-example/README.md +++ b/examples/datafusion-ffi-example/README.md @@ -35,9 +35,7 @@ Separate shared libraries guarantee distinct DataFusion library markers. This ca The example codecs do not inspect the callback `TaskContext`. A production codec that depends on session configuration or registered functions must ensure its exported FFI codec is bound to, and retains, the appropriate host `TaskContextProvider`. -The current Python API installs one external logical codec and one external physical codec. It does not yet compose codecs from several independent plugin owners. This example therefore makes the provider library the sole external codec owner; the planner uses built-in physical nodes and receives the provider codecs from the host. - -Register both provider codecs before installing the planner: +This example makes the provider library the sole external codec owner. Register both provider codecs before installing the planner: ```python ctx = ctx.with_logical_extension_codec(provider_logical_codec) @@ -45,6 +43,6 @@ ctx = ctx.with_physical_extension_codec(provider_physical_codec) ctx = ctx.with_query_planner(planner) ``` -Derived contexts also rebind an installed planner when codecs change, but planner-last order is recommended because it states the ownership flow clearly. +Derived contexts also rebind an installed planner when codecs change, so this order is a recommendation rather than a requirement. Planner-last states the ownership flow more clearly. -Arbitrary custom `LogicalPlan::Extension` nodes are not supported by the current DataFusion FFI logical codec. This example covers foreign table providers, UDFs, and physical execution plans only. +For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. diff --git a/examples/datafusion-ffi-example/pyproject.toml b/examples/datafusion-ffi-example/pyproject.toml index 7f85e9487..c51fa8a8d 100644 --- a/examples/datafusion-ffi-example/pyproject.toml +++ b/examples/datafusion-ffi-example/pyproject.toml @@ -21,7 +21,8 @@ build-backend = "maturin" [project] name = "datafusion_ffi_example" -requires-python = ">=3.9" +# Matches the abi3-py310 feature the crate builds against. +requires-python = ">=3.10" classifiers = [ "Programming Language :: Rust", "Programming Language :: Python :: Implementation :: CPython", diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 53e2bddc5..4d04b7ab8 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -55,6 +55,6 @@ ctx = ctx.with_query_planner(MyQueryPlanner()) `PlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit. -The provider's codec pair is attached to the planner when the derived context is created and is also used to decode the returned physical plan in `datafusion-python`. The API currently supports one external codec owner rather than a registry of independently composed codecs, so this planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; derived contexts rebind codecs after planner installation, but planner-last order is easier to audit. +The provider's codec pair is attached to the planner when the derived context is created and is also used to decode the returned physical plan in `datafusion-python`. This planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; derived contexts rebind codecs after planner installation, but planner-last order is easier to audit. -The pinned FFI logical codec cannot encode arbitrary custom `LogicalPlan::Extension` nodes. The example therefore demonstrates table-provider, UDF, and physical-plan interoperability without claiming custom logical extension support. +For the limits behind that choice — why there is one external codec owner rather than a registry, which node kinds survive the boundary, and what a derived context shares with the context it came from — see [Query Planners Across Multiple Libraries](../../docs/source/contributor-guide/ffi.md#query-planners-across-multiple-libraries) in the contributor guide. diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index 8279879d1..16aa3933e 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -74,6 +74,32 @@ def test_three_library_query_planner(raw_capsule: bool): assert physical_codec.execution_plan_decode_calls() > 0 +def test_second_planner_replaces_the_first(): + """A session holds exactly one planner, so installing another replaces it.""" + ctx, _logical_codec, _physical_codec = configured_context(max_rows=2) + first = MyQueryPlanner() + second = MyQueryPlanner() + ctx = ctx.with_query_planner(first).with_query_planner(second) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert second.plan_calls() > 0 + assert first.plan_calls() == 0 + + +def test_planner_is_not_installed_on_the_original_context(): + """``with_query_planner`` returns a fork; the receiver keeps its planner.""" + ctx, _logical_codec, _physical_codec = configured_context(max_rows=2) + planner = MyQueryPlanner() + derived = ctx.with_query_planner(planner) + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert planner.plan_calls() == 0 + + derived.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert planner.plan_calls() > 0 + + def test_installed_codecs_outlive_python_exporters(): ctx, logical_codec, physical_codec = configured_context(max_rows=2) del logical_codec, physical_codec diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index cb767ffa5..56a9593b0 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -19,6 +19,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use async_trait::async_trait; +use datafusion::common::DataFusionError; use datafusion::execution::TaskContextProvider; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::LogicalPlan; @@ -69,41 +70,54 @@ fn physical_plan_has_foreign_plan(plan: &Arc) -> bool { .any(|child| physical_plan_has_foreign_plan(child)) } +/// The row limit as the host spells it, where `PlannerConfig` is registered as +/// an ordinary config extension under its own `ConfigExtension::PREFIX`. +const MAX_ROWS_KEY: &str = "ffi_query_planner.max_rows"; + +/// The same setting as it appears once the session has crossed the FFI +/// boundary. Rebuilding a `ConfigOptions` on this side parks every foreign +/// extension inside a single `FFI_ExtensionOptions`, which is itself a config +/// extension namespaced under `datafusion_ffi`, so `ConfigOptions::entries` +/// reports the key with both prefixes. +const FFI_MAX_ROWS_KEY: &str = "datafusion_ffi.ffi_query_planner.max_rows"; + fn planner_config(session: &dyn Session) -> datafusion::common::Result { let options = session.config_options(); - // Read the flattened entry first. Some DataFusion revisions add an extra - // `datafusion_ffi` namespace while reconstructing a ForeignSession. Parsing - // it directly also ensures malformed values are reported instead of being - // replaced silently by PlannerConfig::default(). - if let Some(entry) = options + // Prefer the raw entry. `local_or_ffi_extension` discards a value it cannot + // parse and hands back `PlannerConfig::default()`, which would quietly turn + // a typo into a different row limit instead of reporting it. + let config = match options .entries() .into_iter() - .find(|entry| entry.key.ends_with("ffi_query_planner.max_rows")) + .find(|entry| entry.key == MAX_ROWS_KEY || entry.key == FFI_MAX_ROWS_KEY) { - let value = entry.value.ok_or_else(|| { - datafusion::common::DataFusionError::Configuration(format!( - "{} must have a value", - entry.key - )) - })?; - let max_rows = value.parse::().map_err(|err| { - datafusion::common::DataFusionError::Configuration(format!( - "Invalid value '{value}' for {}: {err}", - entry.key - )) - })?; - if max_rows == 0 { - return Err(datafusion::common::DataFusionError::Configuration( - "ffi_query_planner.max_rows must be greater than zero".to_owned(), - )); + Some(entry) => { + let value = entry.value.ok_or_else(|| { + DataFusionError::Configuration(format!("{} must have a value", entry.key)) + })?; + let max_rows = value.parse::().map_err(|err| { + DataFusionError::Configuration(format!( + "Invalid value '{value}' for {}: {err}", + entry.key + )) + })?; + PlannerConfig { max_rows } } - return Ok(PlannerConfig { max_rows }); + None => options + .local_or_ffi_extension::() + .unwrap_or_default(), + }; + + // Validate after both paths so the fallback cannot smuggle in a limit that + // the direct path rejects. + if config.max_rows == 0 { + return Err(DataFusionError::Configuration(format!( + "{MAX_ROWS_KEY} must be greater than zero" + ))); } - Ok(options - .local_or_ffi_extension::() - .unwrap_or_default()) + Ok(config) } #[derive(Debug)] diff --git a/followup-issue.md b/followup-issue.md new file mode 100644 index 000000000..ddec293ae --- /dev/null +++ b/followup-issue.md @@ -0,0 +1,40 @@ +**Title:** Unify the hand-written PyCapsule importers with `from_pycapsule!` and add FFI version checks + +**Body:** + +### Is your feature request related to a problem or challenge? + +`crates/util/src/lib.rs` imports FFI objects from PyCapsules in two different styles. + +Most importers are generated by the `from_pycapsule!` macro (and its fallible sibling `try_from_pycapsule!`). Those take `&Bound`, call the dunder when present, produce a descriptive error when the object is not a capsule, run `validate_pycapsule`, and then `pointer_checked`. `physical_codec_from_pycapsule` and `physical_optimizer_rule_from_pycapsule` are generated this way. + +Two importers are hand-written and diverge from that shape: + +- `ffi_logical_codec_from_pycapsule` takes `Bound` by value, skips `validate_pycapsule`, and produces a bare cast error rather than a named one. +- `ffi_query_planner_from_pycapsule` takes `&Bound` and does run `validate_pycapsule`, but is still hand-written. + +Neither can use the macro today, because the macro returns `Arc<$output_type>` while both of these callers need the raw FFI handle. `PySessionContext::ctx_with_rebound_planner` needs the `FFI_QueryPlanner` value itself so it can rebuild it against replacement codecs, and the logical codec path needs `FFI_LogicalExtensionCodec` to construct a `PythonLogicalCodec`. + +Separately, none of the macro-generated importers check the DataFusion FFI version. `ffi_query_planner_from_pycapsule` is currently the only importer that does: + +```rust +let planner_version = unsafe { (planner.version)() }; +if planner_version != datafusion_ffi::version() { + return Err(PyImportError::new_err(...)); +} +``` + +Without that check, a capsule produced by an extension built against a different `datafusion-ffi` major version is dereferenced anyway. The failure mode is an ABI mismatch — most likely a crash — where a clean `PyImportError` is available for the cost of one comparison. + +### Describe the solution you'd like + +1. Add a macro variant alongside `from_pycapsule!` that returns the FFI type by clone instead of `Arc<$output_type>`, and generate both `ffi_logical_codec_from_pycapsule` and `ffi_query_planner_from_pycapsule` from it. `ffi_logical_codec_from_pycapsule` has two call sites (`crates/core/src/catalog.rs:692` and `crates/core/src/context.rs`), so the signature change from `Bound` to `&Bound` is contained. +2. Add the FFI version check to the shared macro bodies so every importer gets it, not just the query planner. This needs a decision on whether the check belongs on every FFI struct that exposes a `version` function pointer, or only on those where a mismatch is not otherwise caught. + +### Describe alternatives you've considered + +Leaving the two hand-written importers as they are. They work; the cost is that the house style is not obvious from reading the file, and the version check keeps getting added one importer at a time. + +### Additional context + +Came out of review of #1677, which added `ffi_query_planner_from_pycapsule`. That PR deliberately left the existing importers alone to stay scoped to the query planner. diff --git a/pr-body-addition.md b/pr-body-addition.md new file mode 100644 index 000000000..c53af308c --- /dev/null +++ b/pr-body-addition.md @@ -0,0 +1,10 @@ +# What changes are included in this PR? (additions to the existing list) + +- New `docs/source/contributor-guide/ffi.md` subsection "What a derived context shares", documenting that a codec change on a session with a foreign planner installed forks the session state, and which parts of that state are shared (catalogs, tables, runtime environment) versus snapshotted (registered functions, configuration, optimizer rule lists). +- `examples/datafusion-ffi-example/pyproject.toml` declares `requires-python = ">=3.10"` to match the `abi3-py310` feature the crate has always built against. It previously declared `>=3.9`, which advertised support for an interpreter the wheel cannot load. + +# Are there any user-facing changes? + +New public APIs: `SessionContext.with_query_planner` and `SessionContext.__datafusion_query_planner__`. A new example crate ships under `examples/`. No breaking changes to existing APIs. + +`with_query_planner` returns a context whose session state is forked from the receiver. Catalogs, tables, and the runtime environment remain shared with the original context, while registered functions, the session configuration, and the analyzer and optimizer rule lists are snapshotted at the time of the call. Installing a codec on a session that already has a foreign planner forks in the same way, because the planner has to be rebound to the new codec. Sessions with no foreign planner are unaffected and continue to share state as before. This is documented on the affected methods and in the contributor guide. diff --git a/python/datafusion/context.py b/python/datafusion/context.py index c7b73c5da..b1f4f74f0 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1775,12 +1775,26 @@ def with_query_planner( """Create a new session context with a custom query planner. The planner is imported through its ``__datafusion_query_planner__`` - PyCapsule. The returned context preserves the existing session state and - its logical and physical extension codec settings. Codec changes made on - a derived context are rebound to the planner before planning. - - Args: - planner: Object exposing ``__datafusion_query_planner__`` or a raw + PyCapsule. The returned context carries over the current session state + and the logical and physical extension codec settings. Codec changes + made on a derived context are rebound to the planner before planning. + + A session holds exactly one planner, so calling this again replaces the + previous one rather than layering. To chain planners, have the new + planner wrap the capsule from + :meth:`~SessionContext.__datafusion_query_planner__`. + + .. note:: Derived contexts share catalogs, not registrations + The returned context is a fork. Catalogs, schemas, and tables stay + shared with the original context, but registered functions and + configuration are copied at the time of the call. A UDF registered + on the original context afterwards is **not** visible here, while a + table registered on either context is visible to both. Register + functions before deriving, or register them on the derived context. + + Args: + planner: Object exposing ``__datafusion_query_planner__`` (see + :class:`QueryPlannerExportable`) or a raw ``datafusion_query_planner`` PyCapsule. Returns: @@ -2234,6 +2248,12 @@ def with_logical_extension_codec( Only FFI codecs are supported. Pass any object implementing ``__datafusion_logical_extension_codec__`` (see :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`). + + The returned context shares its session state with the original, so a + later registration on either is visible to both. The exception is a + session with a custom query planner installed: that planner has to be + rebound to the new codec, which forks the state. See + :meth:`~SessionContext.with_query_planner` for what a fork shares. """ new_internal = self.ctx.with_logical_extension_codec(codec) new = SessionContext.__new__(SessionContext) @@ -2252,6 +2272,12 @@ def with_physical_extension_codec( Only FFI codecs are supported. Pass any object implementing ``__datafusion_physical_extension_codec__`` (see :py:class:`~datafusion.user_defined.PhysicalExtensionCodecExportable`). + + The returned context shares its session state with the original, so a + later registration on either is visible to both. The exception is a + session with a custom query planner installed: that planner has to be + rebound to the new codec, which forks the state. See + :meth:`~SessionContext.with_query_planner` for what a fork shares. """ new_internal = self.ctx.with_physical_extension_codec(codec) new = SessionContext.__new__(SessionContext) @@ -2294,7 +2320,13 @@ def with_python_udf_inlining(self, *, enabled: bool) -> SessionContext: regardless of the toggle. Returns a new :class:`SessionContext` with the toggle applied; - the original session is unchanged. + the original session is unchanged. The returned context shares + its session state with the original, so a later registration on + either is visible to both. The exception is a session with a + custom query planner installed: that planner has to be rebound + to the new codecs, which forks the state. See + :meth:`~SessionContext.with_query_planner` for what a fork + shares. Examples: >>> import pyarrow as pa diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 6e6eaadbe..5ceb14e0e 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -754,6 +754,58 @@ def test_with_query_planner_capsule(ctx): assert batches[0].column(0) == pa.array([1]) +def test_derived_context_shares_catalogs(ctx): + """Catalogs live behind an Arc, so tables cross the fork in both directions.""" + derived = ctx.with_query_planner(ctx.__datafusion_query_planner__()) + + ctx.register_record_batches( + "registered_on_parent", + [[pa.RecordBatch.from_pydict({"value": [1]})]], + ) + derived.register_record_batches( + "registered_on_derived", + [[pa.RecordBatch.from_pydict({"value": [2]})]], + ) + + assert derived.table_exist("registered_on_parent") + assert ctx.table_exist("registered_on_derived") + + +def test_derived_context_snapshots_functions(ctx): + """Function registries are copied at fork time, unlike catalogs. + + A UDF registered on the parent before the fork is carried over; one + registered afterwards is not. Guards the caveat documented on + ``SessionContext.with_query_planner``. + """ + before = udf( + lambda arr: arr, + [pa.int64()], + pa.int64(), + volatility="immutable", + name="registered_before_fork", + ) + ctx.register_udf(before) + + derived = ctx.with_query_planner(ctx.__datafusion_query_planner__()) + + after = udf( + lambda arr: arr, + [pa.int64()], + pa.int64(), + volatility="immutable", + name="registered_after_fork", + ) + ctx.register_udf(after) + + assert derived.sql("SELECT registered_before_fork(1)").collect() + with pytest.raises(Exception, match="registered_after_fork"): + derived.sql("SELECT registered_after_fork(1)").collect() + + # The parent is unaffected by the fork. + assert ctx.sql("SELECT registered_after_fork(1)").collect() + + def test_table_provider(ctx): batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) ctx.register_record_batches("provider_test", [[batch]]) From 6ef7cf933ec866c7038e5c44b03b479c1a6145f3 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 26 Aug 2026 10:48:06 -0400 Subject: [PATCH 07/30] refactor: drop the runtime adapter and fix exported capsule lifetimes Remove `RuntimeAwareQueryPlanner`. It existed to re-attach our Tokio handle to the session we hand to a foreign planner, on the reasoning that `ForeignQueryPlanner` passes `session_runtime: None`. That handle turns out to have no reachable path: the query planner FFI exchanges serialized bytes rather than plan handles, a provider owned by another library keeps its own runtime because `FFI_TableProvider::new_with_ffi_codec` unwraps a `ForeignTableProvider` back to the original handle, and we execute on our own runtime regardless. Setting the handle to `None` left every test passing. Codec rebinding now downcasts upstream's `ForeignQueryPlanner` directly, which also stops `__datafusion_query_planner__` adding a second layer, since `new_with_ffi_codecs` already unwraps that type. The `datafusion-session` dependency is no longer needed in crates/core. Keep the exporting session alive for codecs handed out in a PyCapsule. `FFI_TaskContextProvider` stores its provider in a `Weak`, so a capsule stopped working as soon as the `SessionContext` that produced it went out of scope. That made the natural spelling of the documented fallback pattern fail: fallback = ctx.__datafusion_query_planner__() ctx = ctx.with_query_planner(MyPlanner(fallback=fallback)) Rebinding `ctx` dropped the exporter and planning then failed with "TaskContextProvider went out of scope over FFI boundary". Both Python codecs gained an opt-in `exported_session`, set only by the three capsule getters. The keep-alive lives in the inner codec because the consumer clones the FFI handle out of the capsule and `clone` clones the inner codec's `Arc`, so a capsule-scoped keep-alive would die too early. It is deliberately opt-in: the same codecs are also attached to providers and catalogs that end up back inside the session, where a strong reference would close a `SessionContext -> SessionState -> query planner -> FFI codec` cycle. Both structs now implement `Debug` by hand, because `SessionContext` is not `Debug`. Add two example tests. One drives a plan containing `RepartitionExec`, which spawns Tokio tasks as it runs, through all three libraries, so the codecs are exercised on a multi-node plan rather than a bare scan. The other layers a planner on top of the session's existing planner using the capsule captured beforehand, which is the delegation pattern upstream prescribes; `Session::create_physical_plan` cannot be used for this, because it dispatches through the installed planner and recurses. Co-Authored-By: Claude Opus 5 (1M context) --- Cargo.lock | 1 - crates/core/Cargo.toml | 1 - crates/core/src/codec.rs | 67 ++++++++- crates/core/src/context.rs | 138 +++++++----------- .../_test_three_library_query_planner.py | 60 ++++++++ .../src/planner.rs | 70 ++++++++- 6 files changed, 239 insertions(+), 98 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a9d54a1f1..11736d8ff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1554,7 +1554,6 @@ dependencies = [ "datafusion-ffi", "datafusion-proto", "datafusion-python-util", - "datafusion-session", "datafusion-spark", "datafusion-substrait", "futures", diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index 91a1d5f77..c5f1e0167 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -54,7 +54,6 @@ datafusion-substrait = { workspace = true, optional = true } datafusion-proto = { workspace = true } datafusion-ffi = { workspace = true } datafusion-spark = { workspace = true, features = ["core"] } -datafusion-session = { workspace = true } prost = { workspace = true } # keep in line with `datafusion-substrait` serde_json = { workspace = true } uuid = { workspace = true, features = ["v4"] } diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs index 26853e69f..1c6658939 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -103,6 +103,7 @@ use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; use datafusion::physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, @@ -233,10 +234,31 @@ fn strip_wire_header<'a>( /// Sitting at the top of the session's logical codec stack means /// every serializer that reads `session.logical_codec()` automatically /// picks up Python-aware encoding for free. -#[derive(Debug)] pub struct PythonLogicalCodec { inner: Arc, python_udf_inlining: bool, + /// Keeps the exporting session alive for a codec handed across the FFI + /// boundary. + /// + /// `FFI_TaskContextProvider` stores its provider in a `Weak`, so an exported + /// codec stops working the moment the object that produced it goes out of + /// scope. Retaining the session in the inner codec survives both that and + /// `clone`, which clones the inner codec's `Arc` and so carries this along. + /// + /// Set this only on codecs that are leaving this library. A codec installed + /// *in* a session must not hold one, or the session would own itself: + /// `SessionContext -> SessionState -> query planner -> FFI codec -> here`. + exported_session: Option>, +} + +impl std::fmt::Debug for PythonLogicalCodec { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PythonLogicalCodec") + .field("inner", &self.inner) + .field("python_udf_inlining", &self.python_udf_inlining) + .finish_non_exhaustive() + } } impl PythonLogicalCodec { @@ -244,9 +266,19 @@ impl PythonLogicalCodec { Self { inner, python_udf_inlining: true, + exported_session: None, } } + /// Retain `ctx` so this codec keeps working after the exporting + /// `SessionContext` goes out of scope. Only for codecs being exported over + /// FFI; see the `exported_session` field for why installed codecs must not + /// use this. + pub fn with_exported_session(mut self, ctx: Arc) -> Self { + self.exported_session = Some(ctx); + self + } + pub fn inner(&self) -> &Arc { &self.inner } @@ -443,10 +475,31 @@ fn refuse_inline_payload(kind: &str, name: &str) -> datafusion::error::DataFusio /// would round-trip at the logical level but break at the physical /// level. Both layers reuse the shared payload framing /// ([`PY_SCALAR_UDF_FAMILY`] et al.) so the wire format is identical. -#[derive(Debug)] pub struct PythonPhysicalCodec { inner: Arc, python_udf_inlining: bool, + /// Keeps the exporting session alive for a codec handed across the FFI + /// boundary. + /// + /// `FFI_TaskContextProvider` stores its provider in a `Weak`, so an exported + /// codec stops working the moment the object that produced it goes out of + /// scope. Retaining the session in the inner codec survives both that and + /// `clone`, which clones the inner codec's `Arc` and so carries this along. + /// + /// Set this only on codecs that are leaving this library. A codec installed + /// *in* a session must not hold one, or the session would own itself: + /// `SessionContext -> SessionState -> query planner -> FFI codec -> here`. + exported_session: Option>, +} + +impl std::fmt::Debug for PythonPhysicalCodec { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PythonPhysicalCodec") + .field("inner", &self.inner) + .field("python_udf_inlining", &self.python_udf_inlining) + .finish_non_exhaustive() + } } impl PythonPhysicalCodec { @@ -454,9 +507,19 @@ impl PythonPhysicalCodec { Self { inner, python_udf_inlining: true, + exported_session: None, } } + /// Retain `ctx` so this codec keeps working after the exporting + /// `SessionContext` goes out of scope. Only for codecs being exported over + /// FFI; see the `exported_session` field for why installed codecs must not + /// use this. + pub fn with_exported_session(mut self, ctx: Arc) -> Self { + self.exported_session = Some(ctx); + self + } + pub fn inner(&self) -> &Arc { &self.inner } diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 59e8b9707..5b9f09a06 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -24,7 +24,6 @@ use std::sync::Arc; use arrow::array::RecordBatchReader; use arrow::ffi_stream::ArrowArrayStreamReader; use arrow::pyarrow::FromPyArrow; -use async_trait::async_trait; use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion::arrow::pyarrow::PyArrowType; use datafusion::arrow::record_batch::RecordBatch; @@ -45,8 +44,6 @@ use datafusion::execution::options::{ArrowReadOptions, ReadOptions}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::execution::{FunctionRegistry, TaskContextProvider}; -use datafusion::logical_expr::LogicalPlan; -use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::{ AvroReadOptions, CsvReadOptions, DataFrame, JsonReadOptions, ParquetReadOptions, }; @@ -56,7 +53,7 @@ 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::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; -use datafusion_ffi::query_planner::FFI_QueryPlanner; +use datafusion_ffi::query_planner::{FFI_QueryPlanner, ForeignQueryPlanner}; use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory; use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::physical_plan::PhysicalExtensionCodec; @@ -67,7 +64,6 @@ use datafusion_python_util::{ physical_codec_from_pycapsule, physical_optimizer_rule_from_pycapsule, spawn_future, wait_for_future, }; -use datafusion_session::Session; use object_store::ObjectStore; use pyo3::IntoPyObjectExt; use pyo3::exceptions::{PyKeyError, PyRuntimeError, PyValueError}; @@ -227,55 +223,6 @@ impl PySessionConfig { } } -/// Adapts an FFI planner to the Tokio runtime owned by datafusion-python. -/// -/// Upstream's `ForeignQueryPlanner` is the consumer-side adapter that lets an -/// `FFI_QueryPlanner` satisfy the `QueryPlanner` trait, which is what makes a -/// planner from another shared library installable in a `SessionState`. Its -/// trait method receives only a `&LogicalPlan` and a `&dyn Session`, so it has -/// nowhere to obtain a runtime handle and calls -/// `create_physical_plan_with_session_runtime` with `None`. -/// -/// Throughout datafusion-ffi each library attaches its *own* runtime to the -/// objects it exports, so a producer-side wrapper can `Handle::enter` before -/// running that library's code. A provider owned by another library keeps its -/// owner's runtime even when it travels through our catalog, because -/// `FFI_TableProvider::new_with_ffi_codec` unwraps a `ForeignTableProvider` -/// back to the original handle and discards the runtime passed alongside it. -/// `session_runtime` is that same rule applied to the session: `FFI_SessionRef` -/// is our object, and every callback on it runs our code. -/// -/// It matters for what those callbacks hand back. A plan produced by our own -/// planner returns as `FFI_ExecutionPlan::new(plan, runtime)`, and `execute` -/// enters that runtime before calling into the plan. The same holds for our -/// physical optimizer rules and for tables we own rather than re-export. The -/// delegation case this type exists for is exactly that shape: a foreign -/// planner that falls back to our planner through -/// `__datafusion_query_planner__` receives a plan whose execution needs our -/// runtime, and datafusion-python owns that runtime as a process global while -/// the Python thread calling in carries no ambient one. -/// -/// This adapter is where the handle gets attached. It wraps the foreign handle -/// and calls `create_physical_plan_with_session_runtime` with `Some(handle)`. -#[derive(Debug)] -struct RuntimeAwareQueryPlanner { - planner: FFI_QueryPlanner, -} - -#[async_trait] -impl QueryPlanner for RuntimeAwareQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session: &dyn Session, - ) -> datafusion::common::Result> { - let runtime = get_tokio_runtime().handle().clone(); - self.planner - .create_physical_plan_with_session_runtime(logical_plan, session, Some(runtime)) - .await - } -} - /// Runtime options for a SessionContext #[pyclass( from_py_object, @@ -1448,27 +1395,23 @@ impl PySessionContext { &self, py: Python<'py>, ) -> PyResult> { - let ffi = self.ffi_logical_codec(); - create_logical_extension_capsule(py, ffi.as_ref()) + create_logical_extension_capsule(py, &self.exported_ffi_logical_codec()) } pub fn __datafusion_query_planner__<'py>( &self, py: Python<'py>, ) -> PyResult> { - // When the installed planner is already foreign, this exports a handle - // that wraps `RuntimeAwareQueryPlanner`, which in turn wraps the - // original foreign handle. That extra layer looks redundant but is not: - // a consumer reaching us through `ForeignQueryPlanner` always calls - // `create_physical_plan_with_session_runtime` with `None`, so the - // adapter is what puts our Tokio handle back on the session before the - // call continues outward. Unwrapping to the inner handle here would - // save one planning-time round trip and silently drop that runtime. + // An already-foreign planner is re-exported as its original handle + // rather than gaining another layer, because `new_with_ffi_codecs` + // unwraps a `ForeignQueryPlanner`. It still adopts the codecs supplied + // here, so a consumer that wraps this capsule decodes our plans with + // our codecs. let planner = Arc::clone(self.ctx.state().query_planner()); let ffi = FFI_QueryPlanner::new_with_ffi_codecs( planner, - self.ffi_logical_codec().as_ref().clone(), - self.ffi_physical_codec().as_ref().clone(), + self.exported_ffi_logical_codec(), + self.exported_ffi_physical_codec(), ); create_query_planner_capsule(py, &ffi) } @@ -1494,8 +1437,7 @@ impl PySessionContext { &self, py: Python<'py>, ) -> PyResult> { - let ffi = self.ffi_physical_codec(); - create_physical_extension_capsule(py, ffi.as_ref()) + create_physical_extension_capsule(py, &self.exported_ffi_physical_codec()) } pub fn with_physical_extension_codec<'py>( @@ -1560,13 +1502,15 @@ impl PySessionContext { ) -> Arc { let state = self.ctx.state(); - // When the caller is only replacing codecs, recover the foreign planner - // already installed so it can be rebound below. + // When the caller is only replacing codecs, recover the handle behind + // the installed planner so it can be rebound below. A planner this + // library owns is not a `ForeignQueryPlanner` and needs no rebinding, + // because it does not carry codecs of its own. let planner = planner.or_else(|| { let installed: &dyn std::any::Any = state.query_planner().as_ref(); installed - .downcast_ref::() - .map(|planner| planner.planner.clone()) + .downcast_ref::() + .map(|planner| planner.0.clone()) }); let Some(planner) = planner else { @@ -1575,13 +1519,12 @@ impl PySessionContext { let ctx = Arc::new(SessionContext::new_with_state(state)); let inner: Arc = (&planner).into(); - let planner = Arc::new(RuntimeAwareQueryPlanner { - planner: FFI_QueryPlanner::new_with_ffi_codecs( - inner, - Self::ffi_logical_codec_for(&ctx, logical_codec), - Self::ffi_physical_codec_for(&ctx, physical_codec), - ), - }); + let planner: Arc = (&FFI_QueryPlanner::new_with_ffi_codecs( + inner, + Self::ffi_logical_codec_for(&ctx, logical_codec), + Self::ffi_physical_codec_for(&ctx, physical_codec), + )) + .into(); let state = SessionStateBuilder::new_from_existing(ctx.state()) .with_query_planner(planner) .build(); @@ -1665,14 +1608,6 @@ impl PySessionContext { FFI_LogicalExtensionCodec::new(codec, Some(runtime), &ctx_provider) } - /// Build an FFI-wrapped clone of the session's physical codec on demand. - pub(crate) fn ffi_physical_codec(&self) -> Arc { - Arc::new(Self::ffi_physical_codec_for( - &self.ctx, - &self.physical_codec, - )) - } - fn ffi_physical_codec_for( ctx: &Arc, codec: &Arc, @@ -1683,6 +1618,35 @@ impl PySessionContext { let ctx_provider = Arc::clone(ctx) as Arc; FFI_PhysicalExtensionCodec::new(codec, Some(runtime), &ctx_provider) } + + /// Build an FFI-wrapped logical codec for handing out in a PyCapsule. + /// + /// Same as [`Self::ffi_logical_codec`] except the inner codec retains this + /// session. The FFI task-context handle is weak, so without that a capsule + /// stops working as soon as the exporting `SessionContext` goes out of + /// scope, which the natural `ctx = ctx.with_query_planner(planner)` does. + /// + /// Only for codecs leaving this library. The plain builder is still correct + /// for codecs attached to objects that end up back inside this session, + /// which would otherwise make the session own itself. + fn exported_ffi_logical_codec(&self) -> FFI_LogicalExtensionCodec { + let codec = Arc::new( + PythonLogicalCodec::new(Arc::clone(self.logical_codec.inner())) + .with_python_udf_inlining(self.logical_codec.python_udf_inlining()) + .with_exported_session(Arc::clone(&self.ctx)), + ); + Self::ffi_logical_codec_for(&self.ctx, &codec) + } + + /// Physical companion to [`Self::exported_ffi_logical_codec`]. + fn exported_ffi_physical_codec(&self) -> FFI_PhysicalExtensionCodec { + let codec = Arc::new( + PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner())) + .with_python_udf_inlining(self.physical_codec.python_udf_inlining()) + .with_exported_session(Arc::clone(&self.ctx)), + ); + Self::ffi_physical_codec_for(&self.ctx, &codec) + } } pub fn parse_file_compression_type( diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index 16aa3933e..15f46bf90 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -74,6 +74,66 @@ def test_three_library_query_planner(raw_capsule: bool): assert physical_codec.execution_plan_decode_calls() > 0 +def test_spawning_plan_across_three_libraries(): + """A plan that spawns Tokio tasks survives the full three-library round trip. + + ``target_partitions`` above one puts a ``RepartitionExec`` under the + aggregate, and that operator spawns tasks while it runs. This exercises the + codecs on a multi-node plan rather than the bare scan the other tests use. + """ + config = SessionConfig().with_extension(PlannerConfig(max_rows=100)) + config = config.with_target_partitions(4) + logical_codec = MyLogicalExtensionCodec() + physical_codec = MyPhysicalExtensionCodec() + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(physical_codec) + ctx.register_table("numbers", MyTableProvider(1, 6, 3)) + + planner = MyQueryPlanner() + ctx = ctx.with_query_planner(planner) + + batches = ctx.sql( + 'SELECT "A" % 2 AS parity, count(*) AS n FROM numbers GROUP BY 1 ORDER BY 1' + ).collect() + counts = { + row[0]: row[1] + for batch in batches + for row in zip( + batch.column(0).to_pylist(), batch.column(1).to_pylist(), strict=True + ) + } + assert sum(counts.values()) == 6 + 7 + 8 + assert planner.plan_calls() > 0 + assert planner.foreign_provider_observed() + + +def test_planner_layers_on_the_session_planner(): + """A planner can wrap the one already installed and delegate to it. + + The capsule has to be captured before this planner is installed, because + ``__datafusion_query_planner__`` exports whichever planner is installed when + it is called. Capturing it afterwards would hand the planner a handle to + itself, and planning would recurse. + """ + ctx, logical_codec, physical_codec = configured_context(max_rows=3) + fallback = ctx.__datafusion_query_planner__() + planner = MyQueryPlanner(fallback=fallback) + # Rebinding `ctx` drops the context that produced the capsule. The exported + # codecs retain it, so the capsule stays usable. Without that the FFI + # task-context handle is weak and planning fails with "TaskContextProvider + # went out of scope over FFI boundary". + ctx = ctx.with_query_planner(planner) + gc.collect() + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert planner.plan_calls() > 0 + assert planner.used_fallback() + assert logical_codec.table_provider_decode_calls() > 0 + assert physical_codec.execution_plan_decode_calls() > 0 + + def test_second_planner_replaces_the_first(): """A session holds exactly one planner, so installing another replaces it.""" ctx, _logical_codec, _physical_codec = configured_context(max_rows=2) diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index 56a9593b0..59791f160 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::fmt; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; @@ -34,20 +35,31 @@ use datafusion_ffi::session::ForeignSession; use datafusion_ffi::table_provider::ForeignTableProvider; use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; -use datafusion_python_util::get_tokio_runtime; +use datafusion_python_util::{ffi_query_planner_from_pycapsule, get_tokio_runtime}; use datafusion_session::{QueryPlanner, Session}; use pyo3::prelude::*; use pyo3::types::PyCapsule; use crate::config::PlannerConfig; -#[derive(Debug, Default)] +#[derive(Default)] struct PlannerObservations { plan_calls: AtomicUsize, last_max_rows: AtomicUsize, foreign_session: AtomicBool, foreign_provider: AtomicBool, foreign_plan: AtomicBool, + used_fallback: AtomicBool, +} + +impl fmt::Debug for PlannerObservations { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PlannerObservations") + .field("plan_calls", &self.plan_calls) + .field("last_max_rows", &self.last_max_rows) + .finish_non_exhaustive() + } } fn logical_plan_has_foreign_provider(plan: &LogicalPlan) -> bool { @@ -123,6 +135,18 @@ fn planner_config(session: &dyn Session) -> datafusion::common::Result, + /// Planner to hand the work to instead of planning here. + /// + /// This is how a real planner layers on top of an existing one. The capsule + /// must be captured from the session *before* this planner is installed: + /// `SessionContext.__datafusion_query_planner__` exports whatever planner + /// is installed at the time it is called, so capturing it afterwards would + /// hand this planner a handle to itself. + /// + /// Note that `Session::create_physical_plan` cannot be used for this. It + /// dispatches through the session's installed query planner, so calling it + /// from inside that planner recurses until the stack overflows. + fallback: Option>, } #[async_trait] @@ -146,9 +170,19 @@ impl QueryPlanner for DistributedQueryPlanner { .last_max_rows .store(config.max_rows, Ordering::SeqCst); - let plan = DefaultPhysicalPlanner::default() - .create_physical_plan(logical_plan, session) - .await?; + let plan = match self.fallback.as_ref() { + Some(fallback) => { + self.observations + .used_fallback + .store(true, Ordering::SeqCst); + fallback.create_physical_plan(logical_plan, session).await? + } + None => { + DefaultPhysicalPlanner::default() + .create_physical_plan(logical_plan, session) + .await? + } + }; self.observations .foreign_plan .store(physical_plan_has_foreign_plan(&plan), Ordering::SeqCst); @@ -170,13 +204,34 @@ impl QueryPlanner for DistributedQueryPlanner { #[derive(Debug, Default, Clone)] pub(crate) struct MyQueryPlanner { observations: Arc, + fallback: Option>, } #[pymethods] impl MyQueryPlanner { + /// Build a planner, optionally layered on top of an existing one. + /// + /// `fallback` takes anything exporting `__datafusion_query_planner__`, + /// including a `SessionContext`. Capture it *before* installing this + /// planner on that context, or the capsule will describe this planner and + /// planning will recurse. #[new] - fn new() -> Self { - Self::default() + #[pyo3(signature = (fallback=None))] + fn new(fallback: Option>) -> PyResult { + let fallback = fallback + .map(|planner| { + ffi_query_planner_from_pycapsule(&planner) + .map(|ffi| -> Arc { (&ffi).into() }) + }) + .transpose()?; + Ok(Self { + fallback, + ..Self::default() + }) + } + + fn used_fallback(&self) -> bool { + self.observations.used_fallback.load(Ordering::SeqCst) } fn plan_calls(&self) -> usize { @@ -205,6 +260,7 @@ impl MyQueryPlanner { ) -> PyResult> { let planner: Arc = Arc::new(DistributedQueryPlanner { observations: Arc::clone(&self.observations), + fallback: self.fallback.clone(), }); let runtime = get_tokio_runtime().handle().clone(); let ctx_provider = Arc::new(SessionContext::new()) as Arc; From c6ba98ff042ff18fa82b3776eee3875d74a536de Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 26 Aug 2026 10:49:49 -0400 Subject: [PATCH 08/30] add override for datafusion version to pre-release testing of upstream fix --- Cargo.lock | 76 +++++++++++++++++++++++++++--------------------------- Cargo.toml | 22 ++++++++-------- 2 files changed, 49 insertions(+), 49 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 11736d8ff..5549ae8ad 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -792,7 +792,7 @@ dependencies = [ [[package]] name = "datafusion" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "arrow-schema", @@ -845,7 +845,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "async-trait", @@ -869,7 +869,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "async-trait", @@ -892,7 +892,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "arrow-ipc", @@ -918,7 +918,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "futures", "log", @@ -928,7 +928,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "async-compression", @@ -964,7 +964,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "arrow-ipc", @@ -988,7 +988,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "arrow-avro", @@ -1006,7 +1006,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "async-trait", @@ -1029,7 +1029,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "async-trait", @@ -1052,7 +1052,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "arrow-schema", @@ -1084,12 +1084,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" [[package]] name = "datafusion-execution" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "arrow-buffer", @@ -1114,7 +1114,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "arrow-schema", @@ -1138,7 +1138,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "datafusion-common", @@ -1149,7 +1149,7 @@ dependencies = [ [[package]] name = "datafusion-ffi" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "arrow-schema", @@ -1220,7 +1220,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "arrow-buffer", @@ -1251,7 +1251,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "datafusion-common", @@ -1271,7 +1271,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "datafusion-common", @@ -1282,7 +1282,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "arrow-ord", @@ -1306,7 +1306,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "async-trait", @@ -1321,7 +1321,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "datafusion-common", @@ -1337,7 +1337,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1346,7 +1346,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "datafusion-doc", "quote", @@ -1356,7 +1356,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "chrono", @@ -1375,7 +1375,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "datafusion-common", @@ -1397,7 +1397,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "datafusion-common", @@ -1411,7 +1411,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "chrono", @@ -1428,7 +1428,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "datafusion-common", @@ -1447,7 +1447,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "arrow-data", @@ -1483,7 +1483,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "datafusion-catalog", @@ -1509,7 +1509,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "datafusion-common", @@ -1519,7 +1519,7 @@ dependencies = [ [[package]] name = "datafusion-proto-models" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "datafusion-common", "datafusion-proto-common", @@ -1529,7 +1529,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "datafusion-common", @@ -1589,7 +1589,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow-schema", "async-trait", @@ -1603,7 +1603,7 @@ dependencies = [ [[package]] name = "datafusion-spark" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "bigdecimal", @@ -1632,7 +1632,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "arrow", "bigdecimal", @@ -1651,7 +1651,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "55.0.0" -source = "git+https://github.com/apache/datafusion?rev=55.0.0-rc3#d5552342012888b7d1a3ab88d92e3d292fc0cde0" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 197bb8eb9..7efe64a7c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -78,14 +78,14 @@ codegen-units = 2 # We cannot publish to crates.io with any patches in the below section. Developers # must remove any entries in this section before creating a release candidate. [patch.crates-io] -datafusion = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } -datafusion-substrait = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } -datafusion-proto = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } -datafusion-ffi = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } -datafusion-catalog = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } -datafusion-common = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } -datafusion-functions-aggregate = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } -datafusion-functions-window = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } -datafusion-spark = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } -datafusion-expr = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } -datafusion-session = { git = "https://github.com/apache/datafusion", rev = "55.0.0-rc3" } +datafusion = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } +datafusion-substrait = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } +datafusion-proto = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } +datafusion-ffi = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } +datafusion-catalog = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } +datafusion-common = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } +datafusion-functions-aggregate = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } +datafusion-functions-window = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } +datafusion-spark = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } +datafusion-expr = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } +datafusion-session = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } From 752243bad69ca4565f84f28638ed2a16d3fea400 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 26 Aug 2026 11:19:05 -0400 Subject: [PATCH 09/30] remove unintentionally committed files --- followup-issue.md | 40 ---------------------------------------- pr-body-addition.md | 10 ---------- 2 files changed, 50 deletions(-) delete mode 100644 followup-issue.md delete mode 100644 pr-body-addition.md diff --git a/followup-issue.md b/followup-issue.md deleted file mode 100644 index ddec293ae..000000000 --- a/followup-issue.md +++ /dev/null @@ -1,40 +0,0 @@ -**Title:** Unify the hand-written PyCapsule importers with `from_pycapsule!` and add FFI version checks - -**Body:** - -### Is your feature request related to a problem or challenge? - -`crates/util/src/lib.rs` imports FFI objects from PyCapsules in two different styles. - -Most importers are generated by the `from_pycapsule!` macro (and its fallible sibling `try_from_pycapsule!`). Those take `&Bound`, call the dunder when present, produce a descriptive error when the object is not a capsule, run `validate_pycapsule`, and then `pointer_checked`. `physical_codec_from_pycapsule` and `physical_optimizer_rule_from_pycapsule` are generated this way. - -Two importers are hand-written and diverge from that shape: - -- `ffi_logical_codec_from_pycapsule` takes `Bound` by value, skips `validate_pycapsule`, and produces a bare cast error rather than a named one. -- `ffi_query_planner_from_pycapsule` takes `&Bound` and does run `validate_pycapsule`, but is still hand-written. - -Neither can use the macro today, because the macro returns `Arc<$output_type>` while both of these callers need the raw FFI handle. `PySessionContext::ctx_with_rebound_planner` needs the `FFI_QueryPlanner` value itself so it can rebuild it against replacement codecs, and the logical codec path needs `FFI_LogicalExtensionCodec` to construct a `PythonLogicalCodec`. - -Separately, none of the macro-generated importers check the DataFusion FFI version. `ffi_query_planner_from_pycapsule` is currently the only importer that does: - -```rust -let planner_version = unsafe { (planner.version)() }; -if planner_version != datafusion_ffi::version() { - return Err(PyImportError::new_err(...)); -} -``` - -Without that check, a capsule produced by an extension built against a different `datafusion-ffi` major version is dereferenced anyway. The failure mode is an ABI mismatch — most likely a crash — where a clean `PyImportError` is available for the cost of one comparison. - -### Describe the solution you'd like - -1. Add a macro variant alongside `from_pycapsule!` that returns the FFI type by clone instead of `Arc<$output_type>`, and generate both `ffi_logical_codec_from_pycapsule` and `ffi_query_planner_from_pycapsule` from it. `ffi_logical_codec_from_pycapsule` has two call sites (`crates/core/src/catalog.rs:692` and `crates/core/src/context.rs`), so the signature change from `Bound` to `&Bound` is contained. -2. Add the FFI version check to the shared macro bodies so every importer gets it, not just the query planner. This needs a decision on whether the check belongs on every FFI struct that exposes a `version` function pointer, or only on those where a mismatch is not otherwise caught. - -### Describe alternatives you've considered - -Leaving the two hand-written importers as they are. They work; the cost is that the house style is not obvious from reading the file, and the version check keeps getting added one importer at a time. - -### Additional context - -Came out of review of #1677, which added `ffi_query_planner_from_pycapsule`. That PR deliberately left the existing importers alone to stay scoped to the query planner. diff --git a/pr-body-addition.md b/pr-body-addition.md deleted file mode 100644 index c53af308c..000000000 --- a/pr-body-addition.md +++ /dev/null @@ -1,10 +0,0 @@ -# What changes are included in this PR? (additions to the existing list) - -- New `docs/source/contributor-guide/ffi.md` subsection "What a derived context shares", documenting that a codec change on a session with a foreign planner installed forks the session state, and which parts of that state are shared (catalogs, tables, runtime environment) versus snapshotted (registered functions, configuration, optimizer rule lists). -- `examples/datafusion-ffi-example/pyproject.toml` declares `requires-python = ">=3.10"` to match the `abi3-py310` feature the crate has always built against. It previously declared `>=3.9`, which advertised support for an interpreter the wheel cannot load. - -# Are there any user-facing changes? - -New public APIs: `SessionContext.with_query_planner` and `SessionContext.__datafusion_query_planner__`. A new example crate ships under `examples/`. No breaking changes to existing APIs. - -`with_query_planner` returns a context whose session state is forked from the receiver. Catalogs, tables, and the runtime environment remain shared with the original context, while registered functions, the session configuration, and the analyzer and optimizer rule lists are snapshotted at the time of the call. Installing a codec on a session that already has a foreign planner forks in the same way, because the planner has to be rebound to the new codec. Sessions with no foreign planner are unaffected and continue to share state as before. This is documented on the affected methods and in the contributor guide. From e9dc22c592654f737f4980d7857a2b60c6fab6df Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 26 Aug 2026 14:12:40 -0400 Subject: [PATCH 10/30] Empty commit to trigger CI From 2169e39f7f4d4fad642ade993ee545fe180c161f Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 26 Aug 2026 15:50:47 -0400 Subject: [PATCH 11/30] fix: keep the example planner's exported task context alive `FFI_TaskContextProvider` downgrades the provider it is given to a `Weak`, so building one inline in `__datafusion_query_planner__` left the capsule carrying a provider that was already dropped by the time it returned. Every codec callback through that capsule would have failed with "TaskContextProvider went out of scope over FFI boundary". The example did not notice because it ships the default codecs and no custom extension nodes, so `try_decode` is never reached. `MyQueryPlanner` now owns the context and hands out clones of it. The `QueryPlanner` the capsule carries holds a reference too, so the capsule stays usable even when the Python object that exported it is dropped first. Document the distinction the inline construction obscured. The `TaskContextProvider` supplied at export time backs the exporting library's own codec callbacks, decoding that library's nodes in its own registry. It is unrelated to the `&dyn Session` that later arrives at `create_physical_plan`, which belongs to the host, and it could not be derived from that session in any case, since the codecs are built before any session exists. Rename `PlannerConfig` to `MyPlannerConfig` to match `MyQueryPlanner`. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/contributor-guide/ffi.md | 26 +++++++ .../README.md | 4 +- .../_test_three_library_query_planner.py | 10 +-- .../src/config.rs | 16 ++-- .../src/lib.rs | 4 +- .../src/planner.rs | 76 ++++++++++++++++--- 6 files changed, 110 insertions(+), 26 deletions(-) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index ded7b4295..8a8b65a8d 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -257,6 +257,32 @@ The current FFI logical codec supports providers and UDFs but not arbitrary cust `LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and local build commands. +### The task context a planner exports is not the session it plans against + +`FFI_QueryPlanner::new` takes a `TaskContextProvider` alongside the two codecs. It is +easy to mistake this for the session the planner will be asked to plan against, but the +two are unrelated and serve opposite directions of the exchange: + +- The `&dyn Session` passed to `QueryPlanner::create_physical_plan` belongs to the + host. It arrives as a `ForeignSession` when the call crossed the boundary, and it is + what the planner reads configuration and catalogs from. +- The `TaskContextProvider` given at export time backs the *exported codecs*. When the + host decodes the plan bytes the planner returned, any extension node in them is + decoded by a callback back into the planner's own library, and that callback needs a + `TaskContext` whose registry can resolve that library's own nodes and functions. + +So a planner library should hand over a context it owns, and register its own UDFs and +extension types on it. It cannot use the host's session for this even if it wanted to: +the codecs are built when the capsule is exported, long before any session shows up. + +That context must be kept alive by the exporter. `FFI_TaskContextProvider` downgrades +it to a `Weak`, so a provider constructed inline in the capsule getter is already +dropped by the time the capsule is used, and every codec callback fails with +`TaskContextProvider went out of scope over FFI boundary`. Hold the reference on +something that lives at least as long as the exported planner — the example stores it +on both the Python-facing planner object and the `QueryPlanner` the capsule carries, so +the capsule keeps working even if the Python object is dropped first. + ### What a derived context shares `with_query_planner`, `with_logical_extension_codec`, `with_physical_extension_codec`, diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 4d04b7ab8..af4d05b1c 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -44,7 +44,7 @@ uv run pytest \ The integration test follows this setup: ```python -config = SessionConfig().with_extension(PlannerConfig(max_rows=3)) +config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) ctx = SessionContext(config) ctx = ctx.with_logical_extension_codec(provider_logical_codec) ctx = ctx.with_physical_extension_codec(provider_physical_codec) @@ -53,7 +53,7 @@ ctx.register_udf(provider_udf) ctx = ctx.with_query_planner(MyQueryPlanner()) ``` -`PlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit. +`MyPlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit. The provider's codec pair is attached to the planner when the derived context is created and is also used to decode the returned physical plan in `datafusion-python`. This planner deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; derived contexts rebind codecs after planner installation, but planner-last order is easier to audit. diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index 15f46bf90..67f5d7460 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -27,11 +27,11 @@ MyPhysicalExtensionCodec, MyTableProvider, ) -from datafusion_ffi_query_planner_example import MyQueryPlanner, PlannerConfig +from datafusion_ffi_query_planner_example import MyPlannerConfig, MyQueryPlanner def configured_context(max_rows: int): - config = SessionConfig().with_extension(PlannerConfig(max_rows=max_rows)) + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows)) logical_codec = MyLogicalExtensionCodec() physical_codec = MyPhysicalExtensionCodec() ctx = SessionContext(config) @@ -81,7 +81,7 @@ def test_spawning_plan_across_three_libraries(): aggregate, and that operator spawns tasks while it runs. This exercises the codecs on a multi-node plan rather than the bare scan the other tests use. """ - config = SessionConfig().with_extension(PlannerConfig(max_rows=100)) + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=100)) config = config.with_target_partitions(4) logical_codec = MyLogicalExtensionCodec() physical_codec = MyPhysicalExtensionCodec() @@ -171,7 +171,7 @@ def test_installed_codecs_outlive_python_exporters(): def test_provider_codecs_can_be_installed_after_planner(): - config = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) planner = MyQueryPlanner() logical_codec = MyLogicalExtensionCodec() physical_codec = MyPhysicalExtensionCodec() @@ -188,7 +188,7 @@ def test_provider_codecs_can_be_installed_after_planner(): def test_query_planner_requires_provider_codec(): - config = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) ctx = SessionContext(config) ctx.register_table("numbers", MyTableProvider(1, 3, 1)) ctx = ctx.with_query_planner(MyQueryPlanner()) diff --git a/examples/datafusion-ffi-query-planner-example/src/config.rs b/examples/datafusion-ffi-query-planner-example/src/config.rs index 801cee9a0..ecfa4b943 100644 --- a/examples/datafusion-ffi-query-planner-example/src/config.rs +++ b/examples/datafusion-ffi-query-planner-example/src/config.rs @@ -28,17 +28,17 @@ use pyo3::types::PyCapsule; #[pyclass( from_py_object, - name = "PlannerConfig", + name = "MyPlannerConfig", module = "datafusion_ffi_query_planner_example", subclass )] #[derive(Clone, Debug)] -pub(crate) struct PlannerConfig { +pub(crate) struct MyPlannerConfig { pub max_rows: usize, } #[pymethods] -impl PlannerConfig { +impl MyPlannerConfig { #[new] #[pyo3(signature = (max_rows=10))] fn new(max_rows: usize) -> Self { @@ -57,17 +57,17 @@ impl PlannerConfig { } } -impl Default for PlannerConfig { +impl Default for MyPlannerConfig { fn default() -> Self { Self { max_rows: 10 } } } -impl ConfigExtension for PlannerConfig { +impl ConfigExtension for MyPlannerConfig { const PREFIX: &'static str = "ffi_query_planner"; } -impl ExtensionOptions for PlannerConfig { +impl ExtensionOptions for MyPlannerConfig { fn as_any(&self) -> &dyn Any { self } @@ -93,7 +93,7 @@ impl ExtensionOptions for PlannerConfig { } } -impl ConfigField for PlannerConfig { +impl ConfigField for MyPlannerConfig { fn visit(&self, visitor: &mut V, _key: &str, _description: &'static str) { self.max_rows.visit( visitor, @@ -106,7 +106,7 @@ impl ConfigField for PlannerConfig { let (key, rem) = key.split_once('.').unwrap_or((key, "")); match key { "max_rows" => self.max_rows.set(rem, value), - _ => config_err!("Config value '{key}' not found on PlannerConfig"), + _ => config_err!("Config value '{key}' not found on MyPlannerConfig"), } } } diff --git a/examples/datafusion-ffi-query-planner-example/src/lib.rs b/examples/datafusion-ffi-query-planner-example/src/lib.rs index 7635c2992..c505c1ce7 100644 --- a/examples/datafusion-ffi-query-planner-example/src/lib.rs +++ b/examples/datafusion-ffi-query-planner-example/src/lib.rs @@ -17,7 +17,7 @@ use pyo3::prelude::*; -use crate::config::PlannerConfig; +use crate::config::MyPlannerConfig; use crate::planner::MyQueryPlanner; mod config; @@ -26,7 +26,7 @@ mod planner; #[pymodule] fn datafusion_ffi_query_planner_example(m: &Bound<'_, PyModule>) -> PyResult<()> { pyo3_log::init(); - m.add_class::()?; + m.add_class::()?; m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index 59791f160..85c0e80d9 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -40,7 +40,7 @@ use datafusion_session::{QueryPlanner, Session}; use pyo3::prelude::*; use pyo3::types::PyCapsule; -use crate::config::PlannerConfig; +use crate::config::MyPlannerConfig; #[derive(Default)] struct PlannerObservations { @@ -82,7 +82,7 @@ fn physical_plan_has_foreign_plan(plan: &Arc) -> bool { .any(|child| physical_plan_has_foreign_plan(child)) } -/// The row limit as the host spells it, where `PlannerConfig` is registered as +/// The row limit as the host spells it, where `MyPlannerConfig` is registered as /// an ordinary config extension under its own `ConfigExtension::PREFIX`. const MAX_ROWS_KEY: &str = "ffi_query_planner.max_rows"; @@ -93,11 +93,11 @@ const MAX_ROWS_KEY: &str = "ffi_query_planner.max_rows"; /// reports the key with both prefixes. const FFI_MAX_ROWS_KEY: &str = "datafusion_ffi.ffi_query_planner.max_rows"; -fn planner_config(session: &dyn Session) -> datafusion::common::Result { +fn planner_config(session: &dyn Session) -> datafusion::common::Result { let options = session.config_options(); // Prefer the raw entry. `local_or_ffi_extension` discards a value it cannot - // parse and hands back `PlannerConfig::default()`, which would quietly turn + // parse and hands back `MyPlannerConfig::default()`, which would quietly turn // a typo into a different row limit instead of reporting it. let config = match options .entries() @@ -114,10 +114,10 @@ fn planner_config(session: &dyn Session) -> datafusion::common::Result options - .local_or_ffi_extension::() + .local_or_ffi_extension::() .unwrap_or_default(), }; @@ -132,9 +132,18 @@ fn planner_config(session: &dyn Session) -> datafusion::common::Result, + /// Keeps the context behind the exported codecs' `TaskContextProvider` + /// alive. See [`MyQueryPlanner::codec_ctx`]. The reference lives here as + /// well as on `MyQueryPlanner` because this is the value the capsule + /// carries, so the provider stays valid even if the Python object that + /// exported it is dropped first. + #[expect( + dead_code, + reason = "strong reference keeping the FFI codecs' weakly held provider alive" + )] + codec_ctx: Arc, /// Planner to hand the work to instead of planning here. /// /// This is how a real planner layers on top of an existing one. The capsule @@ -149,6 +158,16 @@ struct DistributedQueryPlanner { fallback: Option>, } +impl fmt::Debug for DistributedQueryPlanner { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("DistributedQueryPlanner") + .field("observations", &self.observations) + .field("has_fallback", &self.fallback.is_some()) + .finish_non_exhaustive() + } +} + #[async_trait] impl QueryPlanner for DistributedQueryPlanner { async fn create_physical_plan( @@ -201,10 +220,48 @@ impl QueryPlanner for DistributedQueryPlanner { module = "datafusion_ffi_query_planner_example", subclass )] -#[derive(Debug, Default, Clone)] +#[derive(Clone)] pub(crate) struct MyQueryPlanner { observations: Arc, fallback: Option>, + /// Context backing the `TaskContextProvider` handed to the exported + /// codecs. + /// + /// This is *not* the session that arrives at `create_physical_plan`. That + /// one belongs to the host and is what this planner reads config from. The + /// provider here serves the other direction: when the host decodes the + /// plan bytes this library returned, any extension node in them is decoded + /// by a callback back into this library, and that callback needs a + /// `TaskContext` whose registry can resolve *this* library's nodes and + /// functions. A real planner library registers its UDFs and extension + /// types on this context; the example uses the default codecs, so nothing + /// ever calls back. + /// + /// It must be owned rather than built inline: `FFI_TaskContextProvider` + /// downgrades the provider to a `Weak`, so a temporary would already be + /// dropped by the time the capsule is used, and every codec callback would + /// fail with "TaskContextProvider went out of scope over FFI boundary". + codec_ctx: Arc, +} + +impl Default for MyQueryPlanner { + fn default() -> Self { + Self { + observations: Arc::default(), + fallback: None, + codec_ctx: Arc::new(SessionContext::new()), + } + } +} + +impl fmt::Debug for MyQueryPlanner { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("MyQueryPlanner") + .field("observations", &self.observations) + .field("has_fallback", &self.fallback.is_some()) + .finish_non_exhaustive() + } } #[pymethods] @@ -260,10 +317,11 @@ impl MyQueryPlanner { ) -> PyResult> { let planner: Arc = Arc::new(DistributedQueryPlanner { observations: Arc::clone(&self.observations), + codec_ctx: Arc::clone(&self.codec_ctx), fallback: self.fallback.clone(), }); let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider = Arc::new(SessionContext::new()) as Arc; + let ctx_provider = Arc::clone(&self.codec_ctx) as Arc; let ffi = FFI_QueryPlanner::new( planner, Some(runtime), From 5bffd5ed5a70b252d7814b81824633b2598c2837 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 26 Aug 2026 17:30:19 -0400 Subject: [PATCH 12/30] test: cover which session a foreign codec decodes against The example codecs restore objects from a process-local token registry and never read the `TaskContext` their FFI decode callbacks are handed, so which session that context belongs to was untestable. The token path ignores the registry entirely, which is why an empty `SessionContext::new()` has served as the exported provider without anyone noticing. Both codecs now accept `require_udf_on_decode`. When set, every decode call resolves that scalar function out of the task context it was given and fails with the session id if it is absent, which makes the answer observable. Each codec registers a marker function on the context it exports, so a name owned by the codec's library and a name owned by the host can be told apart. Four tests use it. The two library-local cases pass: a foreign codec resolves against the session its own library supplied. The two host-registered cases are `xfail(strict=True)`, because a function registered on the host with `register_udf` is not visible to a foreign codec's decode callback at all. A fifth pins the current error so the failure mode stays legible. Strict xfail means the pair will announce itself if the upstream design changes. Document the rule this establishes, and correct the surrounding section: `with_query_planner` rebuilds a foreign planner against the session that will run the query, so the provider a planner library supplies is replaced on that path. Codecs installed through `with_logical_extension_codec` and `with_physical_extension_codec` keep the provider their own library exported, which is the case these tests exercise. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/contributor-guide/ffi.md | 25 ++++ examples/datafusion-ffi-example/src/lib.rs | 1 + .../src/logical_extension_codec.rs | 32 ++++- .../src/physical_extension_codec.rs | 32 ++++- .../src/required_udf.rs | 127 ++++++++++++++++++ .../_test_three_library_query_planner.py | 102 ++++++++++++++ 6 files changed, 315 insertions(+), 4 deletions(-) create mode 100644 examples/datafusion-ffi-example/src/required_udf.rs diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 8a8b65a8d..e6706d694 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -283,6 +283,31 @@ something that lives at least as long as the exported planner — the example st on both the Python-facing planner object and the `QueryPlanner` the capsule carries, so the capsule keeps working even if the Python object is dropped first. +One wrinkle specific to installing a planner here: `with_query_planner` rebuilds the +foreign planner against the context that will run the query, so the codecs the planner +was exported with — and the provider behind them — are replaced by the session's own. +A planner library still has to supply a provider to construct the capsule, but in this +path it is not the one consulted. A codec installed with +`with_logical_extension_codec` or `with_physical_extension_codec` keeps the provider its +own library exported. + +### A codec decodes against its own library's registry + +Follows from the above, and it is the part most likely to surprise: a decode callback +resolves names against the session the *exporting library* supplied, not the host +session running the query. A scalar function registered on the host with +`ctx.register_udf(...)` is therefore not visible to a foreign codec decoding a node that +references it by name. Register anything a codec has to resolve on the context that +codec exports. + +This is covered directly. Both example codecs accept `require_udf_on_decode`, which +makes every decode call resolve a named scalar function out of the `TaskContext` it was +handed, and +`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py` +asserts both halves: a function registered on the codec's own context resolves, and one +registered only on the host does not. The host-registered cases are `xfail(strict=True)` +so they will announce themselves if the underlying design changes upstream. + ### What a derived context shares `with_query_planner`, `with_logical_extension_codec`, `with_physical_extension_codec`, diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index eccf7b81a..3d00fdb3e 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -35,6 +35,7 @@ pub(crate) mod config; pub(crate) mod logical_extension_codec; pub(crate) mod physical_extension_codec; pub(crate) mod physical_optimizer; +pub(crate) mod required_udf; pub(crate) mod scalar_udf; pub(crate) mod table_function; pub(crate) mod table_provider; diff --git a/examples/datafusion-ffi-example/src/logical_extension_codec.rs b/examples/datafusion-ffi-example/src/logical_extension_codec.rs index 0474a8d39..987ea2dbb 100644 --- a/examples/datafusion-ffi-example/src/logical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/logical_extension_codec.rs @@ -33,6 +33,8 @@ use datafusion_python_util::get_tokio_runtime; use pyo3::prelude::*; use pyo3::types::PyCapsule; +use crate::required_udf::{new_codec_context, resolve_required_udf}; + const TABLE_PROVIDER_TOKEN: &[u8] = b"DFPYEXTP"; static NEXT_TABLE_PROVIDER_ID: AtomicU64 = AtomicU64::new(1); static TABLE_PROVIDERS: OnceLock>>> = OnceLock::new(); @@ -52,6 +54,7 @@ pub(crate) struct CallCounters { pub decode_udf: AtomicUsize, pub encode_table_provider: AtomicUsize, pub decode_table_provider: AtomicUsize, + pub task_ctx_udf_resolutions: AtomicUsize, } /// Example codec for objects owned by this extension library. @@ -63,6 +66,9 @@ pub(crate) struct CallCounters { struct CountingLogicalExtensionCodec { inner: DefaultLogicalExtensionCodec, counters: Arc, + /// Scalar function every table-provider decode must resolve from the + /// `TaskContext` it is handed. See [`crate::required_udf`]. + required_udf: Option, // The FFI task-context handle is weak. Retain its provider for as long as // this codec can be called, even if Python drops the exporter object. _ctx_provider: Arc, @@ -99,6 +105,11 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec { schema: SchemaRef, ctx: &TaskContext, ) -> Result> { + resolve_required_udf( + self.required_udf.as_deref(), + ctx, + &self.counters.task_ctx_udf_resolutions, + )?; if let Some(id) = token_id(buf, TABLE_PROVIDER_TOKEN) { self.counters .decode_table_provider @@ -159,19 +170,35 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec { #[derive(Clone)] pub(crate) struct MyLogicalExtensionCodec { counters: Arc, + required_udf: Option, ctx_provider: Arc, } #[pymethods] impl MyLogicalExtensionCodec { + /// Build the codec. + /// + /// `require_udf_on_decode` names a scalar function that every table + /// provider decode must find in the `TaskContext` it is handed. Leave it + /// unset for the ordinary behaviour; set it to observe *which* session's + /// registry the FFI decode callback actually receives. #[new] - fn new() -> Self { + #[pyo3(signature = (require_udf_on_decode=None))] + fn new(require_udf_on_decode: Option) -> Self { Self { counters: Arc::new(CallCounters::default()), - ctx_provider: Arc::new(SessionContext::new()), + required_udf: require_udf_on_decode, + ctx_provider: new_codec_context(), } } + /// Number of decode calls that resolved `require_udf_on_decode`. + fn task_context_udf_resolutions(&self) -> usize { + self.counters + .task_ctx_udf_resolutions + .load(Ordering::SeqCst) + } + fn encode_udf_calls(&self) -> usize { self.counters.encode_udf.load(Ordering::SeqCst) } @@ -195,6 +222,7 @@ impl MyLogicalExtensionCodec { let inner: Arc = Arc::new(CountingLogicalExtensionCodec { inner: DefaultLogicalExtensionCodec {}, counters: Arc::clone(&self.counters), + required_udf: self.required_udf.clone(), _ctx_provider: Arc::clone(&self.ctx_provider), }); diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs index d1b9ed63b..801d47515 100644 --- a/examples/datafusion-ffi-example/src/physical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -35,6 +35,8 @@ use datafusion_python_util::get_tokio_runtime; use pyo3::prelude::*; use pyo3::types::PyCapsule; +use crate::required_udf::{new_codec_context, resolve_required_udf}; + const EXECUTION_PLAN_TOKEN: &[u8] = b"DFPYEXEP"; static NEXT_EXECUTION_PLAN_ID: AtomicU64 = AtomicU64::new(1); static EXECUTION_PLANS: OnceLock>>> = OnceLock::new(); @@ -54,6 +56,7 @@ pub(crate) struct PhysicalCallCounters { pub decode_udf: AtomicUsize, pub encode_execution_plan: AtomicUsize, pub decode_execution_plan: AtomicUsize, + pub task_ctx_udf_resolutions: AtomicUsize, } /// Physical companion to the logical example codec. @@ -64,6 +67,9 @@ pub(crate) struct PhysicalCallCounters { struct CountingPhysicalExtensionCodec { inner: DefaultPhysicalExtensionCodec, counters: Arc, + /// Scalar function every decode call must resolve from the `TaskContext` + /// it is handed. See [`crate::required_udf`]. + required_udf: Option, // The FFI task-context handle is weak. Keep its provider alive with the // codec rather than relying on the lifetime of the Python exporter. _ctx_provider: Arc, @@ -87,6 +93,11 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { ctx: &TaskContext, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { + resolve_required_udf( + self.required_udf.as_deref(), + ctx, + &self.counters.task_ctx_udf_resolutions, + )?; if let Some(id) = token_id(buf) { self.counters .decode_execution_plan @@ -149,19 +160,35 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { #[derive(Clone)] pub(crate) struct MyPhysicalExtensionCodec { counters: Arc, + required_udf: Option, ctx_provider: Arc, } #[pymethods] impl MyPhysicalExtensionCodec { + /// Build the codec. + /// + /// `require_udf_on_decode` names a scalar function that every decode call + /// must find in the `TaskContext` it is handed. Leave it unset for the + /// ordinary behaviour; set it to observe *which* session's registry the + /// FFI decode callback actually receives. #[new] - fn new() -> Self { + #[pyo3(signature = (require_udf_on_decode=None))] + fn new(require_udf_on_decode: Option) -> Self { Self { counters: Arc::new(PhysicalCallCounters::default()), - ctx_provider: Arc::new(SessionContext::new()), + required_udf: require_udf_on_decode, + ctx_provider: new_codec_context(), } } + /// Number of decode calls that resolved `require_udf_on_decode`. + fn task_context_udf_resolutions(&self) -> usize { + self.counters + .task_ctx_udf_resolutions + .load(Ordering::SeqCst) + } + fn encode_udf_calls(&self) -> usize { self.counters.encode_udf.load(Ordering::SeqCst) } @@ -186,6 +213,7 @@ impl MyPhysicalExtensionCodec { Arc::new(CountingPhysicalExtensionCodec { inner: DefaultPhysicalExtensionCodec {}, counters: Arc::clone(&self.counters), + required_udf: self.required_udf.clone(), _ctx_provider: Arc::clone(&self.ctx_provider), }); diff --git a/examples/datafusion-ffi-example/src/required_udf.rs b/examples/datafusion-ffi-example/src/required_udf.rs new file mode 100644 index 000000000..00e6bd5ee --- /dev/null +++ b/examples/datafusion-ffi-example/src/required_udf.rs @@ -0,0 +1,127 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Support for exercising the `TaskContext` a codec is handed when it decodes. +//! +//! A codec exported over FFI is built with a `TaskContextProvider`, and the +//! decode callbacks in `datafusion-ffi` resolve that provider to a +//! `TaskContext` before calling into the codec. Nothing in the example codecs +//! read anything out of that context, so which session it belongs to was +//! untestable: the token registries they use are keyed by an integer and +//! ignore the registry entirely. +//! +//! The codecs can now be asked to resolve a named scalar function from the +//! context they are given on every decode, which makes the answer observable. +//! Two names matter: +//! +//! - [`LIBRARY_LOCAL_UDF_NAME`] is registered by [`new_codec_context`] on the +//! context this library owns and hands to the FFI codec, so a decode +//! callback resolves it. +//! - A function registered only on the *host* `SessionContext` does not +//! resolve, because the decode callback never sees the host's registry. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use arrow_schema::DataType; +use datafusion::execution::TaskContext; +use datafusion::prelude::SessionContext; +use datafusion_common::error::Result as DataFusionResult; +use datafusion_common::plan_err; +use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, + Volatility, +}; + +/// Name of the scalar function registered on the context this library owns. +pub(crate) const LIBRARY_LOCAL_UDF_NAME: &str = "library_local_marker"; + +/// Placeholder scalar function used only as a registry entry. +/// +/// Decode callbacks look it up by name to prove which `TaskContext` they were +/// handed. It is never invoked, so the body is unreachable in practice. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct LibraryLocalUDF { + signature: Signature, +} + +impl LibraryLocalUDF { + fn new() -> Self { + Self { + signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for LibraryLocalUDF { + fn name(&self) -> &str { + LIBRARY_LOCAL_UDF_NAME + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> DataFusionResult { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> DataFusionResult { + plan_err!("{LIBRARY_LOCAL_UDF_NAME} exists only as a registry entry and cannot be invoked") + } +} + +/// Builds the context a codec keeps for its own FFI `TaskContextProvider`. +/// +/// This is the context every extension library has to conjure in order to +/// export a codec, and it is a plain empty session apart from the marker +/// function. The host's registrations are not in it, which is the point the +/// codec tests make observable. +pub(crate) fn new_codec_context() -> Arc { + let ctx = SessionContext::new(); + ctx.register_udf(ScalarUDF::from(LibraryLocalUDF::new())); + Arc::new(ctx) +} + +/// Resolves `required` against `ctx`, the context the decode callback was given. +/// +/// `Ok(())` when nothing was requested. Otherwise the name must be present in +/// the context's scalar function registry, and `resolutions` counts each +/// success so a test can tell a resolved lookup from a skipped one. +pub(crate) fn resolve_required_udf( + required: Option<&str>, + ctx: &TaskContext, + resolutions: &AtomicUsize, +) -> DataFusionResult<()> { + let Some(name) = required else { + return Ok(()); + }; + + if ctx.scalar_functions().contains_key(name) { + resolutions.fetch_add(1, Ordering::SeqCst); + return Ok(()); + } + + // A fresh SessionContext still carries every built-in, so report the count + // rather than the whole registry. + plan_err!( + "datafusion-ffi-example: decode could not resolve scalar function '{name}' \ + in the task context it was handed (session '{}', {} scalar functions registered)", + ctx.session_id(), + ctx.scalar_functions().len() + ) +} diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index 67f5d7460..446b74f19 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -42,6 +42,108 @@ def configured_context(max_rows: int): return ctx, logical_codec, physical_codec +LIBRARY_LOCAL_UDF = "library_local_marker" +"""Scalar function registered on the context each example codec owns.""" + +HOST_ONLY_UDF = "my_custom_is_null" +"""Scalar function ``configured_context`` registers on the host session only.""" + + +def probe_context( + *, + logical_requires: str | None = None, + physical_requires: str | None = None, + max_rows: int = 3, +): + """Three-library context whose codecs read the task context they are given. + + ``require_udf_on_decode`` makes each codec resolve a scalar function from + the ``TaskContext`` handed to its FFI decode callback, which is otherwise + unobservable: the example codecs restore objects from a token registry and + never look at the registry they are passed. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows)) + logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=logical_requires) + physical_codec = MyPhysicalExtensionCodec(require_udf_on_decode=physical_requires) + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(physical_codec) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + ctx = ctx.with_query_planner(MyQueryPlanner()) + return ctx, logical_codec, physical_codec + + +def test_logical_codec_resolves_its_own_libraries_udf(): + """``try_decode_table_provider`` sees the registry of the codec's own session.""" + ctx, logical_codec, _physical_codec = probe_context( + logical_requires=LIBRARY_LOCAL_UDF + ) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert logical_codec.table_provider_decode_calls() > 0 + assert logical_codec.task_context_udf_resolutions() > 0 + + +def test_physical_codec_resolves_its_own_libraries_udf(): + """``try_decode`` sees the registry of the codec's own session.""" + ctx, _logical_codec, physical_codec = probe_context( + physical_requires=LIBRARY_LOCAL_UDF + ) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert physical_codec.execution_plan_decode_calls() > 0 + assert physical_codec.task_context_udf_resolutions() > 0 + + +@pytest.mark.xfail( + strict=True, + reason=( + "A codec exported over FFI is built with its own TaskContextProvider, so " + "its decode callbacks resolve names against the exporting library's " + "session rather than the host session that is running the query. A " + "function registered only on the host is therefore invisible to them." + ), +) +def test_logical_codec_resolves_a_host_registered_udf(): + ctx, logical_codec, _physical_codec = probe_context(logical_requires=HOST_ONLY_UDF) + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert logical_codec.task_context_udf_resolutions() > 0 + + +@pytest.mark.xfail( + strict=True, + reason=( + "Same as the logical case: the physical codec's decode callback is " + "handed the exporting library's task context, which does not know about " + "functions registered on the host session." + ), +) +def test_physical_codec_resolves_a_host_registered_udf(): + ctx, _logical_codec, physical_codec = probe_context(physical_requires=HOST_ONLY_UDF) + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert physical_codec.task_context_udf_resolutions() > 0 + + +def test_decode_task_context_is_not_the_host_session(): + """Pin the current behaviour so the failure mode above stays legible. + + The host registers ``my_custom_is_null`` and can use it in a query, but the + codec's decode callback cannot see it, and the error names the session it + was given instead. + """ + ctx, _logical_codec, _physical_codec = probe_context(logical_requires=HOST_ONLY_UDF) + + with pytest.raises( + Exception, match=rf"could not resolve scalar function '{HOST_ONLY_UDF}'" + ): + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + + @pytest.mark.parametrize("raw_capsule", [False, True]) def test_three_library_query_planner(raw_capsule: bool): """Host, provider, and planner exchange a real non-empty plan over FFI.""" From b0800feaf3121c2ab48644da487eb9a9917c7cc2 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 26 Aug 2026 21:19:15 -0400 Subject: [PATCH 13/30] feat: pass the session to the capsule getters that need it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `FFI_QueryPlanner::new` and `FFI_{Logical,Physical}ExtensionCodec::new` ask an extension library for a `TaskContextProvider`, and a planner for two codecs on top of that. A library has none of those. Both examples answered with `Arc::new(SessionContext::new())`, an empty session that resolves nothing, held weakly by `FFI_TaskContextProvider` and therefore also a lifetime hazard. The table provider protocol already solved this: the host calls `__datafusion_table_provider__(session)` and the library takes what it needs off the session. Do the same for the other three getters. `__datafusion_query_planner__`, `__datafusion_logical_extension_codec__`, and `__datafusion_physical_extension_codec__` now receive the `SessionContext` they are being installed on. A codec takes the task context provider from it; a planner takes both codecs and uses `new_with_ffi_codecs`, which needs no provider at all. Neither example constructs a `SessionContext` any more. Decode callbacks consequently resolve against the session running the query. The two `xfail(strict=True)` tests from the previous commit now pass unmodified: a scalar function registered on the host with `register_udf` is visible inside a decode callback executing in another library, for both the logical and physical codec. A negative control keeps the check honest, and a further test covers a function registered after the codec was installed, since the provider is a live handle rather than a snapshot. `PySessionContext` gains an `ancestors` list. A foreign codec is built against the session current at the time it is installed and holds it weakly, so installing a foreign planner afterwards — which forks — would strand the codec once the Python name is rebound. The keep-alive lives on `PySessionContext` rather than on the codec because nothing reachable from a `SessionContext` reaches a `PySessionContext`, so it cannot close a cycle. What it does not paper over is the fork itself: a function registered after the fork is not visible to a codec bound to the session before it, which is the existing derived-context caveat seen from the codec's side, and is covered by a test. `SessionContext` accepts and ignores the argument on all three getters, so a session satisfies the same protocol a library implements and `ctx.__datafusion_query_planner__()` keeps working for the delegation pattern. Calling a stale getter that takes no session now reports an incompatible-library error naming the method, matching what `table_provider_from_pycapsule` does. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/catalog.rs | 2 +- crates/core/src/context.rs | 89 ++++++++++++--- crates/util/src/lib.rs | 98 ++++++++++++++--- docs/source/contributor-guide/ffi.md | 103 +++++++++--------- examples/datafusion-ffi-example/README.md | 2 +- .../src/catalog_provider.rs | 6 +- .../src/logical_extension_codec.rs | 23 ++-- .../src/physical_extension_codec.rs | 22 ++-- .../src/required_udf.rs | 77 ++----------- .../src/table_function.rs | 2 +- .../src/table_provider.rs | 2 +- .../src/table_provider_factory.rs | 2 +- .../_test_three_library_query_planner.py | 100 +++++++++-------- .../src/planner.rs | 92 +++------------- python/datafusion/context.py | 36 ++++-- python/datafusion/user_defined.py | 21 +++- 16 files changed, 359 insertions(+), 318 deletions(-) diff --git a/crates/core/src/catalog.rs b/crates/core/src/catalog.rs index 8ad49b098..ff170c257 100644 --- a/crates/core/src/catalog.rs +++ b/crates/core/src/catalog.rs @@ -689,7 +689,7 @@ fn extract_logical_extension_codec( Some(obj) => obj, None => PySessionContext::global_ctx()?.into_bound_py_any(py)?, }; - ffi_logical_codec_from_pycapsule(obj).map(Arc::new) + ffi_logical_codec_from_pycapsule(obj, None).map(Arc::new) } pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 5b9f09a06..db5250d3a 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -60,9 +60,8 @@ use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_python_util::{ create_logical_extension_capsule, create_physical_extension_capsule, create_query_planner_capsule, ffi_logical_codec_from_pycapsule, - ffi_query_planner_from_pycapsule, get_global_ctx, get_tokio_runtime, - physical_codec_from_pycapsule, physical_optimizer_rule_from_pycapsule, spawn_future, - wait_for_future, + ffi_physical_codec_from_pycapsule, ffi_query_planner_from_pycapsule, get_global_ctx, + get_tokio_runtime, physical_optimizer_rule_from_pycapsule, spawn_future, wait_for_future, }; use object_store::ObjectStore; use pyo3::IntoPyObjectExt; @@ -374,6 +373,16 @@ pub struct PySessionContext { pub ctx: Arc, logical_codec: Arc, physical_codec: Arc, + /// Sessions this one was derived from, held only to keep them alive. + /// + /// A foreign codec installed here is built against the session current at + /// the time of the call, and `FFI_TaskContextProvider` holds that session + /// weakly. Installing a foreign query planner afterwards forks, and + /// rebinding the Python name would otherwise drop the session the codec + /// still points at. Retaining ancestors here rather than on the codec + /// avoids a cycle: nothing reachable from a `SessionContext` reaches back + /// to a `PySessionContext`. + ancestors: Vec>, } #[pymethods] @@ -406,6 +415,7 @@ impl PySessionContext { ctx, logical_codec: Arc::new(PythonLogicalCodec::default()), physical_codec: Arc::new(PythonPhysicalCodec::default()), + ancestors: Vec::new(), }) } @@ -414,6 +424,7 @@ impl PySessionContext { ctx: Arc::new(self.ctx.as_ref().clone().enable_url_table()), logical_codec: Arc::clone(&self.logical_codec), physical_codec: Arc::clone(&self.physical_codec), + ancestors: self.ancestors_including_self(), }) } @@ -425,6 +436,7 @@ impl PySessionContext { ctx, logical_codec: Arc::new(PythonLogicalCodec::default()), physical_codec: Arc::new(PythonPhysicalCodec::default()), + ancestors: Vec::new(), }) } @@ -1213,15 +1225,20 @@ impl PySessionContext { Ok(()) } - pub fn with_query_planner(&self, planner: Bound<'_, PyAny>) -> PyDataFusionResult { - let planner = ffi_query_planner_from_pycapsule(&planner)?; + pub fn with_query_planner( + slf: &Bound<'_, Self>, + planner: Bound<'_, PyAny>, + ) -> PyDataFusionResult { + let planner = ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))?; + let this = slf.borrow(); let ctx = - self.ctx_with_rebound_planner(&self.logical_codec, &self.physical_codec, Some(planner)); + this.ctx_with_rebound_planner(&this.logical_codec, &this.physical_codec, Some(planner)); Ok(Self { + ancestors: this.ancestors_for(&ctx), ctx, - logical_codec: Arc::clone(&self.logical_codec), - physical_codec: Arc::clone(&self.physical_codec), + logical_codec: Arc::clone(&this.logical_codec), + physical_codec: Arc::clone(&this.physical_codec), }) } @@ -1391,17 +1408,27 @@ impl PySessionContext { PyCapsule::new_with_value(py, ffi_ctx_provider, cr"datafusion_task_context_provider") } + /// `session` exists so this matches the protocol an extension library + /// implements, where the argument is how the library reaches the session + /// it is being installed on. A session already is one, so it is ignored. + #[pyo3(signature = (session=None))] pub fn __datafusion_logical_extension_codec__<'py>( &self, py: Python<'py>, + session: Option>, ) -> PyResult> { + let _ = session; create_logical_extension_capsule(py, &self.exported_ffi_logical_codec()) } + /// See [`Self::__datafusion_logical_extension_codec__`] for `session`. + #[pyo3(signature = (session=None))] pub fn __datafusion_query_planner__<'py>( &self, py: Python<'py>, + session: Option>, ) -> PyResult> { + let _ = session; // An already-foreign planner is re-exported as its original handle // rather than gaining another layer, because `new_with_ffi_codecs` // unwraps a `ForeignQueryPlanner`. It still adopts the codecs supplied @@ -1417,39 +1444,48 @@ impl PySessionContext { } pub fn with_logical_extension_codec<'py>( - &self, + slf: &Bound<'py, Self>, codec: Bound<'py, PyAny>, ) -> PyDataFusionResult { - let inner_ffi = ffi_logical_codec_from_pycapsule(codec)?; + let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?; let inner: Arc = (&inner_ffi).into(); let logical_codec = Arc::new(PythonLogicalCodec::new(inner)); - let physical_codec = Arc::clone(&self.physical_codec); - let ctx = self.ctx_with_rebound_planner(&logical_codec, &physical_codec, None); + let this = slf.borrow(); + let physical_codec = Arc::clone(&this.physical_codec); + let ctx = this.ctx_with_rebound_planner(&logical_codec, &physical_codec, None); Ok(Self { + ancestors: this.ancestors_for(&ctx), ctx, logical_codec, physical_codec, }) } + /// See [`Self::__datafusion_logical_extension_codec__`] for `session`. + #[pyo3(signature = (session=None))] pub fn __datafusion_physical_extension_codec__<'py>( &self, py: Python<'py>, + session: Option>, ) -> PyResult> { + let _ = session; create_physical_extension_capsule(py, &self.exported_ffi_physical_codec()) } pub fn with_physical_extension_codec<'py>( - &self, + slf: &Bound<'py, Self>, codec: Bound<'py, PyAny>, ) -> PyDataFusionResult { - let inner = physical_codec_from_pycapsule(&codec)?; + let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?; + let inner: Arc = (&inner_ffi).into(); let physical_codec = Arc::new(PythonPhysicalCodec::new(inner)); - let logical_codec = Arc::clone(&self.logical_codec); - let ctx = self.ctx_with_rebound_planner(&logical_codec, &physical_codec, None); + let this = slf.borrow(); + let logical_codec = Arc::clone(&this.logical_codec); + let ctx = this.ctx_with_rebound_planner(&logical_codec, &physical_codec, None); Ok(Self { + ancestors: this.ancestors_for(&ctx), ctx, logical_codec, physical_codec, @@ -1467,6 +1503,7 @@ impl PySessionContext { ); let ctx = self.ctx_with_rebound_planner(&logical_codec, &physical_codec, None); Self { + ancestors: self.ancestors_for(&ctx), ctx, logical_codec, physical_codec, @@ -1475,6 +1512,25 @@ impl PySessionContext { } impl PySessionContext { + /// Keep-alive list for a context derived from this one. + /// + /// Sharing the same `SessionContext` carries the list through unchanged. + /// A fork adds this session, because a foreign codec installed here holds + /// it weakly and would otherwise dangle once the Python name is rebound. + fn ancestors_for(&self, derived: &Arc) -> Vec> { + if Arc::ptr_eq(derived, &self.ctx) { + self.ancestors.clone() + } else { + self.ancestors_including_self() + } + } + + fn ancestors_including_self(&self) -> Vec> { + let mut ancestors = self.ancestors.clone(); + ancestors.push(Arc::clone(&self.ctx)); + ancestors + } + /// Return the [`SessionContext`] a derived `PySessionContext` should use, /// binding a foreign query planner to `logical_codec` and `physical_codec`. /// @@ -1760,6 +1816,7 @@ impl From for PySessionContext { ctx: Arc::new(ctx), logical_codec: Arc::new(PythonLogicalCodec::default()), physical_codec: Arc::new(PythonPhysicalCodec::default()), + ancestors: Vec::new(), } } } diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index 7375a034c..0fa908f5f 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -215,14 +215,47 @@ pub fn create_logical_extension_capsule<'py>( PyCapsule::new_with_value(py, codec, cr"datafusion_logical_extension_codec") } -pub fn ffi_logical_codec_from_pycapsule(obj: Bound) -> PyResult { - let attr_name = "__datafusion_logical_extension_codec__"; - let capsule = if obj.hasattr(attr_name)? { - obj.getattr(attr_name)?.call0()? - } else { - obj +/// Calls `obj.____(session)`, or `obj.____()` when no +/// session is supplied. +/// +/// The session is how an exporting library obtains the codecs and task context +/// of the session it is being installed on, instead of inventing one of its +/// own. `None` is for the reverse direction, where `obj` *is* a session and is +/// being asked for what it holds. +fn call_capsule_getter<'py>( + obj: Bound<'py, PyAny>, + attr_name: &str, + session: Option<&Bound<'py, PyAny>>, +) -> PyResult> { + if !obj.hasattr(attr_name)? { + return Ok(obj); + } + + let getter = obj.getattr(attr_name)?; + let result = match session { + Some(session) => getter.call1((session,)), + None => getter.call0(), }; + result.map_err(|err| { + let py = obj.py(); + if session.is_some() && err.get_type(py).is(PyType::new::(py)) { + PyImportError::new_err(format!( + "Incompatible libraries. `{attr_name}` must accept the SessionContext it \ + is being installed on. Upgrade the library providing this object." + )) + } else { + err + } + }) +} + +pub fn ffi_logical_codec_from_pycapsule<'py>( + obj: Bound<'py, PyAny>, + session: Option<&Bound<'py, PyAny>>, +) -> PyResult { + let capsule = call_capsule_getter(obj, "__datafusion_logical_extension_codec__", session)?; + let capsule = capsule.cast::()?; let data: NonNull = capsule .pointer_checked(Some(c"datafusion_logical_extension_codec"))? @@ -232,6 +265,47 @@ pub fn ffi_logical_codec_from_pycapsule(obj: Bound) -> PyResult( + obj: Bound<'py, PyAny>, + session: Option<&Bound<'py, PyAny>>, +) -> PyResult { + let capsule = call_capsule_getter(obj, "__datafusion_physical_extension_codec__", session)?; + + let capsule = capsule.cast::()?; + validate_pycapsule(capsule, "datafusion_physical_extension_codec")?; + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_physical_extension_codec"))? + .cast(); + let codec = unsafe { data.as_ref() }; + + Ok(codec.clone()) +} + +/// Extracts the `FFI_TaskContextProvider` a session exposes. +/// +/// An extension library exporting a codec needs one for the decode callbacks +/// its codec will receive. Taking the host's means those callbacks resolve +/// names against the session that is actually running the query, and removes +/// any need for the library to construct a `SessionContext` of its own. +pub fn ffi_task_context_provider_from_pycapsule( + session: &Bound, +) -> PyResult { + let capsule = call_capsule_getter( + session.clone(), + "__datafusion_task_context_provider__", + None, + )?; + + let capsule = capsule.cast::()?; + validate_pycapsule(capsule, "datafusion_task_context_provider")?; + let data: NonNull = capsule + .pointer_checked(Some(c"datafusion_task_context_provider"))? + .cast(); + let provider = unsafe { data.as_ref() }; + + Ok(provider.clone()) +} + pub fn create_query_planner_capsule<'py>( py: Python<'py>, planner: &FFI_QueryPlanner, @@ -239,13 +313,11 @@ pub fn create_query_planner_capsule<'py>( PyCapsule::new_with_value(py, planner.clone(), cr"datafusion_query_planner") } -pub fn ffi_query_planner_from_pycapsule(obj: &Bound) -> PyResult { - let attr_name = "__datafusion_query_planner__"; - let capsule = if obj.hasattr(attr_name)? { - obj.getattr(attr_name)?.call0()? - } else { - obj.clone() - }; +pub fn ffi_query_planner_from_pycapsule<'py>( + obj: &Bound<'py, PyAny>, + session: Option<&Bound<'py, PyAny>>, +) -> PyResult { + let capsule = call_capsule_getter(obj.clone(), "__datafusion_query_planner__", session)?; let capsule = capsule.cast::()?; validate_pycapsule(capsule, "datafusion_query_planner")?; diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index e6706d694..886341b89 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -257,56 +257,59 @@ The current FFI logical codec supports providers and UDFs but not arbitrary cust `LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and local build commands. -### The task context a planner exports is not the session it plans against - -`FFI_QueryPlanner::new` takes a `TaskContextProvider` alongside the two codecs. It is -easy to mistake this for the session the planner will be asked to plan against, but the -two are unrelated and serve opposite directions of the exchange: - -- The `&dyn Session` passed to `QueryPlanner::create_physical_plan` belongs to the - host. It arrives as a `ForeignSession` when the call crossed the boundary, and it is - what the planner reads configuration and catalogs from. -- The `TaskContextProvider` given at export time backs the *exported codecs*. When the - host decodes the plan bytes the planner returned, any extension node in them is - decoded by a callback back into the planner's own library, and that callback needs a - `TaskContext` whose registry can resolve that library's own nodes and functions. - -So a planner library should hand over a context it owns, and register its own UDFs and -extension types on it. It cannot use the host's session for this even if it wanted to: -the codecs are built when the capsule is exported, long before any session shows up. - -That context must be kept alive by the exporter. `FFI_TaskContextProvider` downgrades -it to a `Weak`, so a provider constructed inline in the capsule getter is already -dropped by the time the capsule is used, and every codec callback fails with -`TaskContextProvider went out of scope over FFI boundary`. Hold the reference on -something that lives at least as long as the exported planner — the example stores it -on both the Python-facing planner object and the `QueryPlanner` the capsule carries, so -the capsule keeps working even if the Python object is dropped first. - -One wrinkle specific to installing a planner here: `with_query_planner` rebuilds the -foreign planner against the context that will run the query, so the codecs the planner -was exported with — and the provider behind them — are replaced by the session's own. -A planner library still has to supply a provider to construct the capsule, but in this -path it is not the one consulted. A codec installed with -`with_logical_extension_codec` or `with_physical_extension_codec` keeps the provider its -own library exported. - -### A codec decodes against its own library's registry - -Follows from the above, and it is the part most likely to surprise: a decode callback -resolves names against the session the *exporting library* supplied, not the host -session running the query. A scalar function registered on the host with -`ctx.register_udf(...)` is therefore not visible to a foreign codec decoding a node that -references it by name. Register anything a codec has to resolve on the context that -codec exports. - -This is covered directly. Both example codecs accept `require_udf_on_decode`, which -makes every decode call resolve a named scalar function out of the `TaskContext` it was -handed, and -`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py` -asserts both halves: a function registered on the codec's own context resolves, and one -registered only on the host does not. The host-registered cases are `xfail(strict=True)` -so they will announce themselves if the underlying design changes upstream. +### Capsule getters receive the session they are installed on + +`__datafusion_query_planner__`, `__datafusion_logical_extension_codec__`, and +`__datafusion_physical_extension_codec__` all take the `SessionContext` the object is +being installed on, the same way `__datafusion_table_provider__` does: + +```rust +fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, +) -> PyResult> { + let runtime = get_tokio_runtime().handle().clone(); + let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; + let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), ctx_provider); + PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") +} +``` + +This exists because the FFI constructors need things an extension library does not +have. `FFI_{Logical,Physical}ExtensionCodec::new` needs a `TaskContextProvider` for the +decode callbacks the codec will receive, and `FFI_QueryPlanner::new` needs both codecs +on top of that. Taking them from the session is what keeps a library from constructing +a `SessionContext` purely to satisfy a parameter — an empty one resolves nothing, and +`FFI_TaskContextProvider` holds it weakly, so a context built inline in the getter is +already dropped by the time the capsule is used. + +A planner uses `FFI_QueryPlanner::new_with_ffi_codecs` with the two codecs it takes off +the session, and never touches a provider directly. That also matches what installation +does anyway: `with_query_planner` rebinds a foreign planner to the codecs of the session +that will run the query. + +`SessionContext` accepts the argument on all three getters and ignores it, so a session +satisfies the same protocol an extension library implements. When you export the current +planner to wrap it, `ctx.__datafusion_query_planner__()` and +`ctx.__datafusion_query_planner__(ctx)` are both fine. + +### A codec decodes against the session it was installed on + +Because the provider comes from the host, a decode callback running inside an extension +library resolves names against the session that is running the query. A function +registered with `ctx.register_udf(...)` is visible to a foreign codec decoding a node +that references it by name, and the handle is live rather than a snapshot, so a +registration made after the codec is installed is visible too. + +The one boundary is the fork. Installing a foreign query planner forks the session, and +a codec installed beforehand keeps pointing at the session it was built against, so a +function registered after that point is not visible to it. This is the same rule as the +caveat below, seen from the codec's side: register before deriving. Both halves are +covered in +`examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py`, +where the example codecs take a `require_udf_on_decode` name and resolve it out of the +task context they are handed. ### What a derived context shares diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md index f2945251f..f283d1610 100644 --- a/examples/datafusion-ffi-example/README.md +++ b/examples/datafusion-ffi-example/README.md @@ -33,7 +33,7 @@ Separate shared libraries guarantee distinct DataFusion library markers. This ca `MyLogicalExtensionCodec` serializes this example's in-memory table providers, and `MyPhysicalExtensionCodec` serializes provider-owned memory scans and opaque FFI wrappers around them. Both use documented, process-local, one-shot token registries. The registries make ownership and callback routing visible without pretending to be a portable format. They assume trusted in-process payloads and consume each token during decoding. A production provider should instead encode durable metadata from which its provider and plans can be reconstructed. -The example codecs do not inspect the callback `TaskContext`. A production codec that depends on session configuration or registered functions must ensure its exported FFI codec is bound to, and retains, the appropriate host `TaskContextProvider`. +Both codec getters take the `SessionContext` they are being installed on and pull the `TaskContextProvider` off it, so decode callbacks resolve session configuration and registered functions against the session that is running the query. Passing `require_udf_on_decode` to either constructor makes every decode call resolve a named scalar function out of that context, which is how the tests check where the registry came from. This example makes the provider library the sole external codec owner. Register both provider codecs before installing the planner: diff --git a/examples/datafusion-ffi-example/src/catalog_provider.rs b/examples/datafusion-ffi-example/src/catalog_provider.rs index a56b5855c..75890d083 100644 --- a/examples/datafusion-ffi-example/src/catalog_provider.rs +++ b/examples/datafusion-ffi-example/src/catalog_provider.rs @@ -94,7 +94,7 @@ impl FixedSchemaProvider { ) -> PyResult> { let provider = Arc::clone(&self.inner) as Arc; - let codec = ffi_logical_codec_from_pycapsule(session)?; + let codec = ffi_logical_codec_from_pycapsule(session, None)?; let provider = FFI_SchemaProvider::new_with_ffi_codec(provider, None, codec); PyCapsule::new_with_value(py, provider, cr"datafusion_schema_provider") @@ -186,7 +186,7 @@ impl MyCatalogProvider { ) -> PyResult> { let provider = Arc::clone(&self.inner) as Arc; - let codec = ffi_logical_codec_from_pycapsule(session)?; + let codec = ffi_logical_codec_from_pycapsule(session, None)?; let provider = FFI_CatalogProvider::new_with_ffi_codec(provider, None, codec); PyCapsule::new_with_value(py, provider, cr"datafusion_catalog_provider") @@ -245,7 +245,7 @@ impl MyCatalogProviderList { ) -> PyResult> { let provider = Arc::clone(&self.inner) as Arc; - let codec = ffi_logical_codec_from_pycapsule(session)?; + let codec = ffi_logical_codec_from_pycapsule(session, None)?; let provider = FFI_CatalogProviderList::new_with_ffi_codec(provider, None, codec); PyCapsule::new_with_value(py, provider, cr"datafusion_catalog_provider_list") diff --git a/examples/datafusion-ffi-example/src/logical_extension_codec.rs b/examples/datafusion-ffi-example/src/logical_extension_codec.rs index 987ea2dbb..424e953a4 100644 --- a/examples/datafusion-ffi-example/src/logical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/logical_extension_codec.rs @@ -24,16 +24,15 @@ use arrow::datatypes::SchemaRef; use datafusion::catalog::MemTable; use datafusion::common::{DataFusionError, Result, TableReference}; use datafusion::datasource::TableProvider; -use datafusion::execution::{TaskContext, TaskContextProvider}; +use datafusion::execution::TaskContext; use datafusion::logical_expr::{Extension, LogicalPlan, ScalarUDF}; -use datafusion::prelude::SessionContext; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; -use datafusion_python_util::get_tokio_runtime; +use datafusion_python_util::{ffi_task_context_provider_from_pycapsule, get_tokio_runtime}; use pyo3::prelude::*; use pyo3::types::PyCapsule; -use crate::required_udf::{new_codec_context, resolve_required_udf}; +use crate::required_udf::resolve_required_udf; const TABLE_PROVIDER_TOKEN: &[u8] = b"DFPYEXTP"; static NEXT_TABLE_PROVIDER_ID: AtomicU64 = AtomicU64::new(1); @@ -69,9 +68,6 @@ struct CountingLogicalExtensionCodec { /// Scalar function every table-provider decode must resolve from the /// `TaskContext` it is handed. See [`crate::required_udf`]. required_udf: Option, - // The FFI task-context handle is weak. Retain its provider for as long as - // this codec can be called, even if Python drops the exporter object. - _ctx_provider: Arc, } impl fmt::Debug for CountingLogicalExtensionCodec { @@ -171,7 +167,6 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec { pub(crate) struct MyLogicalExtensionCodec { counters: Arc, required_udf: Option, - ctx_provider: Arc, } #[pymethods] @@ -188,7 +183,6 @@ impl MyLogicalExtensionCodec { Self { counters: Arc::new(CallCounters::default()), required_udf: require_udf_on_decode, - ctx_provider: new_codec_context(), } } @@ -215,20 +209,25 @@ impl MyLogicalExtensionCodec { self.counters.decode_table_provider.load(Ordering::SeqCst) } + /// Export the codec, bound to the session it is being installed on. + /// + /// `session` supplies the `TaskContextProvider` the FFI decode callbacks + /// resolve, so this library never constructs a `SessionContext` and the + /// callbacks see the registry of the session running the query. fn __datafusion_logical_extension_codec__<'py>( &self, py: Python<'py>, + session: Bound<'py, PyAny>, ) -> PyResult> { let inner: Arc = Arc::new(CountingLogicalExtensionCodec { inner: DefaultLogicalExtensionCodec {}, counters: Arc::clone(&self.counters), required_udf: self.required_udf.clone(), - _ctx_provider: Arc::clone(&self.ctx_provider), }); let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider: Arc = self.ctx_provider.clone(); - let ffi = FFI_LogicalExtensionCodec::new(inner, Some(runtime), &ctx_provider); + let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; + let ffi = FFI_LogicalExtensionCodec::new(inner, Some(runtime), ctx_provider); PyCapsule::new_with_value(py, ffi, cr"datafusion_logical_extension_codec") } diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs index 801d47515..7d9601912 100644 --- a/examples/datafusion-ffi-example/src/physical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -22,20 +22,19 @@ use std::sync::{Arc, Mutex, OnceLock}; use datafusion::common::{DataFusionError, Result}; use datafusion::datasource::source::DataSourceExec; -use datafusion::execution::{TaskContext, TaskContextProvider}; +use datafusion::execution::TaskContext; use datafusion::logical_expr::ScalarUDF; use datafusion::physical_plan::ExecutionPlan; -use datafusion::prelude::SessionContext; use datafusion_ffi::execution_plan::ForeignExecutionPlan; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, }; -use datafusion_python_util::get_tokio_runtime; +use datafusion_python_util::{ffi_task_context_provider_from_pycapsule, get_tokio_runtime}; use pyo3::prelude::*; use pyo3::types::PyCapsule; -use crate::required_udf::{new_codec_context, resolve_required_udf}; +use crate::required_udf::resolve_required_udf; const EXECUTION_PLAN_TOKEN: &[u8] = b"DFPYEXEP"; static NEXT_EXECUTION_PLAN_ID: AtomicU64 = AtomicU64::new(1); @@ -70,9 +69,6 @@ struct CountingPhysicalExtensionCodec { /// Scalar function every decode call must resolve from the `TaskContext` /// it is handed. See [`crate::required_udf`]. required_udf: Option, - // The FFI task-context handle is weak. Keep its provider alive with the - // codec rather than relying on the lifetime of the Python exporter. - _ctx_provider: Arc, } impl fmt::Debug for CountingPhysicalExtensionCodec { @@ -161,7 +157,6 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { pub(crate) struct MyPhysicalExtensionCodec { counters: Arc, required_udf: Option, - ctx_provider: Arc, } #[pymethods] @@ -178,7 +173,6 @@ impl MyPhysicalExtensionCodec { Self { counters: Arc::new(PhysicalCallCounters::default()), required_udf: require_udf_on_decode, - ctx_provider: new_codec_context(), } } @@ -205,21 +199,25 @@ impl MyPhysicalExtensionCodec { self.counters.decode_execution_plan.load(Ordering::SeqCst) } + /// Export the codec, bound to the session it is being installed on. + /// + /// See [`crate::logical_extension_codec::MyLogicalExtensionCodec`] for why + /// `session` is taken rather than a context this library invents. fn __datafusion_physical_extension_codec__<'py>( &self, py: Python<'py>, + session: Bound<'py, PyAny>, ) -> PyResult> { let inner: Arc = Arc::new(CountingPhysicalExtensionCodec { inner: DefaultPhysicalExtensionCodec {}, counters: Arc::clone(&self.counters), required_udf: self.required_udf.clone(), - _ctx_provider: Arc::clone(&self.ctx_provider), }); let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider: Arc = self.ctx_provider.clone(); - let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), &ctx_provider); + let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; + let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), ctx_provider); PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") } diff --git a/examples/datafusion-ffi-example/src/required_udf.rs b/examples/datafusion-ffi-example/src/required_udf.rs index 00e6bd5ee..e107f0478 100644 --- a/examples/datafusion-ffi-example/src/required_udf.rs +++ b/examples/datafusion-ffi-example/src/required_udf.rs @@ -17,85 +17,22 @@ //! Support for exercising the `TaskContext` a codec is handed when it decodes. //! -//! A codec exported over FFI is built with a `TaskContextProvider`, and the -//! decode callbacks in `datafusion-ffi` resolve that provider to a -//! `TaskContext` before calling into the codec. Nothing in the example codecs -//! read anything out of that context, so which session it belongs to was -//! untestable: the token registries they use are keyed by an integer and -//! ignore the registry entirely. +//! A codec exported over FFI carries a `TaskContextProvider`, and the decode +//! callbacks in `datafusion-ffi` resolve it to a `TaskContext` before calling +//! into the codec. Nothing in the example codecs read anything out of that +//! context, so which session it belongs to was untestable: the token +//! registries they use are keyed by an integer and ignore the registry. //! //! The codecs can now be asked to resolve a named scalar function from the //! context they are given on every decode, which makes the answer observable. -//! Two names matter: -//! -//! - [`LIBRARY_LOCAL_UDF_NAME`] is registered by [`new_codec_context`] on the -//! context this library owns and hands to the FFI codec, so a decode -//! callback resolves it. -//! - A function registered only on the *host* `SessionContext` does not -//! resolve, because the decode callback never sees the host's registry. +//! Because the codecs take their provider from the session they are installed +//! on, a function registered on the host with `register_udf` resolves. -use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use arrow_schema::DataType; use datafusion::execution::TaskContext; -use datafusion::prelude::SessionContext; use datafusion_common::error::Result as DataFusionResult; use datafusion_common::plan_err; -use datafusion_expr::{ - ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, - Volatility, -}; - -/// Name of the scalar function registered on the context this library owns. -pub(crate) const LIBRARY_LOCAL_UDF_NAME: &str = "library_local_marker"; - -/// Placeholder scalar function used only as a registry entry. -/// -/// Decode callbacks look it up by name to prove which `TaskContext` they were -/// handed. It is never invoked, so the body is unreachable in practice. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -struct LibraryLocalUDF { - signature: Signature, -} - -impl LibraryLocalUDF { - fn new() -> Self { - Self { - signature: Signature::new(TypeSignature::Any(1), Volatility::Immutable), - } - } -} - -impl ScalarUDFImpl for LibraryLocalUDF { - fn name(&self) -> &str { - LIBRARY_LOCAL_UDF_NAME - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn return_type(&self, _arg_types: &[DataType]) -> DataFusionResult { - Ok(DataType::Boolean) - } - - fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> DataFusionResult { - plan_err!("{LIBRARY_LOCAL_UDF_NAME} exists only as a registry entry and cannot be invoked") - } -} - -/// Builds the context a codec keeps for its own FFI `TaskContextProvider`. -/// -/// This is the context every extension library has to conjure in order to -/// export a codec, and it is a plain empty session apart from the marker -/// function. The host's registrations are not in it, which is the point the -/// codec tests make observable. -pub(crate) fn new_codec_context() -> Arc { - let ctx = SessionContext::new(); - ctx.register_udf(ScalarUDF::from(LibraryLocalUDF::new())); - Arc::new(ctx) -} /// Resolves `required` against `ctx`, the context the decode callback was given. /// diff --git a/examples/datafusion-ffi-example/src/table_function.rs b/examples/datafusion-ffi-example/src/table_function.rs index 55543cb59..e653aeab1 100644 --- a/examples/datafusion-ffi-example/src/table_function.rs +++ b/examples/datafusion-ffi-example/src/table_function.rs @@ -48,7 +48,7 @@ impl MyTableFunction { session: Bound, ) -> PyResult> { let func = self.clone(); - let codec = ffi_logical_codec_from_pycapsule(session)?; + let codec = ffi_logical_codec_from_pycapsule(session, None)?; let provider = FFI_TableFunction::new_with_ffi_codec(Arc::new(func), None, codec); PyCapsule::new_with_value(py, provider, cr"datafusion_table_function") diff --git a/examples/datafusion-ffi-example/src/table_provider.rs b/examples/datafusion-ffi-example/src/table_provider.rs index 5756e6d02..ef6430e29 100644 --- a/examples/datafusion-ffi-example/src/table_provider.rs +++ b/examples/datafusion-ffi-example/src/table_provider.rs @@ -103,7 +103,7 @@ impl MyTableProvider { .create_table() .map_err(|e: DataFusionError| PyRuntimeError::new_err(e.to_string()))?; - let codec = ffi_logical_codec_from_pycapsule(session)?; + let codec = ffi_logical_codec_from_pycapsule(session, None)?; let provider = FFI_TableProvider::new_with_ffi_codec(Arc::new(provider), false, None, codec); diff --git a/examples/datafusion-ffi-example/src/table_provider_factory.rs b/examples/datafusion-ffi-example/src/table_provider_factory.rs index 71dfd73ca..df0845119 100644 --- a/examples/datafusion-ffi-example/src/table_provider_factory.rs +++ b/examples/datafusion-ffi-example/src/table_provider_factory.rs @@ -77,7 +77,7 @@ impl MyTableProviderFactory { py: Python<'py>, codec: Bound, ) -> PyResult> { - let codec = ffi_logical_codec_from_pycapsule(codec)?; + let codec = ffi_logical_codec_from_pycapsule(codec, None)?; let factory = Arc::clone(&self.inner) as Arc; let factory = FFI_TableProviderFactory::new_with_ffi_codec(factory, None, codec); diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index 446b74f19..6fb8c3f9f 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -42,11 +42,10 @@ def configured_context(max_rows: int): return ctx, logical_codec, physical_codec -LIBRARY_LOCAL_UDF = "library_local_marker" -"""Scalar function registered on the context each example codec owns.""" - HOST_ONLY_UDF = "my_custom_is_null" -"""Scalar function ``configured_context`` registers on the host session only.""" +"""Scalar function registered on the host session and nowhere else.""" + +UNREGISTERED_UDF = "not_registered_anywhere" def probe_context( @@ -74,11 +73,14 @@ def probe_context( return ctx, logical_codec, physical_codec -def test_logical_codec_resolves_its_own_libraries_udf(): - """``try_decode_table_provider`` sees the registry of the codec's own session.""" - ctx, logical_codec, _physical_codec = probe_context( - logical_requires=LIBRARY_LOCAL_UDF - ) +def test_logical_codec_resolves_a_host_registered_udf(): + """``try_decode_table_provider`` sees the host session's registry. + + The codec takes its task context provider from the session it is installed + on, so a function the host registered is resolvable inside a decode + callback running in the other library. + """ + ctx, logical_codec, _physical_codec = probe_context(logical_requires=HOST_ONLY_UDF) batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() assert batches[0].column(0).to_pylist() == [0, 1, 2] @@ -86,11 +88,9 @@ def test_logical_codec_resolves_its_own_libraries_udf(): assert logical_codec.task_context_udf_resolutions() > 0 -def test_physical_codec_resolves_its_own_libraries_udf(): - """``try_decode`` sees the registry of the codec's own session.""" - ctx, _logical_codec, physical_codec = probe_context( - physical_requires=LIBRARY_LOCAL_UDF - ) +def test_physical_codec_resolves_a_host_registered_udf(): + """``try_decode`` sees the host session's registry, as above.""" + ctx, _logical_codec, physical_codec = probe_context(physical_requires=HOST_ONLY_UDF) batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() assert batches[0].column(0).to_pylist() == [0, 1, 2] @@ -98,48 +98,54 @@ def test_physical_codec_resolves_its_own_libraries_udf(): assert physical_codec.task_context_udf_resolutions() > 0 -@pytest.mark.xfail( - strict=True, - reason=( - "A codec exported over FFI is built with its own TaskContextProvider, so " - "its decode callbacks resolve names against the exporting library's " - "session rather than the host session that is running the query. A " - "function registered only on the host is therefore invisible to them." - ), -) -def test_logical_codec_resolves_a_host_registered_udf(): - ctx, logical_codec, _physical_codec = probe_context(logical_requires=HOST_ONLY_UDF) +def test_codec_still_reports_a_name_registered_nowhere(): + """Negative control: resolution really is a lookup, not an unconditional pass.""" + ctx, _logical_codec, _physical_codec = probe_context( + logical_requires=UNREGISTERED_UDF + ) - ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() - assert logical_codec.task_context_udf_resolutions() > 0 + with pytest.raises( + Exception, match=rf"could not resolve scalar function '{UNREGISTERED_UDF}'" + ): + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() -@pytest.mark.xfail( - strict=True, - reason=( - "Same as the logical case: the physical codec's decode callback is " - "handed the exporting library's task context, which does not know about " - "functions registered on the host session." - ), -) -def test_physical_codec_resolves_a_host_registered_udf(): - ctx, _logical_codec, physical_codec = probe_context(physical_requires=HOST_ONLY_UDF) +def test_codec_sees_a_udf_registered_after_it_was_installed(): + """The provider is a live handle to the session, not a snapshot of it. - ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() - assert physical_codec.task_context_udf_resolutions() > 0 + A function registered after the codec is installed is still visible to the + decode callback. It has to be registered before the planner, though: + installing a foreign planner forks the session, and the codec keeps + pointing at the one it was built against. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=HOST_ONLY_UDF) + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + # Registered after the codec was installed and bound to this session. + ctx.register_udf(udf(IsNullUDF())) + ctx = ctx.with_query_planner(MyQueryPlanner()) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert logical_codec.task_context_udf_resolutions() > 0 -def test_decode_task_context_is_not_the_host_session(): - """Pin the current behaviour so the failure mode above stays legible. +def test_codec_does_not_see_a_udf_registered_after_the_planner_fork(): + """Installing a planner forks; the codec still points at the pre-fork session. - The host registers ``my_custom_is_null`` and can use it in a query, but the - codec's decode callback cannot see it, and the error names the session it - was given instead. + This is the same rule as the documented fork caveat, seen from the codec's + side. Register before installing the planner, as the guide advises. """ - ctx, _logical_codec, _physical_codec = probe_context(logical_requires=HOST_ONLY_UDF) + ctx, _logical_codec, _physical_codec = probe_context() + ctx = ctx.with_logical_extension_codec( + MyLogicalExtensionCodec(require_udf_on_decode="registered_after_fork") + ) with pytest.raises( - Exception, match=rf"could not resolve scalar function '{HOST_ONLY_UDF}'" + Exception, match=r"could not resolve scalar function 'registered_after_fork'" ): ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() @@ -150,7 +156,7 @@ def test_three_library_query_planner(raw_capsule: bool): ctx, logical_codec, physical_codec = configured_context(max_rows=3) planner = MyQueryPlanner() exported_planner = ( - planner.__datafusion_query_planner__() if raw_capsule else planner + planner.__datafusion_query_planner__(ctx) if raw_capsule else planner ) ctx = ctx.with_query_planner(exported_planner) diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index 85c0e80d9..7928cedf4 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -21,8 +21,6 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use async_trait::async_trait; use datafusion::common::DataFusionError; -use datafusion::execution::TaskContextProvider; -use datafusion::execution::context::SessionContext; use datafusion::logical_expr::LogicalPlan; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::limit::GlobalLimitExec; @@ -33,9 +31,10 @@ use datafusion_ffi::execution_plan::ForeignExecutionPlan; use datafusion_ffi::query_planner::FFI_QueryPlanner; use datafusion_ffi::session::ForeignSession; use datafusion_ffi::table_provider::ForeignTableProvider; -use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; -use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; -use datafusion_python_util::{ffi_query_planner_from_pycapsule, get_tokio_runtime}; +use datafusion_python_util::{ + ffi_logical_codec_from_pycapsule, ffi_physical_codec_from_pycapsule, + ffi_query_planner_from_pycapsule, +}; use datafusion_session::{QueryPlanner, Session}; use pyo3::prelude::*; use pyo3::types::PyCapsule; @@ -132,18 +131,9 @@ fn planner_config(session: &dyn Session) -> datafusion::common::Result, - /// Keeps the context behind the exported codecs' `TaskContextProvider` - /// alive. See [`MyQueryPlanner::codec_ctx`]. The reference lives here as - /// well as on `MyQueryPlanner` because this is the value the capsule - /// carries, so the provider stays valid even if the Python object that - /// exported it is dropped first. - #[expect( - dead_code, - reason = "strong reference keeping the FFI codecs' weakly held provider alive" - )] - codec_ctx: Arc, /// Planner to hand the work to instead of planning here. /// /// This is how a real planner layers on top of an existing one. The capsule @@ -158,16 +148,6 @@ struct DistributedQueryPlanner { fallback: Option>, } -impl fmt::Debug for DistributedQueryPlanner { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("DistributedQueryPlanner") - .field("observations", &self.observations) - .field("has_fallback", &self.fallback.is_some()) - .finish_non_exhaustive() - } -} - #[async_trait] impl QueryPlanner for DistributedQueryPlanner { async fn create_physical_plan( @@ -220,48 +200,10 @@ impl QueryPlanner for DistributedQueryPlanner { module = "datafusion_ffi_query_planner_example", subclass )] -#[derive(Clone)] +#[derive(Debug, Default, Clone)] pub(crate) struct MyQueryPlanner { observations: Arc, fallback: Option>, - /// Context backing the `TaskContextProvider` handed to the exported - /// codecs. - /// - /// This is *not* the session that arrives at `create_physical_plan`. That - /// one belongs to the host and is what this planner reads config from. The - /// provider here serves the other direction: when the host decodes the - /// plan bytes this library returned, any extension node in them is decoded - /// by a callback back into this library, and that callback needs a - /// `TaskContext` whose registry can resolve *this* library's nodes and - /// functions. A real planner library registers its UDFs and extension - /// types on this context; the example uses the default codecs, so nothing - /// ever calls back. - /// - /// It must be owned rather than built inline: `FFI_TaskContextProvider` - /// downgrades the provider to a `Weak`, so a temporary would already be - /// dropped by the time the capsule is used, and every codec callback would - /// fail with "TaskContextProvider went out of scope over FFI boundary". - codec_ctx: Arc, -} - -impl Default for MyQueryPlanner { - fn default() -> Self { - Self { - observations: Arc::default(), - fallback: None, - codec_ctx: Arc::new(SessionContext::new()), - } - } -} - -impl fmt::Debug for MyQueryPlanner { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("MyQueryPlanner") - .field("observations", &self.observations) - .field("has_fallback", &self.fallback.is_some()) - .finish_non_exhaustive() - } } #[pymethods] @@ -277,7 +219,7 @@ impl MyQueryPlanner { fn new(fallback: Option>) -> PyResult { let fallback = fallback .map(|planner| { - ffi_query_planner_from_pycapsule(&planner) + ffi_query_planner_from_pycapsule(&planner, None) .map(|ffi| -> Arc { (&ffi).into() }) }) .transpose()?; @@ -311,24 +253,24 @@ impl MyQueryPlanner { self.observations.foreign_plan.load(Ordering::SeqCst) } + /// Export the planner, bound to the session it is being installed on. + /// + /// The codecs come off `session` rather than being built here. They carry + /// the host's `TaskContextProvider`, so this library never constructs a + /// `SessionContext`, and `with_query_planner` would rebind them to the + /// running session anyway. fn __datafusion_query_planner__<'py>( &self, py: Python<'py>, + session: Bound<'py, PyAny>, ) -> PyResult> { let planner: Arc = Arc::new(DistributedQueryPlanner { observations: Arc::clone(&self.observations), - codec_ctx: Arc::clone(&self.codec_ctx), fallback: self.fallback.clone(), }); - let runtime = get_tokio_runtime().handle().clone(); - let ctx_provider = Arc::clone(&self.codec_ctx) as Arc; - let ffi = FFI_QueryPlanner::new( - planner, - Some(runtime), - &ctx_provider, - Arc::new(DefaultLogicalExtensionCodec {}), - Arc::new(DefaultPhysicalExtensionCodec {}), - ); + let logical_codec = ffi_logical_codec_from_pycapsule(session.clone(), None)?; + let physical_codec = ffi_physical_codec_from_pycapsule(session, None)?; + let ffi = FFI_QueryPlanner::new_with_ffi_codecs(planner, logical_codec, physical_codec); PyCapsule::new_with_value(py, ffi, cr"datafusion_query_planner") } } diff --git a/python/datafusion/context.py b/python/datafusion/context.py index b1f4f74f0..6e0560f15 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -149,10 +149,12 @@ class QueryPlannerExportable(Protocol): """Type hint for object that has a __datafusion_query_planner__ PyCapsule. The method returns a PyCapsule wrapping an ``FFI_QueryPlanner``, typically - produced by a separate compiled extension. + produced by a separate compiled extension. ``session`` is the + :py:class:`SessionContext` the planner is being installed on; take the + extension codecs from it rather than building your own. """ - def __datafusion_query_planner__(self) -> object: ... # noqa: D105 + def __datafusion_query_planner__(self, session: Any) -> object: ... # noqa: D105 class SessionConfig: @@ -2232,13 +2234,22 @@ def __datafusion_task_context_provider__(self) -> Any: """Access the PyCapsule FFI_TaskContextProvider.""" return self.ctx.__datafusion_task_context_provider__() - def __datafusion_logical_extension_codec__(self) -> Any: - """Access the PyCapsule FFI_LogicalExtensionCodec.""" - return self.ctx.__datafusion_logical_extension_codec__() + def __datafusion_logical_extension_codec__(self, session: Any = None) -> Any: + """Access the PyCapsule FFI_LogicalExtensionCodec. - def __datafusion_query_planner__(self) -> Any: - """Access the ``FFI_QueryPlanner`` PyCapsule for the current planner.""" - return self.ctx.__datafusion_query_planner__() + ``session`` is accepted so a context satisfies the same protocol an + extension library implements, where the argument is how the library + reaches the session it is being installed on. A context already is one, + so the argument is ignored. + """ + return self.ctx.__datafusion_logical_extension_codec__(session) + + def __datafusion_query_planner__(self, session: Any = None) -> Any: + """Access the ``FFI_QueryPlanner`` PyCapsule for the current planner. + + See :meth:`__datafusion_logical_extension_codec__` for ``session``. + """ + return self.ctx.__datafusion_query_planner__(session) def with_logical_extension_codec( self, codec: LogicalExtensionCodecExportable | _PyCapsule @@ -2260,9 +2271,12 @@ def with_logical_extension_codec( new.ctx = new_internal return new - def __datafusion_physical_extension_codec__(self) -> Any: - """Access the PyCapsule FFI_PhysicalExtensionCodec.""" - return self.ctx.__datafusion_physical_extension_codec__() + def __datafusion_physical_extension_codec__(self, session: Any = None) -> Any: + """Access the PyCapsule FFI_PhysicalExtensionCodec. + + See :meth:`__datafusion_logical_extension_codec__` for ``session``. + """ + return self.ctx.__datafusion_physical_extension_codec__(session) def with_physical_extension_codec( self, codec: PhysicalExtensionCodecExportable | _PyCapsule diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index 394c682ae..43b53e469 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -114,15 +114,28 @@ def _is_pycapsule(value: object) -> TypeGuard[_PyCapsule]: class LogicalExtensionCodecExportable(Protocol): - """Type hint for objects exposing ``__datafusion_logical_extension_codec__``.""" + """Type hint for objects exposing ``__datafusion_logical_extension_codec__``. - def __datafusion_logical_extension_codec__(self) -> object: ... # noqa: D105 + ``session`` is the :py:class:`~datafusion.context.SessionContext` the codec + is being installed on. Take the task context provider from it rather than + building a session of your own, so the decode callbacks resolve names + against the session that runs the query. + """ + + def __datafusion_logical_extension_codec__( # noqa: D105 + self, session: Any + ) -> object: ... class PhysicalExtensionCodecExportable(Protocol): - """Type hint for objects exposing ``__datafusion_physical_extension_codec__``.""" + """Type hint for objects exposing ``__datafusion_physical_extension_codec__``. + + See :py:class:`LogicalExtensionCodecExportable` for ``session``. + """ - def __datafusion_physical_extension_codec__(self) -> object: ... # noqa: D105 + def __datafusion_physical_extension_codec__( # noqa: D105 + self, session: Any + ) -> object: ... class ScalarUDF: From 68d4e6da93787341481e80c25c75c63805c51ac0 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 08:02:46 -0400 Subject: [PATCH 14/30] docs: record the FFI capsule protocol as a convention The session-passing rule was already settled for four getters and documented in the 52.0.0 upgrade guide, but nothing pointed an agent or a new contributor at it before they wrote a fifth. Write it down where it will be found. Add the 55.0.0 upgrade guide entry this branch owes. Changing `__datafusion_logical_extension_codec__` and `__datafusion_physical_extension_codec__` to take a session breaks every extension library implementing them, so it needs before/after Rust in the same shape as the 52.0.0 entry. Correct `user-guide/io/table_provider.md`. It still showed the pre-52.0.0 signature with no session and a `PyCapsule::new_bound` call, so the one page a reader is most likely to find contradicted the convention. Add `.ai/skills/ffi-capsule-protocol/`. Its description is written as a trigger rather than a task, because the existing skills are all things to run on request and a convention read as one would be skipped. It leads with enumerating the family, which is the step that makes the rest unnecessary. Point `CLAUDE.md` at it, since that file loads unconditionally and a skill only helps once someone goes looking. Also note that `docs/temp/` is gitignored build output that `grep -r` surfaces with stale copies, and require an upgrade guide section alongside the `api change` label, so a breaking change forces a visit to the file that records the conventions. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/skills/ffi-capsule-protocol/SKILL.md | 141 ++++++++++++++++++++ AGENTS.md | 33 ++++- docs/source/user-guide/io/table_provider.md | 13 +- docs/source/user-guide/upgrade-guides.md | 54 ++++++++ 4 files changed, 235 insertions(+), 6 deletions(-) create mode 100644 .ai/skills/ffi-capsule-protocol/SKILL.md diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md new file mode 100644 index 000000000..72247a26f --- /dev/null +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -0,0 +1,141 @@ + + +--- +name: ffi-capsule-protocol +description: "TRIGGER — read before adding, changing, or reviewing any __datafusion_*__ capsule getter, any FFI_* export that asks for a TaskContextProvider or an extension codec, or any code that calls FFI_QueryPlanner::new / FFI_TableProvider::new / FFI_{Logical,Physical}ExtensionCodec::new. These methods are one protocol with a settled convention. Do not design it fresh; do not construct a SessionContext inside an extension library." +argument-hint: "[getter name] (e.g., \"__datafusion_query_planner__\", \"table provider\", \"codec\", or omit to review the whole family)" +--- + +# FFI Capsule Protocol + +`datafusion-python` shares Rust objects with extension libraries through +PyCapsules. Every hook is a dunder method named `__datafusion___` that +returns a capsule wrapping an FFI-safe struct. They are **one protocol**, not a +collection of unrelated methods, and they have a settled convention that has +already been migrated once (see `docs/source/user-guide/upgrade-guides.md`, +DataFusion 52.0.0 and 55.0.0). + +## Rule 1 — enumerate the family before you change a member + +Do this first, every time. It takes one command and it is the whole point of +this skill: + +```bash +grep -rn "__datafusion_[a-z_]*__" --include="*.rs" crates/ examples/*/src/ +``` + +Compare the signature you are about to write against what the others already +do. If yours is shaped differently, that is a finding about your design, not +about theirs. + +Search `docs/source/`, never `docs/` — `docs/temp/` is gitignored build output +that shadows the real files with stale copies. + +## Rule 2 — a getter takes the session it is being installed on + +```rust +fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, +) -> PyResult> { ... } +``` + +The host calls the getter and passes itself. That argument is how an extension +library reaches things only the session has. + +`SessionContext` implements the same getters and ignores the argument, so a +session satisfies the protocol too — `ctx.__datafusion_query_planner__()` and +`ctx.__datafusion_query_planner__(ctx)` are both valid. + +## Rule 3 — never construct a `SessionContext` in an extension library + +The FFI constructors ask for things a library does not have: + +| Constructor | Wants | Take it from | +|---|---|---| +| `FFI_{Logical,Physical}ExtensionCodec::new` | `TaskContextProvider` | `ffi_task_context_provider_from_pycapsule(&session)` | +| `FFI_TableProvider::new_with_ffi_codec` | logical codec | `ffi_logical_codec_from_pycapsule(session, None)` | +| `FFI_QueryPlanner::new_with_ffi_codecs` | both codecs | `ffi_{logical,physical}_codec_from_pycapsule(session, None)` | + +`Arc::new(SessionContext::new())` is the wrong answer to all three, for two +independent reasons: + +1. **It is the wrong registry.** Decode callbacks resolve names against + whatever provider the codec carries. An empty session resolves nothing, so a + function the host registered with `register_udf` is invisible to a node that + references it by name. +2. **It dangles.** `FFI_TaskContextProvider` downgrades its provider to a + `Weak`. A context built inline in the getter is dropped before the capsule + is ever used, and every callback then fails with `TaskContextProvider went + out of scope over FFI boundary`. + +Prefer the `*_with_ffi_codec(s)` constructors when they exist. They take +prebuilt codecs that already carry the host's provider, so there is no provider +parameter to get wrong. + +## Rule 4 — the helpers live in `crates/util/src/lib.rs` + +`ffi_logical_codec_from_pycapsule`, `ffi_physical_codec_from_pycapsule`, +`ffi_query_planner_from_pycapsule`, `ffi_task_context_provider_from_pycapsule`, +`table_provider_from_pycapsule`. Each takes the object and, where relevant, an +`Option<&Bound>` session: + +- `Some(session)` — importing a *foreign* object; the getter needs the session. +- `None` — the object already *is* a session and is being asked for what it + holds. + +Adding a getter means adding a helper here, not hand-rolling capsule +extraction at the call site. + +## Rule 5 — changing a getter's signature is a breaking change + +Extension libraries implement these methods. A signature change breaks every +one of them, and the failure is a bare `TypeError` from a `call1`. So: + +- Add a section to `docs/source/user-guide/upgrade-guides.md` with before/after + Rust, matching the 52.0.0 and 55.0.0 entries. +- Add the `api change` label to the PR. +- Map the `TypeError` to a diagnosable message. `call_capsule_getter` in + `crates/util/src/lib.rs` already does this; reuse it. +- Update `python/datafusion/context.py` and + `python/datafusion/user_defined.py`, where the `Protocol` type hints for + these methods live. + +## Rule 6 — the session a codec captures is not always the running session + +Installing a foreign query planner **forks** the session +(`PySessionContext::ctx_with_rebound_planner`). A codec installed beforehand +keeps pointing at the pre-fork session, so registrations made after the fork +are invisible to it. Install codecs before the planner, which is what the +contributor guide already advises. `PySessionContext::ancestors` keeps forked +parents alive so the captured session cannot be dropped out from under a codec. + +## Where the truth is + +- `docs/source/contributor-guide/ffi.md` — the protocol, the fork caveat. +- `docs/source/user-guide/upgrade-guides.md` — every past migration. +- `examples/datafusion-ffi-example/src/` — provider, catalog, function, codec + getters, all in current form. +- `examples/datafusion-ffi-query-planner-example/src/planner.rs` — planner + getter. +- `examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py` + — `require_udf_on_decode` proves which session a decode callback resolves + against. Extend these when touching the protocol. diff --git a/AGENTS.md b/AGENTS.md index fda08b23c..1ab28fb5f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -35,7 +35,33 @@ Skills follow the [Agent Skills](https://agentskills.io) open standard. Each ski To discover what skills are available, list `.ai/skills/` and read each `SKILL.md`. The frontmatter `name` and `description` fields summarize the -skill's purpose. +skill's purpose. Some descriptions begin with `TRIGGER —`; those are not tasks +to run on request but conventions to read *before* writing code that meets the +stated condition. + +## FFI Capsule Protocol + +Before adding or changing any `__datafusion_*__` capsule getter, read +[`.ai/skills/ffi-capsule-protocol/SKILL.md`](.ai/skills/ffi-capsule-protocol/SKILL.md) +and enumerate the existing family: + +```bash +grep -rn "__datafusion_[a-z_]*__" --include="*.rs" crates/ examples/*/src/ +``` + +These methods are one protocol with a settled convention: a getter receives the +`SessionContext` it is being installed on and takes what the FFI constructor +needs from that session. Never construct a `SessionContext` inside an extension +library to satisfy a `TaskContextProvider` or codec parameter. + +## Documentation Sources + +Search and edit `docs/source/`. `docs/temp/` is gitignored build output that +`grep -r` will surface with stale copies of the same pages. + +Before changing a public API, check +`docs/source/user-guide/upgrade-guides.md` for how the same API family was +migrated previously. Follow the established pattern rather than inventing one. ## Pull Requests @@ -48,7 +74,10 @@ Every pull request must follow the template in 3. **What changes are included in this PR?** — Summarize the individual changes. 4. **Are there any user-facing changes?** — Note any changes visible to users (new APIs, changed behavior, new files shipped in the package, etc.). If - there are breaking changes to public APIs, add the `api change` label. + there are breaking changes to public APIs, add the `api change` label **and + add a section to `docs/source/user-guide/upgrade-guides.md`** showing the + before and after. This applies to FFI hook method signatures, which + extension libraries implement. ## Pre-commit Checks diff --git a/docs/source/user-guide/io/table_provider.md b/docs/source/user-guide/io/table_provider.md index 3c436ba1d..5dc2dc086 100644 --- a/docs/source/user-guide/io/table_provider.md +++ b/docs/source/user-guide/io/table_provider.md @@ -29,6 +29,11 @@ via [PyCapsule](https://pyo3.rs/main/doc/pyo3/types/struct.pycapsule). A complete example can be found in the [examples folder](https://github.com/apache/datafusion-python/tree/main/examples). +The method takes the `SessionContext` it is being registered on. Take whatever +the FFI constructor needs from that session — here the logical extension codec — +rather than building one inside your library. See the {ref}`ffi` guide for the +full capsule protocol. + ```rust #[pymethods] impl MyTableProvider { @@ -36,13 +41,13 @@ impl MyTableProvider { fn __datafusion_table_provider__<'py>( &self, py: Python<'py>, + session: Bound<'py, PyAny>, ) -> PyResult> { - let name = cr"datafusion_table_provider".into(); - let provider = Arc::new(self.clone()); - let provider = FFI_TableProvider::new(provider, false, None); + let codec = ffi_logical_codec_from_pycapsule(session, None)?; + let provider = FFI_TableProvider::new_with_ffi_codec(provider, false, None, codec); - PyCapsule::new_bound(py, provider, Some(name.clone())) + PyCapsule::new_with_value(py, provider, cr"datafusion_table_provider") } } ``` diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 360e0533c..63129117b 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -19,6 +19,60 @@ # Upgrade Guides +## DataFusion 55.0.0 + +This release extends the change made in 52.0.0 to the remaining {ref}`ffi` hook +methods. Users who contribute their own `LogicalExtensionCodec` or +`PhysicalExtensionCodec` via FFI must update +`__datafusion_logical_extension_codec__` and +`__datafusion_physical_extension_codec__` to accept an additional +`session: Bound` parameter, and take the `TaskContextProvider` from that +session rather than constructing a `SessionContext` of their own. + +Before: + +```rust +fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, +) -> PyResult> { + let ctx_provider: Arc = Arc::clone(&self.ctx_provider); + let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), &ctx_provider); + PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") +} +``` + +After: + +```rust +fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, +) -> PyResult> { + let ctx_provider = ffi_task_context_provider_from_pycapsule(&session)?; + let ffi = FFI_PhysicalExtensionCodec::new(inner, Some(runtime), ctx_provider); + PyCapsule::new_with_value(py, ffi, cr"datafusion_physical_extension_codec") +} +``` + +A codec that keeps its own `SessionContext` still compiles, but its decode +callbacks resolve names against that empty session instead of the one running +the query, so a function registered with `SessionContext.register_udf` is not +visible to it. Taking the provider from `session` also removes a lifetime +hazard: `FFI_TaskContextProvider` holds its provider weakly, so a context +constructed inside the getter is already dropped by the time the capsule is +used. + +`SessionContext` accepts the argument on its own capsule getters and ignores +it, so existing calls such as `ctx.__datafusion_logical_extension_codec__()` +continue to work unchanged. + +New in this release, `__datafusion_query_planner__` follows the same protocol. +It receives the session and takes both extension codecs from it, so a planner +library never builds a `TaskContextProvider` at all. See the {ref}`ffi` guide +for the full protocol. + ## DataFusion 54.0.0 The `Config` class has been removed. It was a standalone wrapper around From ffb35213267736082fd4241d9af44b25395f8bb3 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 08:15:47 -0400 Subject: [PATCH 15/30] fix: put skill frontmatter before the license header MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every `.ai/skills/*/SKILL.md` opened with the ASF header and only then the YAML frontmatter, which has to be the first thing in the file. The result was that no skill's `description` was readable: the skill listing showed ` ---- -name: audit-skill-md -description: Audit the user-facing skill at skills/datafusion_python/SKILL.md against the current public Python API. Find new APIs that should be documented, stale mentions of removed/renamed APIs, examples that drifted from current idiomatic style, and places that need a "requires datafusion-python NN or newer" note. Run after upstream syncs and before each release. -argument-hint: [scope] (e.g., "session-context", "dataframe", "expr", "functions", "patterns", "pitfalls", "version-notes", "all") ---- - # Audit `skills/datafusion_python/SKILL.md` You are auditing the user-facing skill at diff --git a/.ai/skills/check-upstream/SKILL.md b/.ai/skills/check-upstream/SKILL.md index a3d82a670..828f227d8 100644 --- a/.ai/skills/check-upstream/SKILL.md +++ b/.ai/skills/check-upstream/SKILL.md @@ -1,3 +1,9 @@ +--- +name: check-upstream +description: Check if upstream Apache DataFusion features (functions, DataFrame ops, SessionContext methods, FFI types) are exposed in this Python project. Use when adding missing functions, auditing API coverage, or ensuring parity with upstream. +argument-hint: [area] (e.g., "scalar functions", "aggregate functions", "window functions", "dataframe", "session context", "ffi types", "all") +--- + ---- -name: check-upstream -description: Check if upstream Apache DataFusion features (functions, DataFrame ops, SessionContext methods, FFI types) are exposed in this Python project. Use when adding missing functions, auditing API coverage, or ensuring parity with upstream. -argument-hint: [area] (e.g., "scalar functions", "aggregate functions", "window functions", "dataframe", "session context", "ffi types", "all") ---- - # Check Upstream DataFusion Feature Coverage You are auditing the datafusion-python project to find features from the upstream Apache DataFusion Rust library that are **not yet exposed** in this Python binding project. Your goal is to identify gaps and, if asked, implement the missing bindings. diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 72247a26f..3a5dd84d4 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -1,3 +1,9 @@ +--- +name: ffi-capsule-protocol +description: "TRIGGER — read before adding, changing, or reviewing any __datafusion_*__ capsule getter, any FFI_* export that asks for a TaskContextProvider or an extension codec, or any code that calls FFI_QueryPlanner::new / FFI_TableProvider::new / FFI_{Logical,Physical}ExtensionCodec::new. These methods are one protocol with a settled convention. Do not design it fresh; do not construct a SessionContext inside an extension library." +argument-hint: "[getter name] (e.g., \"__datafusion_query_planner__\", \"table provider\", \"codec\", or omit to review the whole family)" +--- + ---- -name: ffi-capsule-protocol -description: "TRIGGER — read before adding, changing, or reviewing any __datafusion_*__ capsule getter, any FFI_* export that asks for a TaskContextProvider or an extension codec, or any code that calls FFI_QueryPlanner::new / FFI_TableProvider::new / FFI_{Logical,Physical}ExtensionCodec::new. These methods are one protocol with a settled convention. Do not design it fresh; do not construct a SessionContext inside an extension library." -argument-hint: "[getter name] (e.g., \"__datafusion_query_planner__\", \"table provider\", \"codec\", or omit to review the whole family)" ---- - # FFI Capsule Protocol `datafusion-python` shares Rust objects with extension libraries through @@ -45,9 +45,6 @@ Compare the signature you are about to write against what the others already do. If yours is shaped differently, that is a finding about your design, not about theirs. -Search `docs/source/`, never `docs/` — `docs/temp/` is gitignored build output -that shadows the real files with stale copies. - ## Rule 2 — a getter takes the session it is being installed on ```rust diff --git a/.ai/skills/make-pythonic/SKILL.md b/.ai/skills/make-pythonic/SKILL.md index 7d490ec03..24c2bb817 100644 --- a/.ai/skills/make-pythonic/SKILL.md +++ b/.ai/skills/make-pythonic/SKILL.md @@ -1,3 +1,9 @@ +--- +name: make-pythonic +description: Audit and improve datafusion-python functions to accept native Python types (int, float, str, bool) instead of requiring explicit lit() or col() wrapping. Analyzes function signatures, checks upstream Rust implementations for type constraints, and applies the appropriate coercion pattern. +argument-hint: [scope] (e.g., "string functions", "datetime functions", "array functions", "math functions", "all", or a specific function name like "split_part") +--- + ---- -name: make-pythonic -description: Audit and improve datafusion-python functions to accept native Python types (int, float, str, bool) instead of requiring explicit lit() or col() wrapping. Analyzes function signatures, checks upstream Rust implementations for type constraints, and applies the appropriate coercion pattern. -argument-hint: [scope] (e.g., "string functions", "datetime functions", "array functions", "math functions", "all", or a specific function name like "split_part") ---- - # Make Python API Functions More Pythonic You are improving the datafusion-python API to feel more natural to Python users. The goal is to allow functions to accept native Python types (int, float, str, bool, etc.) for arguments that are contextually always or typically literal values, instead of requiring users to manually wrap them in `lit()`. diff --git a/AGENTS.md b/AGENTS.md index 1ab28fb5f..327ebd643 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -41,18 +41,9 @@ stated condition. ## FFI Capsule Protocol -Before adding or changing any `__datafusion_*__` capsule getter, read -[`.ai/skills/ffi-capsule-protocol/SKILL.md`](.ai/skills/ffi-capsule-protocol/SKILL.md) -and enumerate the existing family: - -```bash -grep -rn "__datafusion_[a-z_]*__" --include="*.rs" crates/ examples/*/src/ -``` - -These methods are one protocol with a settled convention: a getter receives the -`SessionContext` it is being installed on and takes what the FFI constructor -needs from that session. Never construct a `SessionContext` inside an extension -library to satisfy a `TaskContextProvider` or codec parameter. +The `__datafusion_*__` capsule getters are one protocol with a settled +convention. Before adding or changing one, read +[`.ai/skills/ffi-capsule-protocol/SKILL.md`](.ai/skills/ffi-capsule-protocol/SKILL.md). ## Documentation Sources From 1e1d1aa8d720d101246bc15def51622cd1d67de5 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 13:34:21 -0400 Subject: [PATCH 16/30] Update temporary DF version with corrections in FFI --- Cargo.lock | 76 +++++++++++++++++++++++++++--------------------------- Cargo.toml | 22 ++++++++-------- 2 files changed, 49 insertions(+), 49 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5549ae8ad..a7b81a195 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -792,7 +792,7 @@ dependencies = [ [[package]] name = "datafusion" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-schema", @@ -845,7 +845,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "async-trait", @@ -869,7 +869,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "async-trait", @@ -892,7 +892,7 @@ dependencies = [ [[package]] name = "datafusion-common" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-ipc", @@ -918,7 +918,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "futures", "log", @@ -928,7 +928,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "async-compression", @@ -964,7 +964,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-ipc", @@ -988,7 +988,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-avro", @@ -1006,7 +1006,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "async-trait", @@ -1029,7 +1029,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "async-trait", @@ -1052,7 +1052,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-schema", @@ -1084,12 +1084,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" [[package]] name = "datafusion-execution" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-buffer", @@ -1114,7 +1114,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-schema", @@ -1138,7 +1138,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1149,7 +1149,7 @@ dependencies = [ [[package]] name = "datafusion-ffi" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-schema", @@ -1220,7 +1220,7 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-buffer", @@ -1251,7 +1251,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1271,7 +1271,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1282,7 +1282,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-ord", @@ -1306,7 +1306,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "async-trait", @@ -1321,7 +1321,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1337,7 +1337,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1346,7 +1346,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "datafusion-doc", "quote", @@ -1356,7 +1356,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "chrono", @@ -1375,7 +1375,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1397,7 +1397,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1411,7 +1411,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "chrono", @@ -1428,7 +1428,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1447,7 +1447,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "arrow-data", @@ -1483,7 +1483,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-catalog", @@ -1509,7 +1509,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1519,7 +1519,7 @@ dependencies = [ [[package]] name = "datafusion-proto-models" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "datafusion-common", "datafusion-proto-common", @@ -1529,7 +1529,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "datafusion-common", @@ -1589,7 +1589,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow-schema", "async-trait", @@ -1603,7 +1603,7 @@ dependencies = [ [[package]] name = "datafusion-spark" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "bigdecimal", @@ -1632,7 +1632,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "arrow", "bigdecimal", @@ -1651,7 +1651,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "55.0.0" -source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-session-create-physical-plan-codec-55#24d8903b0e4530642feb53319d7e7de1be244f27" +source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-constructor-argument-drop-55#59740d5de45d20156a8bd8e91db4812a6dbd5197" dependencies = [ "async-recursion", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 7efe64a7c..7e71045a9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -78,14 +78,14 @@ codegen-units = 2 # We cannot publish to crates.io with any patches in the below section. Developers # must remove any entries in this section before creating a release candidate. [patch.crates-io] -datafusion = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } -datafusion-substrait = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } -datafusion-proto = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } -datafusion-ffi = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } -datafusion-catalog = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } -datafusion-common = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } -datafusion-functions-aggregate = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } -datafusion-functions-window = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } -datafusion-spark = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } -datafusion-expr = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } -datafusion-session = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-session-create-physical-plan-codec-55" } +datafusion = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-substrait = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-proto = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-ffi = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-catalog = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-common = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-functions-aggregate = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-functions-window = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-spark = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-expr = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } +datafusion-session = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } From 282e3a5a846606676bbfd672186cbb8df5611657 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 14:00:42 -0400 Subject: [PATCH 17/30] feat: rebind foreign codecs when a planner install forks the session MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Installing a foreign query planner writes to `SessionState`, and `with_query_planner` must not modify its receiver, so it forks. A foreign codec holds an `FFI_TaskContextProvider` pointing at the session it was installed on, and until now the fork could not move it: passing a new provider to `FFI_LogicalExtensionCodec::new` was silently discarded whenever the codec was already foreign. The fork rebound only its own outer wrapper, so decode callbacks in the extension library kept answering from the pre-fork registry, and the pre-fork session had to be retained or the weakly held provider dangled. apache/datafusion#24722 fixes the discard; those constructors now adopt the provider on the already-foreign path. Repoint the patch at the branch carrying it and rebind both codecs onto the fork. Verified the branch carries everything already pinned rather than trusting the commit graph, which reports the two as diverged: across 3811 files the only differences are the four constructors from the fix, and `datafusion/ffi/src/session/mod.rs` is byte-identical, so the `create_physical_plan` codec fix arrives as its branch-55 backport. `ancestors` and its helpers are deleted. They existed only to keep the pre-fork session alive for a codec that could not be moved off it, and a codec bound to the running session needs no such anchor. Three tests, replacing two that were weaker than they looked. One registers a function on the fork after the codec was installed on its parent and resolves it, which is the direct evidence the rebind happened; it failed before this change. One installs a planner twice and asserts the first context still cannot resolve a function registered only on the second, covering the clone-before-adopt half — a rebind that mutated the shared handle would pass the first test and fail this one. The third keeps the live-handle case. The test it replaces required a name registered nowhere, so it passed for the same reason as the negative control and never exercised a fork at all. Note the version floor in `Cargo.toml` rather than raising it now: the patched branch still reports 55.0.0, so the requirement can only move to 55.1.0 when the patch section is removed. Building against 55.0.0 without the patch would compile and silently skip the rebind. Co-Authored-By: Claude Opus 5 (1M context) --- .ai/skills/ffi-capsule-protocol/SKILL.md | 20 ++- Cargo.toml | 5 + crates/core/src/context.rs | 140 ++++++++++-------- docs/source/contributor-guide/ffi.md | 30 ++-- .../_test_three_library_query_planner.py | 63 +++++--- 5 files changed, 161 insertions(+), 97 deletions(-) diff --git a/.ai/skills/ffi-capsule-protocol/SKILL.md b/.ai/skills/ffi-capsule-protocol/SKILL.md index 3a5dd84d4..fc20b2b8e 100644 --- a/.ai/skills/ffi-capsule-protocol/SKILL.md +++ b/.ai/skills/ffi-capsule-protocol/SKILL.md @@ -116,14 +116,22 @@ one of them, and the failure is a bare `TypeError` from a `call1`. So: `python/datafusion/user_defined.py`, where the `Protocol` type hints for these methods live. -## Rule 6 — the session a codec captures is not always the running session +## Rule 6 — a fork must rebind the codecs it carries Installing a foreign query planner **forks** the session -(`PySessionContext::ctx_with_rebound_planner`). A codec installed beforehand -keeps pointing at the pre-fork session, so registrations made after the fork -are invisible to it. Install codecs before the planner, which is what the -contributor guide already advises. `PySessionContext::ancestors` keeps forked -parents alive so the captured session cannot be dropped out from under a codec. +(`PySessionContext::derived_parts`), because installing one writes to +`SessionState` and the receiver must not be modified. A foreign codec holds an +`FFI_TaskContextProvider` pointing at the session it was installed on, so the +fork rebinds each one onto itself via +`PySessionContext::rebound_{logical,physical}_codec`. Skip that and decode +callbacks answer from the pre-fork registry. + +Rebinding relies on `FFI_{Logical,Physical}ExtensionCodec::new` adopting the +provider on the already-foreign path, which needs DataFusion 55.1.0 or newer +(apache/datafusion#24722). It clones the handle before overwriting, so the +receiver keeps its own binding — assert that, not just that the fork works. +A codec this library owns round-trips unchanged, so the rebind is safe to apply +unconditionally. ## Where the truth is diff --git a/Cargo.toml b/Cargo.toml index 7e71045a9..7c86cfd61 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -77,6 +77,11 @@ codegen-units = 2 # We cannot publish to crates.io with any patches in the below section. Developers # must remove any entries in this section before creating a release candidate. +# +# When these are removed, raise the DataFusion requirement above to 55.1.0. The +# FFI codec rebinding in `PySessionContext::derived_parts` is a silent no-op +# before that release (apache/datafusion#24722), so a 55.0.0 build would compile +# and then resolve decode callbacks against the pre-fork session. [patch.crates-io] datafusion = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } datafusion-substrait = { git = "https://github.com/timsaucer/datafusion", branch = "fix/ffi-constructor-argument-drop-55" } diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index db5250d3a..3ea770765 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -373,16 +373,6 @@ pub struct PySessionContext { pub ctx: Arc, logical_codec: Arc, physical_codec: Arc, - /// Sessions this one was derived from, held only to keep them alive. - /// - /// A foreign codec installed here is built against the session current at - /// the time of the call, and `FFI_TaskContextProvider` holds that session - /// weakly. Installing a foreign query planner afterwards forks, and - /// rebinding the Python name would otherwise drop the session the codec - /// still points at. Retaining ancestors here rather than on the codec - /// avoids a cycle: nothing reachable from a `SessionContext` reaches back - /// to a `PySessionContext`. - ancestors: Vec>, } #[pymethods] @@ -415,7 +405,6 @@ impl PySessionContext { ctx, logical_codec: Arc::new(PythonLogicalCodec::default()), physical_codec: Arc::new(PythonPhysicalCodec::default()), - ancestors: Vec::new(), }) } @@ -424,7 +413,6 @@ impl PySessionContext { ctx: Arc::new(self.ctx.as_ref().clone().enable_url_table()), logical_codec: Arc::clone(&self.logical_codec), physical_codec: Arc::clone(&self.physical_codec), - ancestors: self.ancestors_including_self(), }) } @@ -436,7 +424,6 @@ impl PySessionContext { ctx, logical_codec: Arc::new(PythonLogicalCodec::default()), physical_codec: Arc::new(PythonPhysicalCodec::default()), - ancestors: Vec::new(), }) } @@ -1231,14 +1218,16 @@ impl PySessionContext { ) -> PyDataFusionResult { let planner = ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))?; let this = slf.borrow(); - let ctx = - this.ctx_with_rebound_planner(&this.logical_codec, &this.physical_codec, Some(planner)); + let (ctx, logical_codec, physical_codec) = this.derived_parts( + Arc::clone(&this.logical_codec), + Arc::clone(&this.physical_codec), + Some(planner), + ); Ok(Self { - ancestors: this.ancestors_for(&ctx), ctx, - logical_codec: Arc::clone(&this.logical_codec), - physical_codec: Arc::clone(&this.physical_codec), + logical_codec, + physical_codec, }) } @@ -1452,10 +1441,9 @@ impl PySessionContext { let logical_codec = Arc::new(PythonLogicalCodec::new(inner)); let this = slf.borrow(); - let physical_codec = Arc::clone(&this.physical_codec); - let ctx = this.ctx_with_rebound_planner(&logical_codec, &physical_codec, None); + let (ctx, logical_codec, physical_codec) = + this.derived_parts(logical_codec, Arc::clone(&this.physical_codec), None); Ok(Self { - ancestors: this.ancestors_for(&ctx), ctx, logical_codec, physical_codec, @@ -1482,10 +1470,9 @@ impl PySessionContext { let physical_codec = Arc::new(PythonPhysicalCodec::new(inner)); let this = slf.borrow(); - let logical_codec = Arc::clone(&this.logical_codec); - let ctx = this.ctx_with_rebound_planner(&logical_codec, &physical_codec, None); + let (ctx, logical_codec, physical_codec) = + this.derived_parts(Arc::clone(&this.logical_codec), physical_codec, None); Ok(Self { - ancestors: this.ancestors_for(&ctx), ctx, logical_codec, physical_codec, @@ -1501,9 +1488,9 @@ impl PySessionContext { PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner())) .with_python_udf_inlining(enabled), ); - let ctx = self.ctx_with_rebound_planner(&logical_codec, &physical_codec, None); + let (ctx, logical_codec, physical_codec) = + self.derived_parts(logical_codec, physical_codec, None); Self { - ancestors: self.ancestors_for(&ctx), ctx, logical_codec, physical_codec, @@ -1512,50 +1499,35 @@ impl PySessionContext { } impl PySessionContext { - /// Keep-alive list for a context derived from this one. - /// - /// Sharing the same `SessionContext` carries the list through unchanged. - /// A fork adds this session, because a foreign codec installed here holds - /// it weakly and would otherwise dangle once the Python name is rebound. - fn ancestors_for(&self, derived: &Arc) -> Vec> { - if Arc::ptr_eq(derived, &self.ctx) { - self.ancestors.clone() - } else { - self.ancestors_including_self() - } - } - - fn ancestors_including_self(&self) -> Vec> { - let mut ancestors = self.ancestors.clone(); - ancestors.push(Arc::clone(&self.ctx)); - ancestors - } - - /// Return the [`SessionContext`] a derived `PySessionContext` should use, - /// binding a foreign query planner to `logical_codec` and `physical_codec`. + /// Return the pieces a derived `PySessionContext` should hold, binding any + /// foreign query planner and foreign codecs to the session that will run + /// the query. /// /// Pass `Some(planner)` to install one, or `None` to rebind whichever /// planner the session already holds. /// /// With no foreign planner in play there is nothing to rebind, so the - /// existing context is shared and swapping codecs alone does not change - /// what a derived context observes. A foreign planner does have to be - /// rebound, and the FFI codecs capture the context they are built against, - /// so that case forks the session: the codecs are built from the fork and - /// then the fork's state is overwritten in place, because rebuilding the - /// context afterwards would leave the codecs pointing at a session that is - /// no longer used for planning. + /// existing context and codecs are shared unchanged. A foreign planner does + /// have to be rebound, and that forks the session, because installing it + /// writes to `SessionState` and the receiver must not be modified. The + /// codecs are built from the fork and the fork's state is then overwritten + /// in place, since rebuilding the context afterwards would leave them + /// pointing at a session that is no longer used for planning. /// /// A fork is not a deep copy. `SessionState` keeps its catalog list behind /// an `Arc`, so catalogs and tables stay shared with the original session, /// while registered functions and the configuration are snapshotted at the /// time of the call. The session id is deliberately carried over. - fn ctx_with_rebound_planner( + fn derived_parts( &self, - logical_codec: &Arc, - physical_codec: &Arc, + logical_codec: Arc, + physical_codec: Arc, planner: Option, - ) -> Arc { + ) -> ( + Arc, + Arc, + Arc, + ) { let state = self.ctx.state(); // When the caller is only replacing codecs, recover the handle behind @@ -1570,15 +1542,18 @@ impl PySessionContext { }); let Some(planner) = planner else { - return Arc::clone(&self.ctx); + return (Arc::clone(&self.ctx), logical_codec, physical_codec); }; let ctx = Arc::new(SessionContext::new_with_state(state)); + let logical_codec = Arc::new(Self::rebound_logical_codec(&ctx, &logical_codec)); + let physical_codec = Arc::new(Self::rebound_physical_codec(&ctx, &physical_codec)); + let inner: Arc = (&planner).into(); let planner: Arc = (&FFI_QueryPlanner::new_with_ffi_codecs( inner, - Self::ffi_logical_codec_for(&ctx, logical_codec), - Self::ffi_physical_codec_for(&ctx, physical_codec), + Self::ffi_logical_codec_for(&ctx, &logical_codec), + Self::ffi_physical_codec_for(&ctx, &physical_codec), )) .into(); let state = SessionStateBuilder::new_from_existing(ctx.state()) @@ -1586,7 +1561,47 @@ impl PySessionContext { .build(); *ctx.state_ref().write() = state; - ctx + (ctx, logical_codec, physical_codec) + } + + /// Rebind a foreign logical codec to `ctx`. + /// + /// A codec imported from another library holds an `FFI_TaskContextProvider` + /// pointing at whichever session it was installed on. Its decode callbacks + /// resolve names against that session, so a fork has to move them onto the + /// fork or they keep answering from the pre-fork registry. + /// + /// `FFI_LogicalExtensionCodec::new` adopts the provider supplied here when + /// the codec is already foreign, returning a *clone* of the handle, so the + /// context this codec was derived from keeps its own binding. A codec this + /// library owns has no stored provider, and round-trips back to the same + /// `Arc` unchanged, so this is safe to apply unconditionally. + fn rebound_logical_codec( + ctx: &Arc, + codec: &Arc, + ) -> PythonLogicalCodec { + let ctx_provider = Arc::clone(ctx) as Arc; + let runtime = get_tokio_runtime().handle().clone(); + let rebound = + FFI_LogicalExtensionCodec::new(Arc::clone(codec.inner()), Some(runtime), &ctx_provider); + PythonLogicalCodec::new((&rebound).into()) + .with_python_udf_inlining(codec.python_udf_inlining()) + } + + /// Physical counterpart of [`Self::rebound_logical_codec`]. + fn rebound_physical_codec( + ctx: &Arc, + codec: &Arc, + ) -> PythonPhysicalCodec { + let ctx_provider = Arc::clone(ctx) as Arc; + let runtime = get_tokio_runtime().handle().clone(); + let rebound = FFI_PhysicalExtensionCodec::new( + Arc::clone(codec.inner()), + Some(runtime), + &ctx_provider, + ); + PythonPhysicalCodec::new((&rebound).into()) + .with_python_udf_inlining(codec.python_udf_inlining()) } async fn _table(&self, name: &str) -> datafusion::common::Result { @@ -1816,7 +1831,6 @@ impl From for PySessionContext { ctx: Arc::new(ctx), logical_codec: Arc::new(PythonLogicalCodec::default()), physical_codec: Arc::new(PythonPhysicalCodec::default()), - ancestors: Vec::new(), } } } diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 886341b89..4509127ff 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -294,19 +294,27 @@ satisfies the same protocol an extension library implements. When you export the planner to wrap it, `ctx.__datafusion_query_planner__()` and `ctx.__datafusion_query_planner__(ctx)` are both fine. -### A codec decodes against the session it was installed on +### A codec decodes against the session that is running the query Because the provider comes from the host, a decode callback running inside an extension -library resolves names against the session that is running the query. A function -registered with `ctx.register_udf(...)` is visible to a foreign codec decoding a node -that references it by name, and the handle is live rather than a snapshot, so a -registration made after the codec is installed is visible too. - -The one boundary is the fork. Installing a foreign query planner forks the session, and -a codec installed beforehand keeps pointing at the session it was built against, so a -function registered after that point is not visible to it. This is the same rule as the -caveat below, seen from the codec's side: register before deriving. Both halves are -covered in +library resolves names against the session running the query. A function registered with +`ctx.register_udf(...)` is visible to a foreign codec decoding a node that references it +by name, and the handle is live rather than a snapshot, so a registration made after the +codec is installed is visible too. + +This survives a fork. Installing a foreign query planner forks the session, and the fork +rebinds every foreign codec it carries onto the new session, so a function registered +after the fork is still visible to them. `FFI_LogicalExtensionCodec::new` adopts the +provider supplied to it when the codec is already foreign, returning a clone of the +handle, so the context the fork was derived from keeps its own binding and continues to +resolve against its own session. + +That behaviour requires DataFusion 55.1.0 or newer. Before it, those constructors +silently discarded the provider, a foreign codec could not be rebound, and a fork left +it resolving against the pre-fork session +([apache/datafusion#24722](https://github.com/apache/datafusion/issues/24722)). + +All three properties are covered in `examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py`, where the example codecs take a `require_udf_on_decode` name and resolve it out of the task context they are handed. diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index 6fb8c3f9f..0e65f3d19 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -111,13 +111,7 @@ def test_codec_still_reports_a_name_registered_nowhere(): def test_codec_sees_a_udf_registered_after_it_was_installed(): - """The provider is a live handle to the session, not a snapshot of it. - - A function registered after the codec is installed is still visible to the - decode callback. It has to be registered before the planner, though: - installing a foreign planner forks the session, and the codec keeps - pointing at the one it was built against. - """ + """The provider is a live handle to the session, not a snapshot of it.""" config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=HOST_ONLY_UDF) ctx = SessionContext(config) @@ -133,21 +127,56 @@ def test_codec_sees_a_udf_registered_after_it_was_installed(): assert logical_codec.task_context_udf_resolutions() > 0 -def test_codec_does_not_see_a_udf_registered_after_the_planner_fork(): - """Installing a planner forks; the codec still points at the pre-fork session. +def test_codec_follows_the_session_across_a_planner_fork(): + """Installing a planner forks the session, and the codec moves with it. - This is the same rule as the documented fork caveat, seen from the codec's - side. Register before installing the planner, as the guide advises. + The codec is bound to the session before the fork and the function is + registered on the fork afterwards, so resolving it proves the codec's task + context provider was rebound to the forked session rather than left + pointing at the one it was installed on. """ - ctx, _logical_codec, _physical_codec = probe_context() - ctx = ctx.with_logical_extension_codec( - MyLogicalExtensionCodec(require_udf_on_decode="registered_after_fork") - ) + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=HOST_ONLY_UDF) + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx = ctx.with_query_planner(MyQueryPlanner()) + # Registered on the fork, after the codec was installed on its parent. + ctx.register_udf(udf(IsNullUDF())) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert logical_codec.table_provider_decode_calls() > 0 + assert logical_codec.task_context_udf_resolutions() > 0 + + +def test_rebinding_a_fork_leaves_the_receiver_bound_to_its_own_session(): + """Rebinding the fork's codec must not disturb the context it came from. + + ``FFI_LogicalExtensionCodec::new`` clones the handle before adopting the + new provider, so each context keeps its own binding. Both contexts here + have a planner, so both exercise the codec; the function is registered only + on the second, and only the second can resolve it. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + logical_codec = MyLogicalExtensionCodec(require_udf_on_decode=HOST_ONLY_UDF) + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(MyPhysicalExtensionCodec()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + first = ctx.with_query_planner(MyQueryPlanner()) + second = first.with_query_planner(MyQueryPlanner()) + second.register_udf(udf(IsNullUDF())) + + batches = second.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] with pytest.raises( - Exception, match=r"could not resolve scalar function 'registered_after_fork'" + Exception, match=rf"could not resolve scalar function '{HOST_ONLY_UDF}'" ): - ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + first.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() @pytest.mark.parametrize("raw_capsule", [False, True]) From b7fbd342d5fcb24bcfb24a3f1bc352efd94860d7 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 14:49:30 -0400 Subject: [PATCH 18/30] fix: restore base64 0.23.1 in Cargo.lock Regenerating the lock against the patched DataFusion fork silently downgraded base64 from 0.23.1 to 0.23.0. Nothing requires the older version -- neither the fork nor upstream 55.0.0 constrains it -- so this was incidental churn from the lockfile refresh, not a resolution result. Restores the checksum main already had and re-points the three dependents (datafusion-common, datafusion-functions, parquet). No other dependency moves; cargo metadata --locked still resolves cleanly. Co-Authored-By: Claude Opus 5 (1M context) --- Cargo.lock | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a7b81a195..d433d4be2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -202,7 +202,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64 0.23.0", + "base64 0.23.1", "chrono", "comfy-table", "half", @@ -429,9 +429,9 @@ checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" [[package]] name = "base64" -version = "0.23.0" +version = "0.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b25655df2c3cdd83c5e5b293b88acd880332b2ddadd7c30ac43144fdc0033da9" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" [[package]] name = "bigdecimal" @@ -1224,7 +1224,7 @@ source = "git+https://github.com/timsaucer/datafusion?branch=fix%2Fffi-construct dependencies = [ "arrow", "arrow-buffer", - "base64 0.23.0", + "base64 0.23.1", "blake2", "blake3", "chrono", @@ -2684,7 +2684,7 @@ dependencies = [ "arrow-ipc", "arrow-schema", "arrow-select", - "base64 0.23.0", + "base64 0.23.1", "brotli", "bytes", "chrono", From 7ada7c69f21724a9fa835b9eee50e7bc34ed79a9 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 14:50:45 -0400 Subject: [PATCH 19/30] fix: preserve session id across a planner fork Installing a foreign query planner forks the session state, and the fork was minting a new session id. SessionStateBuilder::new_from_existing drops the id and build() replaces it with a fresh UUID, while SessionContext had already cached the original into a field of its own back at new_with_state. Overwriting the state in place afterwards left the two disagreeing: session_id() returned the pre-fork id, every TaskContext handed to a foreign codec carried a different one. Nothing in DataFusion core keys on the session id beyond debug logging, so this broke no in-tree behavior. It matters at the FFI boundary, where session id equality is the idiom for "which session is this codec bound to", and for extension libraries correlating host-side and worker-side state. Upstream hit the same case in SessionContext::enable_url_table and preserves the id explicitly, guarded by preserve_session_context_id. Passing the id through the builder makes the fork, its state, and its TaskContexts agree, which is what the derived_parts doc comment and the FFI contributor guide already claimed. Verified by reading the id out of a decode callback via the example codec's require_udf_on_decode error path -- the only way to observe the state-side id from Python -- with and without the fix. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 10 +++++++++- docs/source/contributor-guide/ffi.md | 5 ++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 3ea770765..324e56911 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1517,7 +1517,8 @@ impl PySessionContext { /// A fork is not a deep copy. `SessionState` keeps its catalog list behind /// an `Arc`, so catalogs and tables stay shared with the original session, /// while registered functions and the configuration are snapshotted at the - /// time of the call. The session id is deliberately carried over. + /// time of the call. The session id is explicitly carried over, so the fork + /// and its `TaskContext`s all report the id the original session had. fn derived_parts( &self, logical_codec: Arc, @@ -1556,7 +1557,14 @@ impl PySessionContext { Self::ffi_physical_codec_for(&ctx, &physical_codec), )) .into(); + // `with_session_id` is load-bearing, not redundant. + // `SessionStateBuilder::new_from_existing` drops the id and `build` + // mints a fresh one, while `SessionContext` cached the original in a + // field of its own at `new_with_state`. Without this the fork would + // report one id from `session_id()` and another from every + // `TaskContext` it hands out. let state = SessionStateBuilder::new_from_existing(ctx.state()) + .with_session_id(ctx.session_id()) .with_query_planner(planner) .build(); *ctx.state_ref().write() = state; diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 4509127ff..f107cfbe2 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -344,7 +344,10 @@ the fork behave differently: context afterwards is not visible to the derived one, and a `SET` applied to one does not reach the other. -The session id is carried over to the fork, so both contexts report the same id. +The session id is carried over to the fork, so both contexts report the same id, and so +does every `TaskContext` either one hands to a foreign codec. A library that identifies a +session by its id — to correlate host-side and worker-side state, or to assert which +session a decode callback was bound to — keeps working across a fork. Register functions before deriving, or register them directly on the derived context: From 283745aea41534b8ed75f82ba4cd26a62fb77541 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 14:56:07 -0400 Subject: [PATCH 20/30] fix: preserve session id in add_physical_optimizer_rule Same drift just fixed in derived_parts, but on a path that never forks. add_physical_optimizer_rule rebuilds SessionState through SessionStateBuilder::new_from_existing and writes it straight back into the caller's own session, so the fresh id build() mints replaces the one SessionContext had already cached at construction. The session the user is holding then reports one id from session_id() and a different one from every TaskContext it hands out, with no derivation to explain it. Reproduced against a foreign codec, reading the id back out of a decode callback: identical setup differing only by an add_physical_optimizer_rule call went from MATCH to DRIFT, and back to MATCH with the id threaded through the builder. This is the last new_from_existing call site in the crate. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 324e56911..865ebd586 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1205,7 +1205,13 @@ impl PySessionContext { let rule = physical_optimizer_rule_from_pycapsule(&rule)?; let state_ref = self.ctx.state_ref(); let mut guard = state_ref.write(); + // Rebuilding through the builder mints a fresh session id, but this + // mutates the caller's own session rather than deriving a new one, so + // the id has to survive. See `derived_parts` for why losing it leaves + // `session_id()` disagreeing with every `TaskContext` the session + // hands out. let new_state = SessionStateBuilder::new_from_existing(guard.clone()) + .with_session_id(guard.session_id().to_string()) .with_physical_optimizer_rule(rule) .build(); *guard = new_state; From 12b35fb01ef6717bfb87ce5e5c608f033cf5180d Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 15:08:06 -0400 Subject: [PATCH 21/30] test: guard the session id a codec decodes against The two session id fixes had no regression guard. A Python-level assertion cannot provide one: session_id() reads a copy SessionContext caches at construction, which stayed correct through both bugs. The id that actually moved was the one inside the TaskContext handed to a foreign codec's decode callback, which nothing exposed. Give the example codecs a TaskContextProbe that records it. This replaces the bare AtomicUsize the require_udf_on_decode support used, so the counter and the session id are recorded together, and the id is recorded on every decode rather than only when a function was requested. Three tests, all against the codec-side id rather than session_id(): a fork agrees with its codecs, add_physical_optimizer_rule does not move the id, and a two-deep fork chain leaves both halves on the parent's id. Confirmed non-vacuous: with both fixes reverted all three fail and the other 17 tests pass; with only the derived_parts fix restored, exactly the add_physical_optimizer_rule test still fails. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/logical_extension_codec.rs | 20 ++--- .../src/physical_extension_codec.rs | 20 ++--- .../src/required_udf.rs | 48 ++++++++-- .../_test_three_library_query_planner.py | 87 +++++++++++++++++++ 4 files changed, 150 insertions(+), 25 deletions(-) diff --git a/examples/datafusion-ffi-example/src/logical_extension_codec.rs b/examples/datafusion-ffi-example/src/logical_extension_codec.rs index 424e953a4..a127407f2 100644 --- a/examples/datafusion-ffi-example/src/logical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/logical_extension_codec.rs @@ -32,7 +32,7 @@ use datafusion_python_util::{ffi_task_context_provider_from_pycapsule, get_tokio use pyo3::prelude::*; use pyo3::types::PyCapsule; -use crate::required_udf::resolve_required_udf; +use crate::required_udf::{TaskContextProbe, resolve_required_udf}; const TABLE_PROVIDER_TOKEN: &[u8] = b"DFPYEXTP"; static NEXT_TABLE_PROVIDER_ID: AtomicU64 = AtomicU64::new(1); @@ -53,7 +53,7 @@ pub(crate) struct CallCounters { pub decode_udf: AtomicUsize, pub encode_table_provider: AtomicUsize, pub decode_table_provider: AtomicUsize, - pub task_ctx_udf_resolutions: AtomicUsize, + pub task_ctx: TaskContextProbe, } /// Example codec for objects owned by this extension library. @@ -101,11 +101,7 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec { schema: SchemaRef, ctx: &TaskContext, ) -> Result> { - resolve_required_udf( - self.required_udf.as_deref(), - ctx, - &self.counters.task_ctx_udf_resolutions, - )?; + resolve_required_udf(self.required_udf.as_deref(), ctx, &self.counters.task_ctx)?; if let Some(id) = token_id(buf, TABLE_PROVIDER_TOKEN) { self.counters .decode_table_provider @@ -188,9 +184,13 @@ impl MyLogicalExtensionCodec { /// Number of decode calls that resolved `require_udf_on_decode`. fn task_context_udf_resolutions(&self) -> usize { - self.counters - .task_ctx_udf_resolutions - .load(Ordering::SeqCst) + self.counters.task_ctx.resolutions() + } + + /// Session id of the `TaskContext` the most recent decode callback ran + /// against, or `None` before any decode. + fn last_task_context_session_id(&self) -> Option { + self.counters.task_ctx.last_session_id() } fn encode_udf_calls(&self) -> usize { diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs index 7d9601912..fe903a68b 100644 --- a/examples/datafusion-ffi-example/src/physical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -34,7 +34,7 @@ use datafusion_python_util::{ffi_task_context_provider_from_pycapsule, get_tokio use pyo3::prelude::*; use pyo3::types::PyCapsule; -use crate::required_udf::resolve_required_udf; +use crate::required_udf::{TaskContextProbe, resolve_required_udf}; const EXECUTION_PLAN_TOKEN: &[u8] = b"DFPYEXEP"; static NEXT_EXECUTION_PLAN_ID: AtomicU64 = AtomicU64::new(1); @@ -55,7 +55,7 @@ pub(crate) struct PhysicalCallCounters { pub decode_udf: AtomicUsize, pub encode_execution_plan: AtomicUsize, pub decode_execution_plan: AtomicUsize, - pub task_ctx_udf_resolutions: AtomicUsize, + pub task_ctx: TaskContextProbe, } /// Physical companion to the logical example codec. @@ -89,11 +89,7 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { ctx: &TaskContext, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - resolve_required_udf( - self.required_udf.as_deref(), - ctx, - &self.counters.task_ctx_udf_resolutions, - )?; + resolve_required_udf(self.required_udf.as_deref(), ctx, &self.counters.task_ctx)?; if let Some(id) = token_id(buf) { self.counters .decode_execution_plan @@ -178,9 +174,13 @@ impl MyPhysicalExtensionCodec { /// Number of decode calls that resolved `require_udf_on_decode`. fn task_context_udf_resolutions(&self) -> usize { - self.counters - .task_ctx_udf_resolutions - .load(Ordering::SeqCst) + self.counters.task_ctx.resolutions() + } + + /// Session id of the `TaskContext` the most recent decode callback ran + /// against, or `None` before any decode. + fn last_task_context_session_id(&self) -> Option { + self.counters.task_ctx.last_session_id() } fn encode_udf_calls(&self) -> usize { diff --git a/examples/datafusion-ffi-example/src/required_udf.rs b/examples/datafusion-ffi-example/src/required_udf.rs index e107f0478..a21362d7f 100644 --- a/examples/datafusion-ffi-example/src/required_udf.rs +++ b/examples/datafusion-ffi-example/src/required_udf.rs @@ -28,28 +28,66 @@ //! Because the codecs take their provider from the session they are installed //! on, a function registered on the host with `register_udf` resolves. +use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; use datafusion::execution::TaskContext; use datafusion_common::error::Result as DataFusionResult; use datafusion_common::plan_err; +/// What the codecs record about the `TaskContext` their decode callbacks run +/// against. +#[derive(Debug, Default)] +pub(crate) struct TaskContextProbe { + resolutions: AtomicUsize, + last_session_id: Mutex>, +} + +impl TaskContextProbe { + /// Successful `require_udf_on_decode` lookups since construction. + pub(crate) fn resolutions(&self) -> usize { + self.resolutions.load(Ordering::SeqCst) + } + + /// Session id of the most recent decode callback, or `None` if the codec + /// has not been asked to decode anything yet. + /// + /// Recorded on every decode, so a test can tell *which* session the + /// callback was bound to rather than only that some session resolved a + /// name. A `SessionContext` that derives a fork must keep reporting the id + /// it reports from `session_id()`. + pub(crate) fn last_session_id(&self) -> Option { + self.last_session_id + .lock() + .expect("task context probe mutex poisoned") + .clone() + } +} + /// Resolves `required` against `ctx`, the context the decode callback was given. /// -/// `Ok(())` when nothing was requested. Otherwise the name must be present in -/// the context's scalar function registry, and `resolutions` counts each -/// success so a test can tell a resolved lookup from a skipped one. +/// Records `ctx`'s session id either way. `Ok(())` when nothing was requested. +/// Otherwise the name must be present in the context's scalar function +/// registry, and `probe` counts each success so a test can tell a resolved +/// lookup from a skipped one. pub(crate) fn resolve_required_udf( required: Option<&str>, ctx: &TaskContext, - resolutions: &AtomicUsize, + probe: &TaskContextProbe, ) -> DataFusionResult<()> { + // Unconditional: the session id is worth observing even when the caller + // asked for no function. + *probe + .last_session_id + .lock() + .expect("task context probe mutex poisoned") = Some(ctx.session_id().to_string()); + let Some(name) = required else { return Ok(()); }; if ctx.scalar_functions().contains_key(name) { - resolutions.fetch_add(1, Ordering::SeqCst); + probe.resolutions.fetch_add(1, Ordering::SeqCst); return Ok(()); } diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index 0e65f3d19..b1b5b21c3 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -25,6 +25,7 @@ IsNullUDF, MyLogicalExtensionCodec, MyPhysicalExtensionCodec, + MyPhysicalOptimizerRule, MyTableProvider, ) from datafusion_ffi_query_planner_example import MyPlannerConfig, MyQueryPlanner @@ -179,6 +180,92 @@ def test_rebinding_a_fork_leaves_the_receiver_bound_to_its_own_session(): first.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() +def codec_context(max_rows: int = 3): + """Context with both example codecs installed and a table to scan. + + Unlike :func:`probe_context` the codecs ask for no function, so the only + thing they record is the session id of the task context they are handed. + No planner yet -- the caller installs one, since that is what forks. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=max_rows)) + logical_codec = MyLogicalExtensionCodec() + physical_codec = MyPhysicalExtensionCodec() + ctx = SessionContext(config) + ctx = ctx.with_logical_extension_codec(logical_codec) + ctx = ctx.with_physical_extension_codec(physical_codec) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + return ctx, logical_codec, physical_codec + + +def test_a_fork_decodes_against_the_session_id_it_reports(): + """A fork and its decode callbacks must agree on the session id. + + ``with_query_planner`` forks the session state by rebuilding it through + ``SessionStateBuilder``, which mints a fresh id unless handed one, and + ``SessionContext`` caches its id in a field of its own. Dropping the id + there leaves the fork reporting one id from ``session_id()`` and a + different one from every ``TaskContext`` it gives a foreign codec. + + Asserting on ``session_id()`` alone cannot catch that: it reads the cached + copy, which stays correct either way. The codec-side id is the only + observable that moves, which is what makes this worth a test rather than a + one-line equality check. + """ + ctx, logical_codec, physical_codec = codec_context() + session_id = ctx.session_id() + + fork = ctx.with_query_planner(MyQueryPlanner()) + assert fork.session_id() == session_id + + fork.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert logical_codec.last_task_context_session_id() == session_id + assert physical_codec.last_task_context_session_id() == session_id + + +def test_adding_a_physical_optimizer_rule_keeps_the_session_id(): + """Mutating a session in place must not move the id it decodes against. + + ``add_physical_optimizer_rule`` rebuilds ``SessionState`` and writes it + back into the caller's own session rather than deriving a new one, so a + regenerated id would desync a context from itself with no fork to explain + it. + """ + ctx, logical_codec, physical_codec = codec_context() + fork = ctx.with_query_planner(MyQueryPlanner()) + session_id = fork.session_id() + + fork.add_physical_optimizer_rule(MyPhysicalOptimizerRule()) + assert fork.session_id() == session_id + + fork.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert logical_codec.last_task_context_session_id() == session_id + assert physical_codec.last_task_context_session_id() == session_id + + +def test_rebinding_a_fork_does_not_move_the_parent_session_id(): + """Each context in a fork chain decodes against its own id. + + Two forks off one parent share the parent's id, so a test that only + compared against the parent would pass even if rebinding leaked one + context's provider into the other. Registering a function on just one fork + is what distinguishes them. + """ + ctx, logical_codec, _physical_codec = codec_context() + session_id = ctx.session_id() + + first = ctx.with_query_planner(MyQueryPlanner()) + second = first.with_query_planner(MyQueryPlanner()) + + assert first.session_id() == session_id + assert second.session_id() == session_id + + first.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert logical_codec.last_task_context_session_id() == session_id + + second.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert logical_codec.last_task_context_session_id() == session_id + + @pytest.mark.parametrize("raw_capsule", [False, True]) def test_three_library_query_planner(raw_capsule: bool): """Host, provider, and planner exchange a real non-empty plan over FFI.""" From 9fa9a83cad6e9d0d1afa001d8a5e4a4f9408e143 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 15:16:35 -0400 Subject: [PATCH 22/30] fix: stop reporting any getter TypeError as an outdated library call_capsule_getter rewrote every TypeError from a capsule getter into "Incompatible libraries ... Upgrade the library providing this object", and dropped the original. Only an arity mismatch means the library is out of date. An extension author whose own getter raised a TypeError -- a bad cast, a wrong argument to something it called -- was told the error was a version problem and lost the error that would have located it. The two are distinguishable without guessing at message text: an arity mismatch is raised by the call machinery before the getter's frame exists, so no frame unwinds and no traceback is attached, while an error from the body carries one. Verified to hold for both pure-Python and pyo3-compiled getters, which is the case that matters here since extension libraries are compiled. Also chains the original as __cause__ on the paths that do report an upgrade, so the arity error stays readable. Tests cover all three outcomes. Confirmed non-vacuous: dropping the traceback check fails only the inside-the-getter test, dropping set_cause fails only the upgrade test. Co-Authored-By: Claude Opus 5 (1M context) --- crates/util/src/lib.rs | 30 ++++++++++++++++----- python/tests/test_context.py | 51 ++++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 7 deletions(-) diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index 0fa908f5f..18ed68c64 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -239,14 +239,30 @@ fn call_capsule_getter<'py>( result.map_err(|err| { let py = obj.py(); - if session.is_some() && err.get_type(py).is(PyType::new::(py)) { - PyImportError::new_err(format!( - "Incompatible libraries. `{attr_name}` must accept the SessionContext it \ - is being installed on. Upgrade the library providing this object." - )) - } else { - err + if session.is_none() || !err.get_type(py).is(PyType::new::(py)) { + return err; } + + // Not every `TypeError` here means the getter refused the argument. One + // raised *inside* a correctly-signed getter would otherwise be reported + // as a version mismatch, sending an extension author to upgrade a + // library that is already correct. + // + // The two are distinguishable: an arity mismatch is raised by the call + // machinery before the getter's frame exists, so nothing unwinds and no + // traceback is attached. An error from the body unwinds that frame and + // carries one. + if err.traceback(py).is_some() { + return err; + } + + let import_err = PyImportError::new_err(format!( + "Incompatible libraries. `{attr_name}` must accept the SessionContext it \ + is being installed on. Upgrade the library providing this object." + )); + // Keep the original reachable as `__cause__` rather than discarding it. + import_err.set_cause(py, Some(err)); + import_err }) } diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 5ceb14e0e..95743efd2 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -737,6 +737,57 @@ def test_with_query_planner_rejects_wrong_capsule(ctx): ctx.with_query_planner(ctx.__datafusion_task_context_provider__()) +def test_pre_55_codec_signature_reports_an_upgrade(ctx): + """A getter that refuses the session is named, not left as a bare TypeError. + + Extension libraries implement these getters, so the pre-55.0.0 signature + is what an out-of-date one still has. The original error stays reachable + as ``__cause__`` rather than being replaced outright. + """ + + class PreSessionCodec: + def __datafusion_logical_extension_codec__(self): + msg = "should never be called" + raise AssertionError(msg) + + with pytest.raises(ImportError, match="__datafusion_logical_extension_codec__"): + ctx.with_logical_extension_codec(PreSessionCodec()) + + with pytest.raises(ImportError) as excinfo: + ctx.with_logical_extension_codec(PreSessionCodec()) + assert isinstance(excinfo.value.__cause__, TypeError) + assert "positional argument" in str(excinfo.value.__cause__) + + +def test_type_error_inside_a_getter_is_not_reported_as_an_upgrade(ctx): + """A correctly-signed getter's own TypeError must survive unchanged. + + Only the call machinery's arity error means the library is out of date. + Rewriting every TypeError would send an author debugging their own getter + off to upgrade a library that is already correct. + """ + + class RaisesTypeError: + def __datafusion_logical_extension_codec__(self, session): + msg = "bad cast inside the getter" + raise TypeError(msg) + + with pytest.raises(TypeError, match="bad cast inside the getter"): + ctx.with_logical_extension_codec(RaisesTypeError()) + + +def test_non_type_errors_from_a_getter_propagate(ctx): + """Anything that is not a TypeError was never a signature problem.""" + + class RaisesValueError: + def __datafusion_logical_extension_codec__(self, session): + msg = "something else entirely" + raise ValueError(msg) + + with pytest.raises(ValueError, match="something else entirely"): + ctx.with_logical_extension_codec(RaisesValueError()) + + def test_with_query_planner_capsule(ctx): capsule = ctx.__datafusion_query_planner__() get_name = ctypes.pythonapi.PyCapsule_GetName From 658e4eada16a364de75b9322cf5c4634b3529174 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 15:23:14 -0400 Subject: [PATCH 23/30] fix!: remove physical_codec_from_pycapsule The from_pycapsule! macros call the getter with no arguments. That is correct for __datafusion_physical_optimizer_rule__ and __datafusion_task_context_provider__, which take no session, but __datafusion_physical_extension_codec__ now takes the session it is being installed on, so this helper was the one member of the family left speaking the old protocol. Nothing in the tree called it, but datafusion-python-util is published by `cargo publish --workspace`, so it was still reachable. Against an updated codec it raised a bare TypeError, bypassing the ImportError that names the method. Against an outdated one it succeeded and produced a codec resolving names against the wrong session -- the silent failure the rest of this work exists to prevent. Removing it is a breaking change to that crate, but the crate already breaks this release: ffi_logical_codec_from_pycapsule gained its session parameter. A compile error pointing at the replacement beats a helper that quietly binds to nothing. Callers move to ffi_physical_codec_from_pycapsule, which passes the session, plus (&ffi).into() where an Arc is wanted -- what crates/core already does. Documents both helper changes in the 55.0.0 upgrade guide, which until now covered only the __datafusion_*__ method signatures and not the Rust helpers the same authors call. Co-Authored-By: Claude Opus 5 (1M context) --- crates/util/src/lib.rs | 14 ++++---- docs/source/user-guide/upgrade-guides.md | 41 ++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 8 deletions(-) diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index 18ed68c64..4e7255969 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -31,7 +31,6 @@ use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; use datafusion_ffi::query_planner::FFI_QueryPlanner; use datafusion_ffi::table_provider::FFI_TableProvider; -use datafusion_proto::physical_plan::PhysicalExtensionCodec; use pyo3::exceptions::{PyImportError, PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyType}; @@ -447,13 +446,12 @@ macro_rules! try_from_pycapsule { #[doc(hidden)] pub use pyo3; -from_pycapsule!( - physical_codec_from_pycapsule, - "datafusion_physical_extension_codec", - FFI_PhysicalExtensionCodec, - dyn PhysicalExtensionCodec -); - +// There is deliberately no `physical_codec_from_pycapsule` here. These macros +// call the getter with no arguments, which is right for the two hooks below but +// wrong for `__datafusion_physical_extension_codec__`, which takes the session +// it is being installed on. Use `ffi_physical_codec_from_pycapsule`, which +// passes the session, and convert with `(&ffi).into()` if you need an +// `Arc`. from_pycapsule!( physical_optimizer_rule_from_pycapsule, "datafusion_physical_optimizer_rule", diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 63129117b..22678c48b 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -73,6 +73,47 @@ It receives the session and takes both extension codecs from it, so a planner library never builds a `TaskContextProvider` at all. See the {ref}`ffi` guide for the full protocol. +### Changes to the `datafusion-python-util` crate + +Extension libraries written in Rust usually depend on the +`datafusion-python-util` crate for the helpers that read these capsules. Two of +those helpers changed, because the getter they call now takes the session. + +`ffi_logical_codec_from_pycapsule` takes a second argument. Pass `Some(session)` +when importing an object from another library, so its getter receives the +session it is being installed on. Pass `None` when the object *is* a session and +you are asking it for what it holds: + +```rust +// Before +let codec = ffi_logical_codec_from_pycapsule(obj)?; + +// After +let codec = ffi_logical_codec_from_pycapsule(obj, Some(session))?; +``` + +`physical_codec_from_pycapsule` has been **removed**. It called +`__datafusion_physical_extension_codec__` with no arguments, which no longer +matches the protocol, so against an updated codec it raised a bare `TypeError` +and against an outdated one it silently produced a codec bound to the wrong +session. Use `ffi_physical_codec_from_pycapsule`, which passes the session: + +```rust +// Before +let codec: Arc = physical_codec_from_pycapsule(&obj)?; + +// After +let ffi = ffi_physical_codec_from_pycapsule(obj, Some(session))?; +let codec: Arc = (&ffi).into(); +``` + +`physical_optimizer_rule_from_pycapsule` and `task_context_from_pycapsule` are +unchanged. Their hooks take no session. + +Calling a getter that still has the old signature now raises an `ImportError` +naming the method, with the original `TypeError` retained as its `__cause__`, +rather than a bare `TypeError`. + ## DataFusion 54.0.0 The `Config` class has been removed. It was a standalone wrapper around From a75c808b062c5b8b7c9a2fc8d1dc3690a2a3d9c4 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 15:29:29 -0400 Subject: [PATCH 24/30] fix: check the FFI major version on every importer that can Only ffi_query_planner_from_pycapsule validated the version a capsule reported. The codec and table provider importers dereference a foreign struct through the same `unsafe { data.as_ref() }` and were happy to accept one built against a different DataFusion. Extracts the planner's inline check into check_ffi_version and applies it to the logical codec, physical codec, and table provider importers as well. The helper is pub so extension libraries writing their own importers can use it. Two things the symmetry cannot reach, both now documented where someone would look: FFI_TaskContextProvider, FFI_TableProviderFactory, and FFI_ExtensionOptions carry no version field, so their importers cannot check. The from_pycapsule!/try_from_pycapsule! macros are #[macro_export] and generic over the FFI type, so requiring a version field there would break downstream users holding one of those three; they stay unchecked and their doc comment now says to call check_ffi_version directly. This is a diagnostic, not a soundness guarantee, and the helper says so: `version` is not the first field on any of these structs, so reading it already assumes the local layout. It turns the realistic failure -- a library compiled against a different DataFusion -- into a clear error instead of undefined behaviour on first use, which is what datafusion_ffi::version is documented to be for. Verified all four sites are wired by inverting the comparison and confirming each one fires from the test suites. Co-Authored-By: Claude Opus 5 (1M context) --- crates/util/src/lib.rs | 44 +++++++++++++++++++++++++++++++++++------- 1 file changed, 37 insertions(+), 7 deletions(-) diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index 4e7255969..ad82e3d51 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -175,6 +175,34 @@ pub fn validate_pycapsule(capsule: &Bound, name: &str) -> PyResult<() Ok(()) } +/// Reject an FFI struct built against a different major version of +/// `datafusion-ffi`. +/// +/// `found` comes from the struct's own `version` function pointer, which +/// reports the major version of the library that produced it. +/// +/// This is a diagnostic, not a soundness guarantee. Reading `version` out of +/// the struct already assumes the local field layout, and `version` is not the +/// first field on any of these types, so a sufficiently different layout can +/// fault before this ever runs. What it buys is a clear error for the case that +/// actually happens -- an extension library compiled against a different +/// DataFusion -- rather than undefined behaviour on first use, which is what +/// `datafusion_ffi::version` exists for. +/// +/// Not every FFI type carries a version. `FFI_TaskContextProvider`, +/// `FFI_TableProviderFactory`, and `FFI_ExtensionOptions` have no such field, +/// so their importers cannot check and are not expected to. +pub fn check_ffi_version(kind: &str, found: u64) -> PyResult<()> { + let expected = datafusion_ffi::version(); + if found != expected { + return Err(PyImportError::new_err(format!( + "Incompatible DataFusion {kind} major version {found}; expected {expected}. \ + Rebuild the library providing this object against a matching DataFusion." + ))); + } + Ok(()) +} + pub fn table_provider_from_pycapsule<'py>( mut obj: Bound<'py, PyAny>, session: Bound<'py, PyAny>, @@ -197,6 +225,7 @@ pub fn table_provider_from_pycapsule<'py>( .pointer_checked(Some(c"datafusion_table_provider"))? .cast(); let provider = unsafe { data.as_ref() }; + check_ffi_version("table provider", unsafe { (provider.version)() })?; let provider: Arc = provider.into(); Ok(Some(provider)) @@ -276,6 +305,7 @@ pub fn ffi_logical_codec_from_pycapsule<'py>( .pointer_checked(Some(c"datafusion_logical_extension_codec"))? .cast(); let codec = unsafe { data.as_ref() }; + check_ffi_version("logical extension codec", unsafe { (codec.version)() })?; Ok(codec.clone()) } @@ -292,6 +322,7 @@ pub fn ffi_physical_codec_from_pycapsule<'py>( .pointer_checked(Some(c"datafusion_physical_extension_codec"))? .cast(); let codec = unsafe { data.as_ref() }; + check_ffi_version("physical extension codec", unsafe { (codec.version)() })?; Ok(codec.clone()) } @@ -340,13 +371,7 @@ pub fn ffi_query_planner_from_pycapsule<'py>( .pointer_checked(Some(c"datafusion_query_planner"))? .cast(); let planner = unsafe { data.as_ref() }; - let planner_version = unsafe { (planner.version)() }; - let expected_version = datafusion_ffi::version(); - if planner_version != expected_version { - return Err(PyImportError::new_err(format!( - "Incompatible DataFusion query planner version {planner_version}; expected major version {expected_version}." - ))); - } + check_ffi_version("query planner", unsafe { (planner.version)() })?; Ok(planner.clone()) } @@ -367,6 +392,11 @@ pub fn create_physical_extension_capsule<'py>( /// Use this when `Arc<$output_type>: From<&$ffi_type>` (infallible /// conversion). For fallible conversions use [`try_from_pycapsule!`] /// instead. +/// +/// The generated extractor does not check the FFI major version, because not +/// every FFI type carries one. If `$ffi_type` has a `version` field, call +/// [`check_ffi_version`] on it yourself, as the hand-written extractors in this +/// crate do. #[macro_export] macro_rules! from_pycapsule { ($fn_name:ident, $capsule_name:literal, $ffi_type:ty, $output_type:ty) => { From 8d138f95a0d7298ba76002ef0089c2408a24bc3b Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 15:36:47 -0400 Subject: [PATCH 25/30] docs: say where to relax the FFI version check Exact equality is only right while datafusion_ffi::version tracks the crate's semver major, which it does today, so the number moves on every major release whether or not the ABI changed. If a version span later becomes compatible, a maintainer needs to know that this one body holds the whole policy -- callers pass a value and no decision -- and that relaxing it at a call site would reintroduce the split the helper was added to remove. Also records the likelier resolution: if the ABI is stable but version still follows the crate major, upstream's compatibility marker is wrong for every consumer, so the fix belongs there rather than in a local range policy. Co-Authored-By: Claude Opus 5 (1M context) --- crates/util/src/lib.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index ad82e3d51..f04f45052 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -192,6 +192,24 @@ pub fn validate_pycapsule(capsule: &Bound, name: &str) -> PyResult<() /// Not every FFI type carries a version. `FFI_TaskContextProvider`, /// `FFI_TableProviderFactory`, and `FFI_ExtensionOptions` have no such field, /// so their importers cannot check and are not expected to. +/// +/// # If the FFI ABI stabilizes +/// +/// Exact equality is the right test only while `datafusion_ffi::version` +/// tracks the DataFusion crate's semver major, which it does today +/// (`env!("CARGO_PKG_VERSION")`, `.major`). That number therefore moves on +/// every major release whether or not the ABI actually changed. +/// +/// Should a version span become compatible, **this function body is the only +/// thing to change** -- callers pass a `found` value and no policy. Relaxing it +/// at a call site instead would reintroduce the split this helper exists to +/// remove. +/// +/// The likelier fix is upstream, not here: if the ABI is stable but `version` +/// still follows the crate major, upstream's own compatibility marker is wrong +/// for every consumer, not just this one. Prefer waiting for +/// `datafusion_ffi::version` to reflect the real ABI over inventing a range +/// policy locally. pub fn check_ffi_version(kind: &str, found: u64) -> PyResult<()> { let expected = datafusion_ffi::version(); if found != expected { From 3eb86305f44f7f5d5627a1ba36cbd707efb05727 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 15:40:53 -0400 Subject: [PATCH 26/30] fix: route every capsule getter through call_capsule_getter The table provider and table function importers each carried their own copy of the TypeError-to-ImportError mapping, predating call_capsule_getter and never folded into it. Both therefore missed the correction it since received: they rewrote a TypeError raised inside a correctly-signed getter into "upgrade your library", and discarded the original. Three copies of one mapping, two of them stale, is the reason to have one. Both now call the shared helper, so they pick up the traceback discrimination and the __cause__ chain, and any later correction reaches all three by construction. Their messages named DataFusion 52.0.0. The shared message names the method that refused the argument instead, which points at the specific hook rather than a release, and the upgrade guide carries the version detail. call_capsule_getter is now pub, with a doc comment saying to use it rather than calling getattr directly. Tests cover both outcomes on both paths. Verified against the previous build that they are non-vacuous: before this change the raises-inside case produced the same misleading ImportError as the old-signature case. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/udtf.rs | 18 +++++++----------- crates/util/src/lib.rs | 19 ++++++------------- python/tests/test_context.py | 29 +++++++++++++++++++++++++++++ python/tests/test_udtf.py | 30 ++++++++++++++++++++++++++++++ 4 files changed, 72 insertions(+), 24 deletions(-) diff --git a/crates/core/src/udtf.rs b/crates/core/src/udtf.rs index cffa0c12a..51ea8fa4f 100644 --- a/crates/core/src/udtf.rs +++ b/crates/core/src/udtf.rs @@ -24,10 +24,10 @@ use datafusion::execution::context::SessionContext; use datafusion::execution::session_state::SessionState; use datafusion::logical_expr::Expr; use datafusion_ffi::udtf::FFI_TableFunction; +use datafusion_python_util::call_capsule_getter; use pyo3::IntoPyObjectExt; -use pyo3::exceptions::{PyImportError, PyTypeError}; use pyo3::prelude::*; -use pyo3::types::{PyCapsule, PyDict, PyTuple, PyType}; +use pyo3::types::{PyCapsule, PyDict, PyTuple}; use crate::context::PySessionContext; use crate::errors::{py_datafusion_err, to_datafusion_err}; @@ -76,15 +76,11 @@ impl PyTableFunction { Some(session) => session, None => PySessionContext::global_ctx()?.into_bound_py_any(py)?, }; - let capsule = func - .getattr("__datafusion_table_function__")? - .call1((session,)).map_err(|err| { - if err.get_type(py).is(PyType::new::(py)) { - PyImportError::new_err("Incompatible libraries. DataFusion 52.0.0 introduced an incompatible signature change for table functions. Either downgrade DataFusion or upgrade your function library.") - } else { - err - } - })?; + let capsule = call_capsule_getter( + func.clone(), + "__datafusion_table_function__", + Some(&session), + )?; let capsule = capsule.cast::()?; let data: NonNull = capsule .pointer_checked(Some(c"datafusion_table_function"))? diff --git a/crates/util/src/lib.rs b/crates/util/src/lib.rs index f04f45052..741a9d04c 100644 --- a/crates/util/src/lib.rs +++ b/crates/util/src/lib.rs @@ -225,18 +225,7 @@ pub fn table_provider_from_pycapsule<'py>( mut obj: Bound<'py, PyAny>, session: Bound<'py, PyAny>, ) -> PyResult>> { - if obj.hasattr("__datafusion_table_provider__")? { - obj = obj - .getattr("__datafusion_table_provider__")? - .call1((session,)).map_err(|err| { - let py = obj.py(); - if err.get_type(py).is(PyType::new::(py)) { - PyImportError::new_err("Incompatible libraries. DataFusion 52.0.0 introduced an incompatible signature change for table providers. Either downgrade DataFusion or upgrade your function library.") - } else { - err - } - })?; - } + obj = call_capsule_getter(obj, "__datafusion_table_provider__", Some(&session))?; if let Ok(capsule) = obj.cast::() { let data: NonNull = capsule @@ -268,7 +257,11 @@ pub fn create_logical_extension_capsule<'py>( /// of the session it is being installed on, instead of inventing one of its /// own. `None` is for the reverse direction, where `obj` *is* a session and is /// being asked for what it holds. -fn call_capsule_getter<'py>( +/// Every capsule getter must go through here rather than calling `getattr` +/// directly, so that the mapping from a refused argument to a useful error +/// lives in one place. Three importers previously each had their own copy and +/// each missed later corrections to it. +pub fn call_capsule_getter<'py>( obj: Bound<'py, PyAny>, attr_name: &str, session: Option<&Bound<'py, PyAny>>, diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 95743efd2..a8d5a1161 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1275,3 +1275,32 @@ def test_read_csv_with_options(tmp_path, as_read, global_ctx): read_csv_with_options_inner( tmp_path, csv_content, options, expected, as_read, global_ctx ) + + +def test_pre_52_table_provider_signature_reports_an_upgrade(ctx): + """The table provider hook reports an upgrade the same way codecs do. + + The 52.0.0 signature change added the session argument. This path had its + own copy of the error mapping and so missed later corrections to it. + """ + + class PreSessionProvider: + def __datafusion_table_provider__(self): + msg = "should never be called" + raise AssertionError(msg) + + with pytest.raises(ImportError, match="__datafusion_table_provider__") as excinfo: + ctx.register_table_provider("old_sig", PreSessionProvider()) + assert isinstance(excinfo.value.__cause__, TypeError) + + +def test_type_error_inside_a_table_provider_getter_propagates(ctx): + """A correctly-signed provider getter's own TypeError survives unchanged.""" + + class RaisesTypeError: + def __datafusion_table_provider__(self, session): + msg = "bad cast inside the getter" + raise TypeError(msg) + + with pytest.raises(TypeError, match="bad cast inside the getter"): + ctx.register_table_provider("raises", RaisesTypeError()) diff --git a/python/tests/test_udtf.py b/python/tests/test_udtf.py index dcb2bacc3..aa0599ffa 100644 --- a/python/tests/test_udtf.py +++ b/python/tests/test_udtf.py @@ -233,3 +233,33 @@ class FakeFFITableFunction: with pytest.raises(TypeError, match="FFI-exported table functions"): TableFunction("fake_ffi", fake, with_session=True) + + +def test_pre_52_table_function_signature_reports_an_upgrade() -> None: + """A getter that refuses the session is named, not left as a bare TypeError. + + The 52.0.0 signature change added the session argument. An out-of-date + library still has the old one, and the original error stays reachable as + ``__cause__``. + """ + + class PreSessionTableFunction: + def __datafusion_table_function__(self): + msg = "should never be called" + raise AssertionError(msg) + + with pytest.raises(ImportError, match="__datafusion_table_function__") as excinfo: + TableFunction("old_sig", PreSessionTableFunction(), None) + assert isinstance(excinfo.value.__cause__, TypeError) + + +def test_type_error_inside_a_table_function_getter_propagates() -> None: + """A correctly-signed getter's own TypeError must survive unchanged.""" + + class RaisesTypeError: + def __datafusion_table_function__(self, session): + msg = "bad cast inside the getter" + raise TypeError(msg) + + with pytest.raises(TypeError, match="bad cast inside the getter"): + TableFunction("raises", RaisesTypeError(), None) From 5d68f7b6b951e2279c503912448f69245f90918a Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 15:44:47 -0400 Subject: [PATCH 27/30] docs: explain the dropped & in the codec migration snippet The before and after snippets pass the task context provider differently, by reference in one and by value in the other, with nothing saying why. Read as a diff it looks like a typo in one of them, and a reader correcting it would be puzzled when both versions compile. Both are valid: the parameter is impl Into, which is satisfied by &Arc and by FFI_TaskContextProvider itself, and the latter is what ffi_task_context_provider_from_pycapsule returns. The argument changes because the provider now comes from the session instead of a field, which is the point of the migration. The contributor guide shows only the post-migration form, so it needs no equivalent note. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/user-guide/upgrade-guides.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/source/user-guide/upgrade-guides.md b/docs/source/user-guide/upgrade-guides.md index 22678c48b..1ded21921 100644 --- a/docs/source/user-guide/upgrade-guides.md +++ b/docs/source/user-guide/upgrade-guides.md @@ -56,6 +56,13 @@ fn __datafusion_physical_extension_codec__<'py>( } ``` +The dropped `&` on the last argument is not a typo. That parameter is +`impl Into`, so it accepts either an +`&Arc`, as before, or an `FFI_TaskContextProvider`, +which is what `ffi_task_context_provider_from_pycapsule` hands back. Both forms +compile; the argument changes because the provider now comes from the session +rather than from a field. + A codec that keeps its own `SessionContext` still compiles, but its decode callbacks resolve names against that empty session instead of the one running the query, so a function registered with `SessionContext.register_udf` is not From 5448e8da4411a17e74ff1f36b04dc68b59717431 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 15:47:14 -0400 Subject: [PATCH 28/30] docs: spell out the token registry lifecycle in the example codecs The README already describes these as one-shot registries that consume each token during decoding, but the source comments did not, and the source is what someone reuses the pattern from. The existing comment warned that the registry is process-local without saying that a decode removes its entry, which is the constraint most likely to bite. Documents both consequences on the registry accessors, where the mechanism lives, with a pointer from each struct doc: - Decode consumes the token, so the same encoded bytes cannot be decoded twice. Fine here because every plan is encoded immediately before the one decode that consumes it, but it rules out replaying a stored plan, retrying a decode, or fanning one plan out to several readers. - An encode that never reaches a decoder leaks for the life of the process. Normal operation does not: encode and decode counts balance exactly across repeated queries, which is what makes remove-on-decode the right trade here rather than a leak on every call. Comments only. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/logical_extension_codec.rs | 22 +++++++++++++++++++ .../src/physical_extension_codec.rs | 9 ++++++++ 2 files changed, 31 insertions(+) diff --git a/examples/datafusion-ffi-example/src/logical_extension_codec.rs b/examples/datafusion-ffi-example/src/logical_extension_codec.rs index a127407f2..1fcaaef4c 100644 --- a/examples/datafusion-ffi-example/src/logical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/logical_extension_codec.rs @@ -38,6 +38,24 @@ const TABLE_PROVIDER_TOKEN: &[u8] = b"DFPYEXTP"; static NEXT_TABLE_PROVIDER_ID: AtomicU64 = AtomicU64::new(1); static TABLE_PROVIDERS: OnceLock>>> = OnceLock::new(); +/// Hands a provider to another library in this process by token. +/// +/// Encoding inserts, decoding removes. Two consequences worth knowing before +/// copying this: +/// +/// - **Decode consumes the token.** Decoding the same encoded bytes twice +/// fails the second time with `Unknown ... table provider token`. That is +/// fine here because every plan is encoded immediately before the single +/// decode that consumes it, but it rules out anything that replays a stored +/// plan, retries a decode, or fans one encoded plan out to several readers. +/// - **An encode that is never decoded leaks.** Nothing expires entries, so a +/// plan that fails to reach its decoder keeps its provider alive for the +/// life of the process. +/// +/// Both are acceptable for an example whose job is to show that Rust type +/// identity survives a trip through two other libraries. Neither is acceptable +/// in a real codec, which should encode metadata sufficient to rebuild the +/// provider rather than parking the object here. fn table_providers() -> &'static Mutex>> { TABLE_PROVIDERS.get_or_init(|| Mutex::new(HashMap::new())) } @@ -62,6 +80,10 @@ pub(crate) struct CallCounters { /// example of preserving Rust type identity across three loaded libraries, not a /// network serialization format. Production libraries should encode reconstructible /// provider metadata rather than retaining objects in a global registry. +/// +/// See [`table_providers`] for the token lifecycle, which is narrower than it +/// looks: a decode consumes its token, so the same encoded plan cannot be +/// decoded twice. struct CountingLogicalExtensionCodec { inner: DefaultLogicalExtensionCodec, counters: Arc, diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs index fe903a68b..f9e96382e 100644 --- a/examples/datafusion-ffi-example/src/physical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -40,6 +40,11 @@ const EXECUTION_PLAN_TOKEN: &[u8] = b"DFPYEXEP"; static NEXT_EXECUTION_PLAN_ID: AtomicU64 = AtomicU64::new(1); static EXECUTION_PLANS: OnceLock>>> = OnceLock::new(); +/// Execution-plan counterpart of the logical codec's provider registry, with +/// the same lifecycle: encoding inserts, decoding removes, so a decode +/// consumes its token and an encode that is never decoded leaks. See +/// [`crate::logical_extension_codec`] for why that is acceptable here and not +/// in a real codec. fn execution_plans() -> &'static Mutex>> { EXECUTION_PLANS.get_or_init(|| Mutex::new(HashMap::new())) } @@ -63,6 +68,10 @@ pub(crate) struct PhysicalCallCounters { /// Provider-owned memory scan plans use a same-process token registry so the /// owning cdylib can restore their concrete Rust type after the plan travels /// through the independent query-planner and datafusion-python libraries. +/// +/// See [`execution_plans`] for the token lifecycle, which is narrower than it +/// looks: a decode consumes its token, so the same encoded plan cannot be +/// decoded twice. struct CountingPhysicalExtensionCodec { inner: DefaultPhysicalExtensionCodec, counters: Arc, From ca3264e129c1d899cf23ecefc96504b3389ea8c9 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 15:58:44 -0400 Subject: [PATCH 29/30] fix: resolve a planner fallback when it is installed, not constructed MyQueryPlanner::new imported its fallback immediately, with no session to pass, so the fallback's getter was called with no arguments. That works for a SessionContext, whose getter takes the session optionally, and for a raw capsule, which has no getter at all. It fails for another foreign planner, which implements the same protocol this type does and requires the argument -- and layering on another planner is the case a distributed engine actually needs. The docstring claimed fallback "takes anything exporting __datafusion_query_planner__", which was not true. Holds the Python object instead and imports it in __datafusion_query_planner__, where the session is in hand and can be forwarded. All three fallback kinds now work. Deferring also removes a footgun rather than adding one. Passing a SessionContext now delegates to whichever planner it holds at install time, and since with_query_planner calls the getter before installing, the context still reports its previous planner, so wrapping a context in a planner installed on that same context does not recurse. Arc> rather than Py because pyo3 0.29 gates Py: Clone behind the py-clone feature, and this type derives Clone. Matches how PythonTableFunctionCallable holds its callable. Co-Authored-By: Claude Opus 5 (1M context) --- .../_test_three_library_query_planner.py | 40 ++++++++++++++ .../src/planner.rs | 52 +++++++++++++------ 2 files changed, 76 insertions(+), 16 deletions(-) diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index b1b5b21c3..2de48a8c9 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -358,6 +358,46 @@ def test_planner_layers_on_the_session_planner(): assert physical_codec.execution_plan_decode_calls() > 0 +def test_a_planner_can_fall_back_to_another_planner_library(): + """A fallback may be another foreign planner, not only a session. + + The fallback is imported when this planner is installed rather than when + it is constructed, so its own getter receives the session. Importing it at + construction time would mean calling that getter with no session, which + only a ``SessionContext`` or a raw capsule tolerates -- and layering on + another planner is the case a distributed engine actually needs. + """ + ctx, logical_codec, physical_codec = configured_context(max_rows=3) + inner = MyQueryPlanner() + outer = MyQueryPlanner(fallback=inner) + ctx = ctx.with_query_planner(outer) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert outer.plan_calls() > 0 + assert outer.used_fallback() + # The delegation reached the inner planner rather than stopping at the + # default physical planner. + assert inner.plan_calls() > 0 + assert logical_codec.table_provider_decode_calls() > 0 + assert physical_codec.execution_plan_decode_calls() > 0 + + +def test_a_session_fallback_delegates_to_its_installed_planner(): + """Passing a SessionContext delegates to whatever planner it holds.""" + ctx, _logical_codec, _physical_codec = configured_context(max_rows=3) + first = MyQueryPlanner() + ctx = ctx.with_query_planner(first) + + second = MyQueryPlanner(fallback=ctx) + ctx = ctx.with_query_planner(second) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert second.used_fallback() + assert first.plan_calls() > 0 + + def test_second_planner_replaces_the_first(): """A session holds exactly one planner, so installing another replaces it.""" ctx, _logical_codec, _physical_codec = configured_context(max_rows=2) diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index 7928cedf4..e34e9812c 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -203,30 +203,38 @@ impl QueryPlanner for DistributedQueryPlanner { #[derive(Debug, Default, Clone)] pub(crate) struct MyQueryPlanner { observations: Arc, - fallback: Option>, + /// Held as the Python object rather than an imported planner, and resolved + /// in `__datafusion_query_planner__` where a session is in hand. + /// + /// Importing it here would mean calling its getter with no session, which + /// only a `SessionContext` or a raw capsule accepts. Another foreign + /// planner -- the case that matters, since layering is the whole point of + /// a fallback -- implements the same protocol this type does and requires + /// the argument. + fallback: Option>>, } #[pymethods] impl MyQueryPlanner { /// Build a planner, optionally layered on top of an existing one. /// - /// `fallback` takes anything exporting `__datafusion_query_planner__`, - /// including a `SessionContext`. Capture it *before* installing this - /// planner on that context, or the capsule will describe this planner and - /// planning will recurse. + /// `fallback` takes anything exporting `__datafusion_query_planner__`: + /// another planner library, a `SessionContext`, or a raw capsule. It is + /// imported when this planner is installed, not here, so that the session + /// can be handed to its getter. + /// + /// Passing a `SessionContext` delegates to whichever planner that context + /// holds at install time. If you instead capture a capsule with + /// `ctx.__datafusion_query_planner__()`, capture it *before* installing + /// this planner on that context, or the capsule will describe this planner + /// and planning will recurse. #[new] #[pyo3(signature = (fallback=None))] - fn new(fallback: Option>) -> PyResult { - let fallback = fallback - .map(|planner| { - ffi_query_planner_from_pycapsule(&planner, None) - .map(|ffi| -> Arc { (&ffi).into() }) - }) - .transpose()?; - Ok(Self { - fallback, + fn new(fallback: Option>) -> Self { + Self { + fallback: fallback.map(|obj| Arc::new(obj.unbind())), ..Self::default() - }) + } } fn used_fallback(&self) -> bool { @@ -264,9 +272,21 @@ impl MyQueryPlanner { py: Python<'py>, session: Bound<'py, PyAny>, ) -> PyResult> { + // Resolved here rather than in `new` so the fallback's own getter + // receives the session, which is what the protocol requires of every + // implementation other than a `SessionContext`. + let fallback = self + .fallback + .as_ref() + .map(|planner| { + ffi_query_planner_from_pycapsule(planner.bind(py), Some(&session)) + .map(|ffi| -> Arc { (&ffi).into() }) + }) + .transpose()?; + let planner: Arc = Arc::new(DistributedQueryPlanner { observations: Arc::clone(&self.observations), - fallback: self.fallback.clone(), + fallback, }); let logical_codec = ffi_logical_codec_from_pycapsule(session.clone(), None)?; let physical_codec = ffi_physical_codec_from_pycapsule(session, None)?; From a2a2bb20e0da0c47dcb2f8d8d5c4049fea46c3fd Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 27 Aug 2026 16:33:26 -0400 Subject: [PATCH 30/30] fix: accumulate planner observations instead of overwriting them foreign_session, foreign_provider, and foreign_plan were written with store, so each one described only the most recent plan. Their accessors are named foreign_*_observed, which asks whether the thing was ever seen, and the tests assert them after running more than one query. The existing tests passed by luck. Reproduced: after scanning a foreign provider and then running SELECT 1, foreign_provider_observed goes from True back to False. Writes them with fetch_or so a later plan cannot retract what an earlier one observed. plan_calls already accumulated, used_fallback only ever stores true so it was already cumulative, and last_max_rows is deliberately last-wins as its name says. Documents that split on the struct, since it is the kind of thing that gets "tidied" back. Confirmed non-vacuous: with store restored, exactly the new test fails and the other 22 pass. Co-Authored-By: Claude Opus 5 (1M context) --- .../_test_three_library_query_planner.py | 30 +++++++++++++++++++ .../src/planner.rs | 21 +++++++++++-- 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index 2de48a8c9..9a5d78c44 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -398,6 +398,36 @@ def test_a_session_fallback_delegates_to_its_installed_planner(): assert first.plan_calls() > 0 +def test_observations_accumulate_across_queries(): + """A later plain query must not retract what an earlier query observed. + + The ``*_observed`` accessors answer "was this ever seen". They are written + with ``fetch_or`` rather than ``store`` so a query that touches no foreign + object cannot clear a flag an earlier one set. Written with ``store``, + ``SELECT 1`` here clears ``foreign_provider_observed``, and every other + test asserting these flags after more than one query is a coincidence away + from failing. + """ + ctx, _logical_codec, _physical_codec = configured_context(max_rows=3) + planner = MyQueryPlanner() + ctx = ctx.with_query_planner(planner) + + ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert planner.foreign_session_observed() + assert planner.foreign_provider_observed() + assert planner.foreign_plan_observed() + + # Touches no table, so this plan has no foreign provider of its own. + ctx.sql("SELECT 1").collect() + assert planner.foreign_session_observed() + assert planner.foreign_provider_observed() + assert planner.foreign_plan_observed() + + # `last_max_rows` is deliberately not cumulative; it reports the last plan. + assert planner.last_max_rows() == 3 + assert planner.plan_calls() >= 2 + + def test_second_planner_replaces_the_first(): """A session holds exactly one planner, so installing another replaces it.""" ctx, _logical_codec, _physical_codec = configured_context(max_rows=2) diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index e34e9812c..67262e39c 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -41,6 +41,16 @@ use pyo3::types::PyCapsule; use crate::config::MyPlannerConfig; +/// What the planner saw, accumulated across every call rather than reset each +/// time. +/// +/// Two kinds of field here, and mixing them up is easy. `last_max_rows` +/// reports the most recent value, as its name says. Everything else is +/// cumulative: a count, or a "did this ever happen" flag written with +/// `fetch_or` so a later plan cannot retract an earlier observation. Tests +/// assert after running more than one query, so a flag that only described the +/// most recent plan would be answering a different question than the one its +/// accessor name asks. #[derive(Default)] struct PlannerObservations { plan_calls: AtomicUsize, @@ -48,6 +58,7 @@ struct PlannerObservations { foreign_session: AtomicBool, foreign_provider: AtomicBool, foreign_plan: AtomicBool, + /// Only ever set to `true`, so it is already cumulative. used_fallback: AtomicBool, } @@ -156,10 +167,14 @@ impl QueryPlanner for DistributedQueryPlanner { session: &dyn Session, ) -> datafusion::common::Result> { self.observations.plan_calls.fetch_add(1, Ordering::SeqCst); + // `fetch_or`, not `store`: these answer "was this ever seen", so a + // later plan that happens not to touch a foreign object must not + // retract what an earlier one observed. A bare `SELECT 1` after a + // scan of a foreign provider would otherwise clear the flag. self.observations .foreign_session - .store(session.as_any().is::(), Ordering::SeqCst); - self.observations.foreign_provider.store( + .fetch_or(session.as_any().is::(), Ordering::SeqCst); + self.observations.foreign_provider.fetch_or( logical_plan_has_foreign_provider(logical_plan), Ordering::SeqCst, ); @@ -184,7 +199,7 @@ impl QueryPlanner for DistributedQueryPlanner { }; self.observations .foreign_plan - .store(physical_plan_has_foreign_plan(&plan), Ordering::SeqCst); + .fetch_or(physical_plan_has_foreign_plan(&plan), Ordering::SeqCst); Ok(Arc::new(GlobalLimitExec::new( plan,