From 9ba147c1697a0cf8a2f4d09e618b86064fbcb2b7 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 31 Jul 2026 15:38:53 -0400 Subject: [PATCH 01/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] Empty commit to trigger CI