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 8e02881b545b20066e6b6e7fa99ff32a7f9a252e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 10:34:45 -0400 Subject: [PATCH 04/10] Make extension codecs composable Installing a logical or physical extension codec now prepends it to a codec chain instead of replacing the prior codec. The most recently installed codec is consulted first, falling through codec by codec to the default codec. This lets multiple independent extension libraries install codecs on the same session, and removes the codec registration ordering requirement between libraries. Chain dispatch treats a codec error as "not mine". Encoding runs each codec against a scratch buffer so failed attempts leave no partial bytes, and treats Ok-with-no-bytes (encode by name) as no opinion so later codecs still get a chance. When every codec fails, the errors are aggregated so the owning codec's diagnostic is not masked by the default codec's generic error. Also preserves the python_udf_inlining setting when installing a codec; previously it was silently reset to enabled. Documents the remaining planner constraint: a session holds one query planner, layering is explicit via fallback capsules, and codecs must be installed before exporting or chaining planners because a planner capsule captures the codecs at export time. Co-Authored-By: Claude Fable 5 --- crates/core/src/codec.rs | 444 +++++++++++++++--- crates/core/src/context.rs | 16 +- docs/source/contributor-guide/ffi.md | 58 ++- examples/datafusion-ffi-example/README.md | 2 +- .../tests/_test_logical_extension_codec.py | 27 +- .../tests/_test_physical_extension_codec.py | 23 + .../README.md | 2 +- .../_test_three_library_query_planner.py | 23 + python/datafusion/context.py | 26 +- 9 files changed, 552 insertions(+), 69 deletions(-) diff --git a/crates/core/src/codec.rs b/crates/core/src/codec.rs index 26853e69f..6dd38ae46 100644 --- a/crates/core/src/codec.rs +++ b/crates/core/src/codec.rs @@ -29,16 +29,16 @@ //! //! [`PythonLogicalCodec`] is the [`LogicalExtensionCodec`] that //! datafusion-python parks on every `SessionContext`. It wraps a -//! user-supplied (or default) inner codec and adds Python-aware -//! in-band encoding on top: when the encoder sees a Python-defined -//! UDF, the codec cloudpickles the callable + signature into the -//! `fun_definition` proto field; when the decoder sees a payload it -//! produced, it reconstructs the UDF from the bytes alone — no -//! pre-registration on the receiver. UDFs the codec does not -//! recognise are delegated to `inner`, which is typically -//! `DefaultLogicalExtensionCodec` but may be a downstream-supplied -//! FFI codec installed via -//! `SessionContext.with_logical_extension_codec(...)`. +//! chain of composable codecs and adds Python-aware in-band encoding +//! on top: when the encoder sees a Python-defined UDF, the codec +//! cloudpickles the callable + signature into the `fun_definition` +//! proto field; when the decoder sees a payload it produced, it +//! reconstructs the UDF from the bytes alone — no pre-registration on +//! the receiver. Everything the codec does not recognise is delegated +//! to the chain: each downstream FFI codec installed via +//! `SessionContext.with_logical_extension_codec(...)` is consulted in +//! most-recently-installed-first order, with +//! `DefaultLogicalExtensionCodec` as the terminal fallback. //! //! [`PythonPhysicalCodec`] is the symmetric wrapper around //! [`PhysicalExtensionCodec`]. Logical and physical layers each have @@ -58,7 +58,7 @@ //! actionable error instead of an opaque `marshal` failure on load //! (cloudpickle payloads are not portable across Python minor //! versions). Dispatch precedence on decode: **family match + -//! supported version + matching Python version → `inner` codec → +//! supported version + matching Python version → codec chain → //! caller's `FunctionRegistry` fallback.** //! //! ## Wire-format family registry @@ -81,10 +81,11 @@ //! for an older shape. //! //! Downstream FFI codecs should pick non-colliding family prefixes -//! (use a `DF` namespace plus a crate-specific suffix). The codec -//! implementations in this module currently delegate every method to -//! `inner`; the encoder/decoder hooks for each kind are added as the -//! corresponding Python-side type becomes serializable. +//! (use a `DF` namespace plus a crate-specific suffix) and return an +//! error for payloads and objects they do not own — that error is the +//! chain's "not mine" signal, letting the next codec take a turn. A +//! codec that answers `Ok` for objects outside its family shadows +//! every codec installed before it. use std::sync::Arc; @@ -167,7 +168,7 @@ fn write_wire_header(buf: &mut Vec, family: &[u8], py_version: (u8, u8)) { /// Inspect the framing on `buf`. /// /// * `Ok(None)` — `buf` does not carry `family`. The caller should -/// delegate to its `inner` codec. +/// delegate to its codec chain. /// * `Ok(Some(payload))` — `buf` carries `family` at a version this /// build accepts and a Python `(major, minor)` matching /// `expected_py`; `payload` is the cloudpickle blob. @@ -223,32 +224,129 @@ fn strip_wire_header<'a>( Ok(Some(&buf[py_minor_idx + 1..])) } +/// Run `f` against each codec in `chain`, returning the first `Ok`. +/// +/// A codec signals "not mine" by returning an error, so the chain +/// keeps trying until a codec succeeds. When every codec fails and the +/// chain has more than one entry, the errors are aggregated into a +/// single message — returning only the last error would surface the +/// terminal `Default*ExtensionCodec` "not provided" message and mask +/// the more specific diagnostic from an installed codec (e.g. a +/// corrupt-token error from the codec that owns the payload family). +fn chain_try(chain: &[Arc], what: &str, f: impl Fn(&C) -> Result) -> Result { + let mut errors: Vec = Vec::new(); + for codec in chain { + match f(codec) { + Ok(value) => return Ok(value), + Err(err) => errors.push(err), + } + } + Err(aggregate_chain_errors(what, errors)) +} + +/// Collapse per-codec failures into one error. A single failure is +/// returned as-is so the one-codec (default-only) chain behaves +/// exactly like the pre-chain implementation. +fn aggregate_chain_errors( + what: &str, + mut errors: Vec, +) -> datafusion::error::DataFusionError { + match errors.len() { + 0 => datafusion::error::DataFusionError::Internal(format!( + "Empty extension codec chain while handling {what}" + )), + 1 => errors.swap_remove(0), + _ => { + let joined = errors + .iter() + .map(|err| err.to_string()) + .collect::>() + .join("; "); + datafusion::error::DataFusionError::Execution(format!( + "None of the {} composed extension codecs handled {what}: {joined}", + errors.len() + )) + } + } +} + +/// Encode variant of [`chain_try`] for methods that write into a +/// caller-provided buffer. +/// +/// Each codec encodes into a scratch buffer so a failed attempt cannot +/// leave partial bytes behind. `Ok` with bytes written commits those +/// bytes and ends the chain. `Ok` with an empty buffer is treated as +/// "no opinion" — the standard `Default*ExtensionCodec` behavior of +/// encoding a UDF by name writes nothing — so later codecs still get a +/// chance to emit a richer payload. If no codec writes bytes but at +/// least one returned `Ok`, the overall result is `Ok` with nothing +/// written (encode by name). +fn chain_encode( + chain: &[Arc], + buf: &mut Vec, + what: &str, + f: impl Fn(&C, &mut Vec) -> Result<()>, +) -> Result<()> { + let mut saw_empty_ok = false; + let mut errors: Vec = Vec::new(); + for codec in chain { + let mut scratch = Vec::new(); + match f(codec, &mut scratch) { + Ok(()) if !scratch.is_empty() => { + buf.extend_from_slice(&scratch); + return Ok(()); + } + Ok(()) => saw_empty_ok = true, + Err(err) => errors.push(err), + } + } + if saw_empty_ok { + return Ok(()); + } + Err(aggregate_chain_errors(what, errors)) +} + /// `LogicalExtensionCodec` parked on every `SessionContext`. Holds /// the Python-aware encoding hooks for logical-layer types /// (`LogicalPlan`, `Expr`) and delegates everything it does not -/// handle to the composable `inner` codec — typically -/// `DefaultLogicalExtensionCodec`, or a downstream FFI codec -/// installed via `SessionContext.with_logical_extension_codec(...)`. +/// handle to a chain of composable codecs. The chain starts as just +/// `DefaultLogicalExtensionCodec`; each downstream FFI codec installed +/// via `SessionContext.with_logical_extension_codec(...)` is prepended, +/// so the most recently installed codec is consulted first and the +/// default codec always runs last. +/// +/// Chain dispatch relies on each codec recognizing its own payloads +/// (distinct family prefixes — see the module docs) and returning an +/// error for everything else so the next codec gets a chance. /// /// 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)] +#[derive(Debug, Clone)] pub struct PythonLogicalCodec { - inner: Arc, + chain: Vec>, python_udf_inlining: bool, } impl PythonLogicalCodec { pub fn new(inner: Arc) -> Self { Self { - inner, + chain: vec![inner], python_udf_inlining: true, } } - pub fn inner(&self) -> &Arc { - &self.inner + /// Return a copy of this codec with `codec` prepended to the + /// chain, preserving the Python-UDF-inlining setting. The new + /// codec is consulted before every previously installed codec. + pub fn with_additional_codec(&self, codec: Arc) -> Self { + let mut chain = Vec::with_capacity(self.chain.len() + 1); + chain.push(codec); + chain.extend(self.chain.iter().map(Arc::clone)); + Self { + chain, + python_udf_inlining: self.python_udf_inlining, + } } /// Toggle inline encoding of Python UDFs. See @@ -289,11 +387,18 @@ impl LogicalExtensionCodec for PythonLogicalCodec { inputs: &[LogicalPlan], ctx: &TaskContext, ) -> Result { - self.inner.try_decode(buf, inputs, ctx) + chain_try(&self.chain, "an extension logical plan node", |codec| { + codec.try_decode(buf, inputs, ctx) + }) } fn try_encode(&self, node: &Extension, buf: &mut Vec) -> Result<()> { - self.inner.try_encode(node, buf) + chain_encode( + &self.chain, + buf, + "an extension logical plan node", + |codec, buf| codec.try_encode(node, buf), + ) } fn try_decode_table_provider( @@ -303,8 +408,9 @@ impl LogicalExtensionCodec for PythonLogicalCodec { schema: SchemaRef, ctx: &TaskContext, ) -> Result> { - self.inner - .try_decode_table_provider(buf, table_ref, schema, ctx) + chain_try(&self.chain, "a table provider", |codec| { + codec.try_decode_table_provider(buf, table_ref, Arc::clone(&schema), ctx) + }) } fn try_encode_table_provider( @@ -313,7 +419,9 @@ impl LogicalExtensionCodec for PythonLogicalCodec { node: Arc, buf: &mut Vec, ) -> Result<()> { - self.inner.try_encode_table_provider(table_ref, node, buf) + chain_encode(&self.chain, buf, "a table provider", |codec, buf| { + codec.try_encode_table_provider(table_ref, Arc::clone(&node), buf) + }) } fn try_decode_file_format( @@ -321,7 +429,9 @@ impl LogicalExtensionCodec for PythonLogicalCodec { buf: &[u8], ctx: &TaskContext, ) -> Result> { - self.inner.try_decode_file_format(buf, ctx) + chain_try(&self.chain, "a file format", |codec| { + codec.try_decode_file_format(buf, ctx) + }) } fn try_encode_file_format( @@ -329,14 +439,18 @@ impl LogicalExtensionCodec for PythonLogicalCodec { buf: &mut Vec, node: Arc, ) -> Result<()> { - self.inner.try_encode_file_format(buf, node) + chain_encode(&self.chain, buf, "a file format", |codec, buf| { + codec.try_encode_file_format(buf, Arc::clone(&node)) + }) } fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? { return Ok(()); } - self.inner.try_encode_udf(node, buf) + chain_encode(&self.chain, buf, "a scalar UDF", |codec, buf| { + codec.try_encode_udf(node, buf) + }) } fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { @@ -347,14 +461,18 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } else { refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?; } - self.inner.try_decode_udf(name, buf) + chain_try(&self.chain, "a scalar UDF", |codec| { + codec.try_decode_udf(name, buf) + }) } fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udaf(node, buf)? { return Ok(()); } - self.inner.try_encode_udaf(node, buf) + chain_encode(&self.chain, buf, "an aggregate UDF", |codec, buf| { + codec.try_encode_udaf(node, buf) + }) } fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { @@ -365,14 +483,18 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } else { refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?; } - self.inner.try_decode_udaf(name, buf) + chain_try(&self.chain, "an aggregate UDF", |codec| { + codec.try_decode_udaf(name, buf) + }) } fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udwf(node, buf)? { return Ok(()); } - self.inner.try_encode_udwf(node, buf) + chain_encode(&self.chain, buf, "a window UDF", |codec, buf| { + codec.try_encode_udwf(node, buf) + }) } fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { @@ -383,13 +505,15 @@ impl LogicalExtensionCodec for PythonLogicalCodec { } else { refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?; } - self.inner.try_decode_udwf(name, buf) + chain_try(&self.chain, "a window UDF", |codec| { + codec.try_decode_udwf(name, buf) + }) } } /// Strict-mode gate: if `buf` is a well-framed inline payload for /// `family`, return the strict-refusal error; otherwise return -/// `Ok(())` so the caller can delegate to its `inner` codec. +/// `Ok(())` so the caller can delegate to its codec chain. /// /// Routing through [`read_framed_payload`] (rather than a bare /// `starts_with` probe) means malformed inline bytes — wrong @@ -434,7 +558,8 @@ fn refuse_inline_payload(kind: &str, name: &str) -> datafusion::error::DataFusio /// `PhysicalExtensionCodec` mirror of [`PythonLogicalCodec`] parked /// on the same `SessionContext`. Carries the Python-aware encoding /// hooks for physical-layer types (`ExecutionPlan`, `PhysicalExpr`) -/// and delegates the rest to `inner`. +/// and delegates the rest to the composable codec chain (see +/// [`PythonLogicalCodec`] for chain ordering and dispatch rules). /// /// The `PhysicalExtensionCodec` trait has its own `try_encode_udf` /// / `try_decode_udf` pair distinct from the logical one, so a @@ -443,22 +568,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)] +#[derive(Debug, Clone)] pub struct PythonPhysicalCodec { - inner: Arc, + chain: Vec>, python_udf_inlining: bool, } impl PythonPhysicalCodec { pub fn new(inner: Arc) -> Self { Self { - inner, + chain: vec![inner], python_udf_inlining: true, } } - pub fn inner(&self) -> &Arc { - &self.inner + /// Return a copy of this codec with `codec` prepended to the + /// chain, preserving the Python-UDF-inlining setting. The new + /// codec is consulted before every previously installed codec. + pub fn with_additional_codec(&self, codec: Arc) -> Self { + let mut chain = Vec::with_capacity(self.chain.len() + 1); + chain.push(codec); + chain.extend(self.chain.iter().map(Arc::clone)); + Self { + chain, + python_udf_inlining: self.python_udf_inlining, + } } /// Toggle inline encoding of Python UDFs on this physical codec. @@ -489,7 +623,9 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { ctx: &TaskContext, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - self.inner.try_decode(buf, inputs, ctx, proto_converter) + chain_try(&self.chain, "an execution plan", |codec| { + codec.try_decode(buf, inputs, ctx, proto_converter) + }) } fn try_encode( @@ -498,14 +634,18 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { buf: &mut Vec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { - self.inner.try_encode(node, buf, proto_converter) + chain_encode(&self.chain, buf, "an execution plan", |codec, buf| { + codec.try_encode(Arc::clone(&node), buf, proto_converter) + }) } fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_scalar_udf(node, buf)? { return Ok(()); } - self.inner.try_encode_udf(node, buf) + chain_encode(&self.chain, buf, "a scalar UDF", |codec, buf| { + codec.try_encode_udf(node, buf) + }) } fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { @@ -516,7 +656,9 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } else { refuse_if_inline(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", name)?; } - self.inner.try_decode_udf(name, buf) + chain_try(&self.chain, "a scalar UDF", |codec| { + codec.try_decode_udf(name, buf) + }) } fn try_encode_expr( @@ -525,7 +667,9 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { buf: &mut Vec, ctx: &PhysicalExprEncodeCtx<'_>, ) -> Result<()> { - self.inner.try_encode_expr(node, buf, ctx) + chain_encode(&self.chain, buf, "a physical expression", |codec, buf| { + codec.try_encode_expr(node, buf, ctx) + }) } fn try_decode_expr( @@ -534,14 +678,18 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { inputs: &[Arc], ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { - self.inner.try_decode_expr(buf, inputs, ctx) + chain_try(&self.chain, "a physical expression", |codec| { + codec.try_decode_expr(buf, inputs, ctx) + }) } fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udaf(node, buf)? { return Ok(()); } - self.inner.try_encode_udaf(node, buf) + chain_encode(&self.chain, buf, "an aggregate UDF", |codec, buf| { + codec.try_encode_udaf(node, buf) + }) } fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { @@ -552,14 +700,18 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } else { refuse_if_inline(buf, PY_AGG_UDF_FAMILY, "aggregate UDF", name)?; } - self.inner.try_decode_udaf(name, buf) + chain_try(&self.chain, "an aggregate UDF", |codec| { + codec.try_decode_udaf(name, buf) + }) } fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { if self.python_udf_inlining && try_encode_python_udwf(node, buf)? { return Ok(()); } - self.inner.try_encode_udwf(node, buf) + chain_encode(&self.chain, buf, "a window UDF", |codec, buf| { + codec.try_encode_udwf(node, buf) + }) } fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { @@ -570,7 +722,9 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { } else { refuse_if_inline(buf, PY_WINDOW_UDF_FAMILY, "window UDF", name)?; } - self.inner.try_decode_udwf(name, buf) + chain_try(&self.chain, "a window UDF", |codec| { + codec.try_decode_udwf(name, buf) + }) } } @@ -587,7 +741,7 @@ impl PhysicalExtensionCodec for PythonPhysicalCodec { /// `Ok(true)` when the payload (`DFPYUDF` family prefix, version byte, /// cloudpickled tuple) was written and the caller should skip its /// inner codec. Returns `Ok(false)` for any non-Python UDF, signalling -/// the caller to delegate to its `inner`. +/// the caller to delegate to its codec chain. pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec) -> Result { let Some(py_udf) = node.inner().downcast_ref::() else { return Ok(false); @@ -602,7 +756,7 @@ pub(crate) fn try_encode_python_scalar_udf(node: &ScalarUDF, buf: &mut Vec) /// Decode an inline Python scalar UDF payload. Returns `Ok(None)` /// when `buf` does not carry the `DFPYUDF` family prefix, signalling -/// the caller to delegate to its `inner` codec (and eventually the +/// the caller to delegate to its codec chain (and eventually the /// `FunctionRegistry`). pub(crate) fn try_decode_python_scalar_udf(buf: &[u8]) -> Result>> { if !buf.starts_with(PY_SCALAR_UDF_FAMILY) { @@ -1162,3 +1316,181 @@ mod wire_header_tests { )); } } + +#[cfg(test)] +mod codec_chain_tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use datafusion::catalog::MemTable; + use datafusion::common::exec_err; + + use super::*; + + /// Codec that owns a single byte token for table providers and + /// errors on everything else, mirroring the family-prefix + /// discipline expected of downstream FFI codecs. + #[derive(Debug)] + struct TokenCodec { + token: &'static [u8], + /// Return `Ok` from `try_encode_table_provider` without + /// writing bytes, imitating a "no opinion" codec. + encode_by_name: bool, + decode_hits: AtomicUsize, + encode_hits: AtomicUsize, + } + + impl TokenCodec { + fn new(token: &'static [u8]) -> Arc { + Arc::new(Self { + token, + encode_by_name: false, + decode_hits: AtomicUsize::new(0), + encode_hits: AtomicUsize::new(0), + }) + } + + fn new_by_name(token: &'static [u8]) -> Arc { + Arc::new(Self { + token, + encode_by_name: true, + decode_hits: AtomicUsize::new(0), + encode_hits: AtomicUsize::new(0), + }) + } + } + + impl LogicalExtensionCodec for TokenCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[LogicalPlan], + _ctx: &TaskContext, + ) -> Result { + exec_err!("TokenCodec does not decode extension nodes") + } + + fn try_encode(&self, _node: &Extension, _buf: &mut Vec) -> Result<()> { + exec_err!("TokenCodec does not encode extension nodes") + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + _table_ref: &TableReference, + schema: SchemaRef, + _ctx: &TaskContext, + ) -> Result> { + if buf != self.token { + return exec_err!("Unknown table provider token for TokenCodec"); + } + self.decode_hits.fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(MemTable::try_new(schema, vec![vec![]])?)) + } + + fn try_encode_table_provider( + &self, + _table_ref: &TableReference, + _node: Arc, + buf: &mut Vec, + ) -> Result<()> { + self.encode_hits.fetch_add(1, Ordering::SeqCst); + if !self.encode_by_name { + buf.extend_from_slice(self.token); + } + Ok(()) + } + } + + fn mem_table() -> Arc { + Arc::new(MemTable::try_new(Arc::new(Schema::empty()), vec![vec![]]).unwrap()) + } + + fn table_ref() -> TableReference { + TableReference::bare("t") + } + + #[test] + fn decode_falls_through_to_earlier_installed_codec() { + let first = TokenCodec::new(b"AAAA"); + let second = TokenCodec::new(b"BBBB"); + let codec = PythonLogicalCodec::default() + .with_additional_codec(first.clone()) + .with_additional_codec(second.clone()); + + let ctx = TaskContext::default(); + codec + .try_decode_table_provider(b"AAAA", &table_ref(), Arc::new(Schema::empty()), &ctx) + .unwrap(); + + assert_eq!(first.decode_hits.load(Ordering::SeqCst), 1); + assert_eq!(second.decode_hits.load(Ordering::SeqCst), 0); + } + + #[test] + fn most_recently_installed_codec_encodes_first() { + let first = TokenCodec::new(b"AAAA"); + let second = TokenCodec::new(b"BBBB"); + let codec = PythonLogicalCodec::default() + .with_additional_codec(first.clone()) + .with_additional_codec(second.clone()); + + let mut buf = Vec::new(); + codec + .try_encode_table_provider(&table_ref(), mem_table(), &mut buf) + .unwrap(); + + assert_eq!(buf, b"BBBB"); + assert_eq!(first.encode_hits.load(Ordering::SeqCst), 0); + } + + #[test] + fn empty_ok_encode_lets_later_codec_write_payload() { + let writer = TokenCodec::new(b"AAAA"); + let by_name = TokenCodec::new_by_name(b"BBBB"); + let codec = PythonLogicalCodec::default() + .with_additional_codec(writer.clone()) + .with_additional_codec(by_name.clone()); + + let mut buf = Vec::new(); + codec + .try_encode_table_provider(&table_ref(), mem_table(), &mut buf) + .unwrap(); + + assert_eq!(buf, b"AAAA"); + assert_eq!(by_name.encode_hits.load(Ordering::SeqCst), 1); + assert_eq!(writer.encode_hits.load(Ordering::SeqCst), 1); + } + + #[test] + fn decode_failure_aggregates_every_codec_error() { + let codec = PythonLogicalCodec::default() + .with_additional_codec(TokenCodec::new(b"AAAA")) + .with_additional_codec(TokenCodec::new(b"BBBB")); + + let ctx = TaskContext::default(); + let err = codec + .try_decode_table_provider(b"????", &table_ref(), Arc::new(Schema::empty()), &ctx) + .unwrap_err(); + + let msg = err.to_string(); + assert!(msg.contains("None of the 3 composed extension codecs")); + assert!(msg.contains("Unknown table provider token")); + } + + #[test] + fn single_codec_chain_error_is_returned_verbatim() { + let codec = PythonLogicalCodec::default(); + let ctx = TaskContext::default(); + let err = codec + .try_decode_table_provider(b"????", &table_ref(), Arc::new(Schema::empty()), &ctx) + .unwrap_err(); + assert!(!err.to_string().contains("composed extension codecs")); + } + + #[test] + fn with_additional_codec_preserves_udf_inlining_setting() { + let strict = PythonLogicalCodec::default().with_python_udf_inlining(false); + let extended = strict.with_additional_codec(TokenCodec::new(b"AAAA")); + assert!(!extended.python_udf_inlining()); + } +} diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index a1dc0169f..ab64f0010 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1472,7 +1472,9 @@ impl PySessionContext { ) -> PyDataFusionResult { let inner_ffi = ffi_logical_codec_from_pycapsule(codec)?; let inner: Arc = (&inner_ffi).into(); - let logical_codec = Arc::new(PythonLogicalCodec::new(inner)); + // Prepend rather than replace: previously installed codecs stay + // active, with the most recently installed one consulted first. + let logical_codec = Arc::new(self.logical_codec.with_additional_codec(inner)); let physical_codec = Arc::clone(&self.physical_codec); let ctx = self @@ -1497,7 +1499,9 @@ impl PySessionContext { codec: Bound<'py, PyAny>, ) -> PyDataFusionResult { let inner = physical_codec_from_pycapsule(&codec)?; - let physical_codec = Arc::new(PythonPhysicalCodec::new(inner)); + // Prepend rather than replace: previously installed codecs stay + // active, with the most recently installed one consulted first. + let physical_codec = Arc::new(self.physical_codec.with_additional_codec(inner)); let logical_codec = Arc::clone(&self.logical_codec); let ctx = self @@ -1511,11 +1515,15 @@ impl PySessionContext { pub fn with_python_udf_inlining(&self, enabled: bool) -> Self { let logical_codec = Arc::new( - PythonLogicalCodec::new(Arc::clone(self.logical_codec.inner())) + self.logical_codec + .as_ref() + .clone() .with_python_udf_inlining(enabled), ); let physical_codec = Arc::new( - PythonPhysicalCodec::new(Arc::clone(self.physical_codec.inner())) + self.physical_codec + .as_ref() + .clone() .with_python_udf_inlining(enabled), ); let ctx = self diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index d13eda4d8..3b5b8b91b 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -248,15 +248,65 @@ foreign planner. This lets the planner decode provider-owned objects and lets 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. +### Composable codecs + +Extension codecs compose. Each call to `with_logical_extension_codec` or +`with_physical_extension_codec` adds the codec to the front of the session's codec +chain rather than replacing prior codecs. During encoding and decoding, the most +recently installed codec is consulted first, falling through codec by codec to +DataFusion's default codec. A codec signals "not mine" by returning an error, which +sends the chain on to the next codec. Two conventions keep this dispatch sound: + +- Frame your payloads with a distinct byte prefix (pick a `DF` namespace plus a + crate-specific suffix) and only decode payloads carrying your prefix. +- Return an error for objects and payloads you do not own. A codec that answers + success for objects outside its family shadows every codec installed before it. + +Because dispatch keys off payload prefixes rather than install position, codec +registration order between independent libraries does not matter. 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. +### One planner per session, with explicit fallback + +Unlike codecs, a `SessionState` holds exactly one query planner — installing another +replaces it. Planner layering is therefore explicit: a planner that wants to handle +only some queries should accept a fallback planner and delegate the rest to it. The +current planner can be exported for that purpose with +`ctx.__datafusion_query_planner__()`. + +One ordering rule applies: a planner capsule captures the session's codecs at export +time and cannot be rebound afterward. Codec changes made after installing a single +planner are rebound automatically, but a planner wrapped inside another planner as a +fallback is opaque and keeps the codecs it was exported with. **Install all extension +codecs before exporting or chaining planners.** + +Putting it together for a session using two extension libraries that each provide +tables, functions, and a query planner: + +```python +ctx = SessionContext(config) + +# 1. Codecs from both libraries. Order between libraries does not matter. +ctx = ctx.with_logical_extension_codec(lib_a.codec()) +ctx = ctx.with_logical_extension_codec(lib_b.codec()) +ctx = ctx.with_physical_extension_codec(lib_a.physical_codec()) +ctx = ctx.with_physical_extension_codec(lib_b.physical_codec()) + +# 2. Planners, innermost fallback first. Library A's planner falls back to +# DataFusion's default planner; library B's planner falls back to A's. +ctx = ctx.with_query_planner(lib_a.Planner()) +ctx = ctx.with_query_planner( + lib_b.Planner(fallback=ctx.__datafusion_query_planner__()) +) + +# 3. Tables and functions — any time before the first query. +ctx.register_table("t", lib_a.TableProvider()) +ctx.register_udf(udf(lib_b.SomeUDF())) +``` + ## 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 4a35e75da..cbd599698 100644 --- a/examples/datafusion-ffi-example/README.md +++ b/examples/datafusion-ffi-example/README.md @@ -16,7 +16,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. +Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call prepends the codec to the session's codec chain, with the most recently installed codec consulted first and DataFusion's default codec as the terminal fallback. A codec signals "not mine" by returning an error, so several independent plugin libraries can install codecs on the same session as long as each only answers for payloads it owns (frame them with a distinct byte prefix). In this example the provider library is the only 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: diff --git a/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py b/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py index cd0c5a61a..0fd1d7431 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py +++ b/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py @@ -18,7 +18,7 @@ from __future__ import annotations from datafusion import LogicalPlan, SessionContext -from datafusion_ffi_example import MyLogicalExtensionCodec +from datafusion_ffi_example import MyLogicalExtensionCodec, MyTableProvider def _setup_session_with_codec() -> tuple[SessionContext, MyLogicalExtensionCodec]: @@ -80,3 +80,28 @@ def test_ffi_logical_codec_roundtrip(): restored = LogicalPlan.from_bytes(ctx, blob) df_round_trip = ctx.create_dataframe_from_logical_plan(restored) assert df.collect() == df_round_trip.collect() + + +def test_ffi_logical_codec_composes_with_later_install(): + """Codecs compose: installing a second codec prepends it to the + session's codec chain instead of replacing the first. The second + codec here (a default-backed codec exported from a fresh session) + cannot encode this library's table provider, so encoding falls + through to the user codec installed first. Under replace semantics + this test fails with `LogicalExtensionCodec is not provided`.""" + ctx, codec = _setup_session_with_codec() + ctx = ctx.with_logical_extension_codec( + SessionContext().__datafusion_logical_extension_codec__() + ) + + ctx.register_table("numbers", MyTableProvider(1, 4, 1)) + df = ctx.sql('SELECT "A" FROM numbers') + plan = df.logical_plan() + + before = codec.table_provider_encode_calls() + blob = plan.to_bytes(ctx) + assert codec.table_provider_encode_calls() > before + + restored = LogicalPlan.from_bytes(ctx, blob) + df_round_trip = ctx.create_dataframe_from_logical_plan(restored) + assert df.collect() == df_round_trip.collect() diff --git a/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py b/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py index 28eaaf2d9..82116bef7 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py +++ b/examples/datafusion-ffi-example/python/tests/_test_physical_extension_codec.py @@ -76,3 +76,26 @@ def test_ffi_physical_codec_roundtrip(): restored = ExecutionPlan.from_bytes(ctx, blob) assert str(original) == str(restored) + + +def test_ffi_physical_codec_composes_with_later_install(): + """Codecs compose: a second install prepends to the chain instead + of replacing the first codec. The second codec here (default-backed + export from a fresh session) encodes UDFs by name without writing + bytes, which the chain treats as "no opinion" — so the user codec + installed first is still consulted. Under replace semantics its + counter stays at zero.""" + ctx, codec = _setup_session_with_codec() + ctx = ctx.with_physical_extension_codec( + SessionContext().__datafusion_physical_extension_codec__() + ) + + df = ctx.sql("SELECT abs(a) AS x FROM t") + original = df.execution_plan() + + before = codec.encode_udf_calls() + blob = original.to_bytes(ctx) + assert codec.encode_udf_calls() > before + + restored = ExecutionPlan.from_bytes(ctx, blob) + assert str(original) == str(restored) diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 02bfa75fe..2f0e04f9d 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 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`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call prepends to the session's codec chain, so several libraries can install codecs on the same session. This planner owns no serializable types of its own and deliberately uses only built-in physical nodes. Install codecs before the planner; derived contexts rebind codecs after a planner is installed directly, but a planner exported as a fallback for another planner keeps the codecs captured at export time. 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/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..0991c5c9e 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 @@ -118,3 +118,26 @@ def test_query_planner_rejects_invalid_config(max_rows: str): with pytest.raises(Exception, match=r"max_rows|Invalid value"): ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect() + + +def test_composed_codecs_with_query_planner(): + """A second pair of codecs installed on top of the provider codecs + composes with them instead of replacing them. The extra codecs + (default-backed exports from a fresh session) decline everything, + so planner-driven encode/decode falls through to the provider + codecs and the query still succeeds end to end.""" + ctx, logical_codec, physical_codec = configured_context(max_rows=2) + other = SessionContext() + ctx = ctx.with_logical_extension_codec( + other.__datafusion_logical_extension_codec__() + ) + ctx = ctx.with_physical_extension_codec( + other.__datafusion_physical_extension_codec__() + ) + 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] + assert logical_codec.table_provider_encode_calls() > 0 + assert logical_codec.table_provider_decode_calls() > 0 + assert physical_codec.execution_plan_decode_calls() > 0 diff --git a/python/datafusion/context.py b/python/datafusion/context.py index c7b73c5da..b4214fdd5 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1779,6 +1779,15 @@ def with_query_planner( its 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 query planner; installing another replaces + it. To layer planners, construct the new planner with the current + planner as its fallback (export it via + :py:meth:`__datafusion_query_planner__`) before installing. A planner + exported this way captures the codecs installed at export time and + cannot be rebound afterward, so install all extension codecs before + chaining planners. See the FFI extensions guide for the full + multi-library registration recipe. + Args: planner: Object exposing ``__datafusion_query_planner__`` or a raw ``datafusion_query_planner`` PyCapsule. @@ -2229,11 +2238,19 @@ def __datafusion_query_planner__(self) -> Any: def with_logical_extension_codec( self, codec: LogicalExtensionCodecExportable | _PyCapsule ) -> SessionContext: - """Create a new session context with specified codec. + """Create a new session context with an additional logical codec. Only FFI codecs are supported. Pass any object implementing ``__datafusion_logical_extension_codec__`` (see :py:class:`~datafusion.user_defined.LogicalExtensionCodecExportable`). + + Codecs compose: each call adds the codec to the front of the + session's codec chain rather than replacing prior codecs. During + encoding and decoding, the most recently installed codec is + consulted first, falling through codec by codec to DataFusion's + default codec. Codecs signal "not mine" by returning an error, so + extension codecs should only answer for payloads they own — + typically identified by a distinct byte prefix. """ new_internal = self.ctx.with_logical_extension_codec(codec) new = SessionContext.__new__(SessionContext) @@ -2247,11 +2264,16 @@ def __datafusion_physical_extension_codec__(self) -> Any: def with_physical_extension_codec( self, codec: PhysicalExtensionCodecExportable | _PyCapsule ) -> SessionContext: - """Create a new session context with the specified physical codec. + """Create a new session context with an additional physical codec. Only FFI codecs are supported. Pass any object implementing ``__datafusion_physical_extension_codec__`` (see :py:class:`~datafusion.user_defined.PhysicalExtensionCodecExportable`). + + Codecs compose the same way as in + :py:meth:`with_logical_extension_codec`: each call prepends to the + session's codec chain, and the most recently installed codec is + consulted first. """ new_internal = self.ctx.with_physical_extension_codec(codec) new = SessionContext.__new__(SessionContext) From 55f43b5e231da5591e7bbd632dac85f304a22cca Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 11:05:35 -0400 Subject: [PATCH 05/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 cbd599698..1f0775516 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 2f0e04f9d..7e1ec899a 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 c1bb4bf1d2f226e9a75af87eacd3a2be25895425 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 14:09:16 -0400 Subject: [PATCH 06/10] Add atomic SessionContext.with_extensions API Installing FFI extension codecs and query planners by chaining the existing with_* methods can bind task-context providers to intermediate contexts that are later collected, breaking the weak provider reference over the FFI boundary. with_extensions creates one destination context, passes it to each extension factory so components bind to that exact context, and installs everything in a single state write. Co-Authored-By: Claude Fable 5 --- crates/core/src/context.rs | 84 ++++++++++++++++++++ python/datafusion/__init__.py | 4 + python/datafusion/context.py | 141 ++++++++++++++++++++++++++++++++++ python/tests/test_context.py | 115 +++++++++++++++++++++++++++ 4 files changed, 344 insertions(+) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index ab64f0010..cdbc929b2 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1534,6 +1534,90 @@ impl PySessionContext { physical_codec, } } + + /// Create the destination context for a `with_extensions` transaction. + /// + /// Private support method for `SessionContext.with_extensions`. The + /// returned context is the single `Arc` that every FFI + /// task-context provider created during the transaction must target; + /// `_install_extensions` later mutates its state in place rather than + /// deriving a new context. + pub fn _derive_for_extensions(&self) -> Self { + Self { + ctx: Arc::new(SessionContext::new_with_state(self.ctx.state())), + logical_codec: Arc::clone(&self.logical_codec), + physical_codec: Arc::clone(&self.physical_codec), + } + } + + /// Commit a `with_extensions` transaction onto this context. + /// + /// Private support method for `SessionContext.with_extensions`; `self` + /// must be a context produced by `_derive_for_extensions`. Codec capsules + /// are imported and validated before any state change, so a failure + /// leaves the context untouched. The final state is written through this + /// context's own `state_ref()`, never a derived context, so FFI + /// task-context providers bound to it stay valid. + #[pyo3(signature = (logical_codecs, physical_codecs, planner=None))] + pub fn _install_extensions<'py>( + &self, + logical_codecs: Vec>, + physical_codecs: Vec>, + planner: Option>, + ) -> PyDataFusionResult { + let mut logical_codec = self.logical_codec.as_ref().clone(); + for codec in logical_codecs { + let inner_ffi = ffi_logical_codec_from_pycapsule(codec)?; + let inner: Arc = (&inner_ffi).into(); + logical_codec = logical_codec.with_additional_codec(inner); + } + let logical_codec = Arc::new(logical_codec); + + let mut physical_codec = self.physical_codec.as_ref().clone(); + for codec in physical_codecs { + let inner = physical_codec_from_pycapsule(&codec)?; + physical_codec = physical_codec.with_additional_codec(inner); + } + let physical_codec = Arc::new(physical_codec); + + // Bind the planner only after the codec chains are final. Both FFI + // codec wrappers target this exact context so their weak task-context + // providers stay valid for as long as the returned context lives. + let ffi_logical = Self::ffi_logical_codec_for(&self.ctx, &logical_codec); + let ffi_physical = Self::ffi_physical_codec_for(&self.ctx, &physical_codec); + let query_planner: Option> = match planner { + Some(planner) => { + let planner = ffi_query_planner_from_pycapsule(&planner)?; + let planner: Arc = (&planner).into(); + let planner = + FFI_QueryPlanner::new_with_ffi_codecs(planner, ffi_logical, ffi_physical); + Some(Arc::new(RuntimeAwareQueryPlanner { planner })) + } + None => { + let state = self.ctx.state(); + let planner_any: &dyn std::any::Any = state.query_planner().as_ref(); + planner_any + .downcast_ref::() + .map(|p| { + Arc::new(p.with_ffi_codecs(ffi_logical, ffi_physical)) + as Arc + }) + } + }; + + if let Some(query_planner) = query_planner { + let state = SessionStateBuilder::new_from_existing(self.ctx.state()) + .with_query_planner(query_planner) + .build(); + *self.ctx.state_ref().write() = state; + } + + Ok(Self { + ctx: Arc::clone(&self.ctx), + logical_codec, + physical_codec, + }) + } } impl PySessionContext { diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 9c55f446c..86d0054b3 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -80,6 +80,8 @@ RuntimeEnvBuilder, SessionConfig, SessionContext, + SessionExtensionComponents, + SessionExtensionExportable, SQLOptions, ) from .dataframe import ( @@ -134,6 +136,8 @@ "ScalarUDF", "SessionConfig", "SessionContext", + "SessionExtensionComponents", + "SessionExtensionExportable", "Table", "TableFunction", "TableProviderFactory", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index b4214fdd5..ce6bbd53e 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -46,6 +46,7 @@ import uuid import warnings +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Protocol try: @@ -155,6 +156,49 @@ class QueryPlannerExportable(Protocol): def __datafusion_query_planner__(self) -> object: ... # noqa: D105 +@dataclass(frozen=True) +class SessionExtensionComponents: + """Components an extension contributes to a session context. + + Returned by :py:meth:`SessionExtensionExportable.__datafusion_session_extension__` + and consumed by :py:meth:`SessionContext.with_extensions`. Every component + must be created against the context passed to that method; components bound + to any other context hold a task-context provider for the wrong session and + cannot be rebound. + + Attributes: + logical_extension_codecs: Logical codecs to add to the session's codec + chain, in declaration order. + physical_extension_codecs: Physical codecs to add to the session's + codec chain, in declaration order. + query_planner: Optional query planner. At most one extension per + :py:meth:`SessionContext.with_extensions` call may supply one. + """ + + logical_extension_codecs: tuple[ + LogicalExtensionCodecExportable | _PyCapsule, ... + ] = () + physical_extension_codecs: tuple[ + PhysicalExtensionCodecExportable | _PyCapsule, ... + ] = () + query_planner: QueryPlannerExportable | _PyCapsule | None = None + + +class SessionExtensionExportable(Protocol): + """Type hint for extension bundles installable via ``with_extensions``. + + Implementations are reusable configuration objects: they must not retain a + :py:class:`SessionContext` and must create fresh components on every call + using the context supplied by :py:meth:`SessionContext.with_extensions`. + They should also avoid mutating global state during binding, since a + failed installation discards the destination context. + """ + + def __datafusion_session_extension__( # noqa: D105 + self, ctx: SessionContext + ) -> SessionExtensionComponents: ... + + class SessionConfig: """Session configuration options.""" @@ -1808,6 +1852,103 @@ def with_query_planner( new.ctx = new_internal return new + def with_extensions( + self, *extensions: SessionExtensionExportable + ) -> SessionContext: + """Create a new session context with the given extension bundles. + + This is the preferred way to install FFI extensions that need a + task-context provider (extension codecs and query planners). Each + extension's ``__datafusion_session_extension__`` method is called with + the destination context so it can bind its components to that exact + context, then all components are installed in one step. This avoids + the pitfalls of chaining :py:meth:`with_logical_extension_codec`, + :py:meth:`with_physical_extension_codec`, and + :py:meth:`with_query_planner` by hand, where components can end up + bound to an intermediate context that is later garbage collected. + + Codecs compose with the existing chain and with each other: extensions + are processed left to right and prepend to the codec chain, so codecs + from later extensions are consulted first. At most one extension may + supply a query planner. If none does, an existing FFI planner on the + source context is rebound to the final codec chains. + + If any extension raises or returns invalid components, the source + context's state is left unchanged and the partially built destination + is discarded. Extension factories must treat the context they receive + as configuration-only: catalogs are shared with the source context, so + registering tables or otherwise mutating the context during binding is + not rolled back on failure. + + Args: + extensions: One or more objects implementing + ``__datafusion_session_extension__`` (see + :py:class:`SessionExtensionExportable`). + + Returns: + A new context with all extension components installed. + + Raises: + TypeError: If an argument does not implement the protocol or + returns something other than a + :py:class:`SessionExtensionComponents`. + ValueError: If no extensions are given or more than one extension + supplies a query planner. + + Examples: + >>> from my_extension import DistributedEngineExtension # doctest: +SKIP + >>> ctx = SessionContext().with_extensions( + ... DistributedEngineExtension("scheduler:50050") + ... ) # doctest: +SKIP + >>> ctx.sql("SELECT 1").collect() # doctest: +SKIP + """ + if not extensions: + msg = "with_extensions requires at least one extension" + raise ValueError(msg) + for extension in extensions: + if not hasattr(extension, "__datafusion_session_extension__"): + msg = ( + "Extension does not implement __datafusion_session_extension__: " + f"{extension!r}" + ) + raise TypeError(msg) + + # Single destination context. Every component the extensions create + # must bind to this context; _install_extensions later mutates its + # state in place so those bindings stay valid. + destination = SessionContext.__new__(SessionContext) + destination.ctx = self.ctx._derive_for_extensions() + + logical_codecs: list[LogicalExtensionCodecExportable | _PyCapsule] = [] + physical_codecs: list[PhysicalExtensionCodecExportable | _PyCapsule] = [] + planner: QueryPlannerExportable | _PyCapsule | None = None + for extension in extensions: + components = extension.__datafusion_session_extension__(destination) + if not isinstance(components, SessionExtensionComponents): + msg = ( + "__datafusion_session_extension__ must return " + "SessionExtensionComponents, got " + f"{type(components).__name__} from {extension!r}" + ) + raise TypeError(msg) + logical_codecs.extend(components.logical_extension_codecs) + physical_codecs.extend(components.physical_extension_codecs) + if components.query_planner is not None: + if planner is not None: + msg = ( + "Multiple extensions supplied a query planner; a " + "session context has exactly one. Layer planners " + "explicitly instead." + ) + raise ValueError(msg) + planner = components.query_planner + + new = SessionContext.__new__(SessionContext) + new.ctx = destination.ctx._install_extensions( + logical_codecs, physical_codecs, planner + ) + return new + def table_provider(self, name: str) -> Table: """Return the :py:class:`~datafusion.catalog.Table` for the given table name. diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 6e6eaadbe..207234851 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -16,6 +16,7 @@ # under the License. import ctypes import datetime as dt +import gc import gzip import pathlib import shutil @@ -29,6 +30,7 @@ RuntimeEnvBuilder, SessionConfig, SessionContext, + SessionExtensionComponents, SQLOptions, Table, column, @@ -754,6 +756,119 @@ def test_with_query_planner_capsule(ctx): assert batches[0].column(0) == pa.array([1]) +class _CodecOnlyExtension: + """Contributes decline-all codecs exported from an unrelated session.""" + + def __init__(self): + self.exporter = SessionContext() + self.bound_ctx = None + + def __datafusion_session_extension__(self, ctx): + self.bound_ctx = ctx + return SessionExtensionComponents( + logical_extension_codecs=( + self.exporter.__datafusion_logical_extension_codec__(), + ), + physical_extension_codecs=( + self.exporter.__datafusion_physical_extension_codec__(), + ), + ) + + +class _PlannerExtension: + """Contributes the destination context's own exported planner.""" + + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + query_planner=ctx.__datafusion_query_planner__() + ) + + +def test_with_extensions_requires_an_extension(ctx): + with pytest.raises(ValueError, match="at least one extension"): + ctx.with_extensions() + + +def test_with_extensions_rejects_non_extension(ctx): + with pytest.raises(TypeError, match="__datafusion_session_extension__"): + ctx.with_extensions(object()) + + +def test_with_extensions_rejects_bad_components(ctx): + class BadExtension: + def __datafusion_session_extension__(self, ctx): + return 42 + + with pytest.raises(TypeError, match="SessionExtensionComponents"): + ctx.with_extensions(BadExtension()) + + +def test_with_extensions_rejects_multiple_planners(ctx): + with pytest.raises(ValueError, match="query planner"): + ctx.with_extensions(_PlannerExtension(), _PlannerExtension()) + + +def test_with_extensions_rejects_bad_codec_capsule(ctx): + class BadCodecExtension: + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + logical_extension_codecs=(ctx.__datafusion_task_context_provider__(),), + ) + + with pytest.raises(ValueError, match="incorrect name"): + ctx.with_extensions(BadCodecExtension()) + + +def test_with_extensions_installs_codecs_and_planner(ctx): + ctx.register_record_batches( + "extensions_test", + [[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]], + ) + extension = _CodecOnlyExtension() + result = ctx.with_extensions(extension, _PlannerExtension()) + + assert result.table_exist("extensions_test") + # In-memory tables need a real extension codec to round-trip through the + # FFI planner, so query plans that don't serialize a table provider. + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_with_extensions_binds_to_returned_context(ctx): + extension = _CodecOnlyExtension() + result = ctx.with_extensions(extension) + + # The context passed to the factory shares the same underlying session + # as the returned context: registrations made through it are visible. + extension.bound_ctx.register_record_batches( + "bound_test", + [[pa.RecordBatch.from_pydict({"value": [1]})]], + ) + assert result.table_exist("bound_test") + + +def test_with_extensions_survives_source_collection(): + extension = _CodecOnlyExtension() + result = SessionContext().with_extensions(extension, _PlannerExtension()) + gc.collect() + + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_with_extensions_failure_leaves_source_usable(ctx): + class BoomExtension: + def __datafusion_session_extension__(self, ctx): + msg = "boom" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="boom"): + ctx.with_extensions(_CodecOnlyExtension(), BoomExtension()) + + batches = ctx.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 c868447c7a2b919d07388e67aa91488b269c2cca Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 14:17:45 -0400 Subject: [PATCH 07/10] Add extension-bundle example and with_extensions FFI tests MyPlannerExtension in the query-planner example crate implements the __datafusion_session_extension__ protocol from Rust: it extracts the destination context's task-context provider, binds fresh observing codecs and a planner to it, and returns SessionExtensionComponents. Its codecs record the max_rows config value resolved through the weak provider, letting tests prove the provider targets the returned context rather than the source. Documents with_extensions as the preferred API in the FFI guide. Co-Authored-By: Claude Fable 5 --- docs/source/contributor-guide/ffi.md | 62 ++++- .../README.md | 17 +- .../_test_three_library_query_planner.py | 170 ++++++++++- .../src/extension.rs | 263 ++++++++++++++++++ .../src/lib.rs | 3 + .../src/planner.rs | 22 +- 6 files changed, 523 insertions(+), 14 deletions(-) create mode 100644 examples/datafusion-ffi-query-planner-example/src/extension.rs diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 3b5b8b91b..17c58826f 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -283,8 +283,60 @@ planner are rebound automatically, but a planner wrapped inside another planner fallback is opaque and keeps the codecs it was exported with. **Install all extension codecs before exporting or chaining planners.** -Putting it together for a session using two extension libraries that each provide -tables, functions, and a query planner: +### Extension bundles: `with_extensions` + +Codec and planner capsules carry an `FFI_TaskContextProvider` holding a *weak* +reference to the `SessionContext` they were created against. A capsule does not keep +that context alive, and a component bound to one context cannot be rebound to another. +Chaining the low-level `with_*` methods by hand therefore risks binding components to +an intermediate context that is later garbage collected, which fails at query time +with `TaskContextProvider went out of scope over FFI boundary` — or worse, silently +reads stale session state. + +`SessionContext.with_extensions` avoids this by construction. An extension library +exposes a bundle object implementing the `__datafusion_session_extension__` protocol: + +```python +class MyEngineExtension: + def __datafusion_session_extension__(self, ctx: SessionContext) -> SessionExtensionComponents: + # Create fresh components bound to `ctx` on every call. `ctx` is the + # exact context the host will return from with_extensions. + return SessionExtensionComponents( + logical_extension_codecs=(self._make_logical_codec(ctx),), + physical_extension_codecs=(self._make_physical_codec(ctx),), + query_planner=self._make_planner(ctx), + ) +``` + +The host creates one destination context, passes it to every factory, installs all +codecs, binds the planner against the final codec chains, and returns the context in +a single step: + +```python +ctx = SessionContext(config).with_extensions(lib_a.Extension(), lib_b.Extension()) +ctx.register_table("t", lib_a.TableProvider()) +ctx.register_udf(udf(lib_b.SomeUDF())) +``` + +Extensions are processed left to right and prepend to the codec chain, so codecs from +later extensions are consulted first. At most one extension may supply a query +planner. If any factory fails, the source context's state is unchanged. + +Bundle objects must be configuration-only: create fresh components on each call, never +cache bound components, and do not retain the context passed in. Catalogs are shared +with the source context, so registrations made during binding are not rolled back on +failure. + +`MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete Rust +implementation of this protocol, including extracting the task-context provider from +the supplied context and constructing a Python `SessionExtensionComponents`. + +### Advanced: chaining the low-level methods + +The `with_logical_extension_codec`, `with_physical_extension_codec`, and +`with_query_planner` methods remain available for advanced use. Putting them together +for a session using two extension libraries that each provide tables, functions, and +a query planner: ```python ctx = SessionContext(config) @@ -307,6 +359,12 @@ ctx.register_table("t", lib_a.TableProvider()) ctx.register_udf(udf(lib_b.SomeUDF())) ``` +When chaining by hand, keep the final context assigned to `ctx` as the single owner: +components created against earlier intermediate contexts (for example a codec +constructed with a context that is later discarded) hold weak references that break +once that intermediate context is collected. Prefer `with_extensions` whenever the +extension library provides a bundle. + ## Alternative Approach Suppose you needed to expose some other features of DataFusion and you could not wait diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 7e1ec899a..55e898693 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -41,7 +41,22 @@ uv run pytest \ examples/datafusion-ffi-query-planner-example/python/tests/_test*.py ``` -The integration test follows this setup: +The preferred setup uses `SessionContext.with_extensions` with extension bundles: + +```python +config = SessionConfig().with_extension(PlannerConfig(max_rows=3)) +ctx = SessionContext(config).with_extensions(provider_bundle, MyPlannerExtension()) +ctx.register_table("numbers", provider) +ctx.register_udf(provider_udf) +``` + +`MyPlannerExtension` implements the `__datafusion_session_extension__` protocol: it +receives the destination context, binds fresh codec and planner components to that +context's task-context provider, and returns them as `SessionExtensionComponents`. +The host installs everything in one step, so no component can end up bound to an +intermediate context that is later collected. + +The integration tests also cover the low-level chaining setup: ```python config = SessionConfig().with_extension(PlannerConfig(max_rows=3)) 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 0991c5c9e..5846661f5 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 @@ -20,14 +20,23 @@ import gc import pytest -from datafusion import SessionConfig, SessionContext, udf +from datafusion import ( + SessionConfig, + SessionContext, + SessionExtensionComponents, + udf, +) from datafusion_ffi_example import ( IsNullUDF, MyLogicalExtensionCodec, MyPhysicalExtensionCodec, MyTableProvider, ) -from datafusion_ffi_query_planner_example import MyQueryPlanner, PlannerConfig +from datafusion_ffi_query_planner_example import ( + MyPlannerExtension, + MyQueryPlanner, + PlannerConfig, +) def configured_context(max_rows: int): @@ -120,6 +129,163 @@ def test_query_planner_rejects_invalid_config(max_rows: str): ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect() +class ProviderCodecsExtension: + """Bundles the provider library's codecs for ``with_extensions``. + + These codecs keep their own private task-context provider, so they only + need to be created once; the bundle can hand out the same exporters on + every call. + """ + + def __init__(self) -> None: + self.logical_codec = MyLogicalExtensionCodec() + self.physical_codec = MyPhysicalExtensionCodec() + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return SessionExtensionComponents( + logical_extension_codecs=(self.logical_codec,), + physical_extension_codecs=(self.physical_codec,), + ) + + +def test_with_extensions_three_library_query(): + """One with_extensions call installs provider codecs and a planner bundle, + and a real non-empty plan flows across the three libraries.""" + config = SessionConfig().with_extension(PlannerConfig(max_rows=3)) + provider_ext = ProviderCodecsExtension() + planner_ext = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(provider_ext, planner_ext) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + + 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_ext.plan_calls() >= 1 + assert planner_ext.last_max_rows() == 3 + assert planner_ext.foreign_session_observed() + assert planner_ext.foreign_provider_observed() + assert planner_ext.foreign_plan_observed() + assert provider_ext.logical_codec.table_provider_encode_calls() > 0 + assert provider_ext.logical_codec.table_provider_decode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_encode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_decode_calls() > 0 + + +def test_with_extensions_provider_targets_returned_context(): + """The bundle's task-context provider reads current state from the + returned context, not the source it was derived from.""" + config = SessionConfig().with_extension(PlannerConfig(max_rows=3)) + source = SessionContext(config) + source.register_table("numbers", MyTableProvider(1, 6, 1)) + planner_ext = MyPlannerExtension() + result = source.with_extensions(ProviderCodecsExtension(), planner_ext) + + # Diverge the two live contexts. Config state is copied at derivation, + # so after these statements source and result disagree. + source.sql("SET ffi_query_planner.max_rows = 5").collect() + result.sql("SET ffi_query_planner.max_rows = 2").collect() + + batches = result.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner_ext.last_max_rows() == 2 + + # Codec decode calls resolve the task context through the weak provider + # bound during with_extensions. Seeing 2 (never 5) proves the provider + # targets the returned context rather than the source. + seen = planner_ext.decode_max_rows_seen() + assert seen, "expected the bundle codecs to observe at least one decode" + assert set(seen) == {2} + + +def test_with_extensions_survives_dropping_source_and_bundles(): + """Neither the source context nor the bundle objects are needed to keep + the installed components' task-context provider alive.""" + config = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + ctx = SessionContext(config).with_extensions( + ProviderCodecsExtension(), MyPlannerExtension() + ) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + gc.collect() + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + + +def test_with_extensions_sees_state_changes_after_install(): + """Tables, UDFs, and config changes made after installation are visible + to the planner and to provider callbacks.""" + config = SessionConfig().with_extension(PlannerConfig(max_rows=4)) + planner_ext = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(ProviderCodecsExtension(), planner_ext) + + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + ctx.sql("SET ffi_query_planner.max_rows = 2").collect() + + 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] + assert planner_ext.last_max_rows() == 2 + seen = planner_ext.decode_max_rows_seen() + assert seen + assert set(seen) == {2} + + +def test_with_extensions_bundle_is_reusable(): + """Installing the same bundle into two contexts binds fresh components to + each destination.""" + planner_ext = MyPlannerExtension() + + config_a = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + ctx_a = SessionContext(config_a).with_extensions( + ProviderCodecsExtension(), planner_ext + ) + ctx_a.register_table("numbers", MyTableProvider(1, 6, 1)) + + config_b = SessionConfig().with_extension(PlannerConfig(max_rows=3)) + ctx_b = SessionContext(config_b).with_extensions( + ProviderCodecsExtension(), planner_ext + ) + ctx_b.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx_a.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner_ext.last_max_rows() == 2 + + batches = ctx_b.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert planner_ext.last_max_rows() == 3 + + +def test_with_extensions_failure_leaves_source_usable(): + """A failing factory after a successful one leaves the source context + fully functional.""" + + class BoomExtension: + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + msg = "boom" + raise RuntimeError(msg) + + config = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + source = SessionContext(config) + source.register_table("numbers", MyTableProvider(1, 6, 1)) + + with pytest.raises(RuntimeError, match="boom"): + source.with_extensions(MyPlannerExtension(), BoomExtension()) + + # No planner was installed, so the default planner runs unrestricted. + batches = source.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2, 3, 4, 5] + + def test_composed_codecs_with_query_planner(): """A second pair of codecs installed on top of the provider codecs composes with them instead of replacing them. The extra codecs diff --git a/examples/datafusion-ffi-query-planner-example/src/extension.rs b/examples/datafusion-ffi-query-planner-example/src/extension.rs new file mode 100644 index 000000000..714d3b324 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -0,0 +1,263 @@ +// 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::fmt; +use std::ptr::NonNull; +use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::{Result, TableReference}; +use datafusion::datasource::TableProvider; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::{Extension, LogicalPlan}; +use datafusion::physical_plan::ExecutionPlan; +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_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, +}; +use datafusion_python_util::get_tokio_runtime; +use datafusion_session::QueryPlanner; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyDict}; + +use crate::planner::{DistributedQueryPlanner, PlannerObservations, planner_config_from_options}; + +/// Values of `ffi_query_planner.max_rows` observed through the task-context +/// provider bound at installation time. Decoding reads the provider's current +/// session state, so these prove which context the provider targets. +type ObservedMaxRows = Arc>>; + +fn record_task_ctx(observed: &ObservedMaxRows, ctx: &TaskContext) { + if let Ok(config) = planner_config_from_options(ctx.session_config().options()) + && let Ok(mut observed) = observed.lock() + { + observed.push(config.max_rows); + } +} + +/// Records the task context resolved by the FFI wrapper, then declines by +/// delegating to the default codec so the host's codec chain falls through to +/// the codec that owns the payload. +struct ObservingLogicalExtensionCodec { + inner: DefaultLogicalExtensionCodec, + observed: ObservedMaxRows, +} + +impl fmt::Debug for ObservingLogicalExtensionCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ObservingLogicalExtensionCodec") + .finish_non_exhaustive() + } +} + +impl LogicalExtensionCodec for ObservingLogicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[LogicalPlan], + ctx: &TaskContext, + ) -> Result { + record_task_ctx(&self.observed, ctx); + self.inner.try_decode(buf, inputs, ctx) + } + + fn try_encode(&self, node: &Extension, buf: &mut Vec) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + table_ref: &TableReference, + schema: SchemaRef, + ctx: &TaskContext, + ) -> Result> { + record_task_ctx(&self.observed, ctx); + self.inner + .try_decode_table_provider(buf, table_ref, schema, ctx) + } + + fn try_encode_table_provider( + &self, + table_ref: &TableReference, + node: Arc, + buf: &mut Vec, + ) -> Result<()> { + self.inner.try_encode_table_provider(table_ref, node, buf) + } +} + +/// Physical companion to [`ObservingLogicalExtensionCodec`]. +struct ObservingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec, + observed: ObservedMaxRows, +} + +impl fmt::Debug for ObservingPhysicalExtensionCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ObservingPhysicalExtensionCodec") + .finish_non_exhaustive() + } +} + +impl PhysicalExtensionCodec for ObservingPhysicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + record_task_ctx(&self.observed, ctx); + self.inner.try_decode(buf, inputs, ctx, proto_converter) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + self.inner.try_encode(node, buf, proto_converter) + } +} + +fn task_ctx_provider_from_session(ctx: &Bound<'_, PyAny>) -> PyResult { + let capsule = ctx.call_method0("__datafusion_task_context_provider__")?; + let capsule = capsule.cast::()?; + let provider: NonNull = capsule + .pointer_checked(Some(c"datafusion_task_context_provider"))? + .cast(); + // The FFI provider holds a weak reference; cloning it does not keep the + // session context alive. The host's returned context is the strong owner. + Ok(unsafe { provider.as_ref() }.clone()) +} + +/// Extension bundle for `SessionContext.with_extensions`. +/// +/// Mirrors how a distributed engine such as Ballista packages its session +/// extensions: the object itself is reusable configuration, and every +/// `__datafusion_session_extension__` call creates fresh codec and planner +/// components bound to the task-context provider of the context it receives. +#[pyclass( + from_py_object, + name = "MyPlannerExtension", + module = "datafusion_ffi_query_planner_example", + subclass +)] +#[derive(Debug, Default, Clone)] +pub(crate) struct MyPlannerExtension { + observations: Arc, + observed_max_rows: ObservedMaxRows, +} + +#[pymethods] +impl MyPlannerExtension { + #[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) + } + + /// `ffi_query_planner.max_rows` values seen through the bound + /// task-context provider during codec decode calls. + fn decode_max_rows_seen(&self) -> Vec { + self.observed_max_rows + .lock() + .map(|observed| observed.clone()) + .unwrap_or_default() + } + + fn __datafusion_session_extension__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + ) -> PyResult> { + // Bind every component to the destination context supplied by the + // host. Components must not be cached across calls: each installation + // targets a different context. + let provider = task_ctx_provider_from_session(&ctx)?; + let runtime = get_tokio_runtime().handle().clone(); + + let logical: Arc = Arc::new(ObservingLogicalExtensionCodec { + inner: DefaultLogicalExtensionCodec {}, + observed: Arc::clone(&self.observed_max_rows), + }); + let ffi_logical = + FFI_LogicalExtensionCodec::new(logical, Some(runtime.clone()), provider.clone()); + let logical_capsule = + PyCapsule::new_with_value(py, ffi_logical, cr"datafusion_logical_extension_codec")?; + + let physical: Arc = + Arc::new(ObservingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec {}, + observed: Arc::clone(&self.observed_max_rows), + }); + let ffi_physical = + FFI_PhysicalExtensionCodec::new(physical, Some(runtime.clone()), provider.clone()); + let physical_capsule = + PyCapsule::new_with_value(py, ffi_physical, cr"datafusion_physical_extension_codec")?; + + let planner: Arc = Arc::new(DistributedQueryPlanner { + observations: Arc::clone(&self.observations), + }); + let ffi_planner = FFI_QueryPlanner::new( + planner, + Some(runtime), + provider, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::new(DefaultPhysicalExtensionCodec {}), + ); + let planner_capsule = + PyCapsule::new_with_value(py, ffi_planner, cr"datafusion_query_planner")?; + + let components = py + .import("datafusion")? + .getattr("SessionExtensionComponents")?; + let kwargs = PyDict::new(py); + kwargs.set_item("logical_extension_codecs", (logical_capsule,))?; + kwargs.set_item("physical_extension_codecs", (physical_capsule,))?; + kwargs.set_item("query_planner", planner_capsule)?; + components.call((), Some(&kwargs)) + } +} diff --git a/examples/datafusion-ffi-query-planner-example/src/lib.rs b/examples/datafusion-ffi-query-planner-example/src/lib.rs index 7635c2992..ff8ce7848 100644 --- a/examples/datafusion-ffi-query-planner-example/src/lib.rs +++ b/examples/datafusion-ffi-query-planner-example/src/lib.rs @@ -18,15 +18,18 @@ use pyo3::prelude::*; use crate::config::PlannerConfig; +use crate::extension::MyPlannerExtension; use crate::planner::MyQueryPlanner; mod config; +mod extension; mod planner; #[pymodule] fn datafusion_ffi_query_planner_example(m: &Bound<'_, PyModule>) -> PyResult<()> { pyo3_log::init(); m.add_class::()?; + m.add_class::()?; m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index cb767ffa5..3477f8c0f 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -41,12 +41,12 @@ 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, +pub(crate) struct PlannerObservations { + pub(crate) plan_calls: AtomicUsize, + pub(crate) last_max_rows: AtomicUsize, + pub(crate) foreign_session: AtomicBool, + pub(crate) foreign_provider: AtomicBool, + pub(crate) foreign_plan: AtomicBool, } fn logical_plan_has_foreign_provider(plan: &LogicalPlan) -> bool { @@ -70,8 +70,12 @@ fn physical_plan_has_foreign_plan(plan: &Arc) -> bool { } fn planner_config(session: &dyn Session) -> datafusion::common::Result { - let options = session.config_options(); + planner_config_from_options(session.config_options()) +} +pub(crate) fn planner_config_from_options( + options: &datafusion::common::config::ConfigOptions, +) -> datafusion::common::Result { // 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 @@ -107,8 +111,8 @@ fn planner_config(session: &dyn Session) -> datafusion::common::Result, +pub(crate) struct DistributedQueryPlanner { + pub(crate) observations: Arc, } #[async_trait] From 825b748295725f414bde49bbd5ec7037c639f5b1 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 14:19:06 -0400 Subject: [PATCH 08/10] Document and test the context-outlives-DataFrame contract A DataFrame does not keep its SessionContext alive. FFI components hold a weak task-context provider, so operations that reach an FFI codec after the context is collected fail with a clean out-of-scope error rather than crashing. Lock that behavior in with a test and document the ownership contract in the FFI guide and with_extensions docstring. Co-Authored-By: Claude Fable 5 --- docs/source/contributor-guide/ffi.md | 7 +++++++ .../_test_three_library_query_planner.py | 19 +++++++++++++++++++ python/datafusion/context.py | 5 +++++ 3 files changed, 31 insertions(+) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 17c58826f..a51e979e6 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -327,6 +327,13 @@ cache bound components, and do not retain the context passed in. Catalogs are sh with the source context, so registrations made during binding are not rolled back on failure. +The returned context is the strong owner of every installed component's task-context +provider, and dependent objects do not extend its lifetime. A `DataFrame`, logical +plan, or capsule can outlive the context, but any operation that reaches an FFI codec +after the context is collected fails with `TaskContextProvider went out of scope over +FFI boundary`. Keep the context alive for as long as objects derived from it are in +use. + `MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete Rust implementation of this protocol, including extracting the task-context provider from the supplied context and constructing a Python `SessionExtensionComponents`. 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 5846661f5..2fde3e2d1 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 @@ -286,6 +286,25 @@ def __datafusion_session_extension__( assert batches[0].column(0).to_pylist() == [0, 1, 2, 3, 4, 5] +def test_dataframe_outliving_context_fails_cleanly(): + """A DataFrame does not keep its SessionContext alive. FFI components + resolve the task context through a weak reference, so using the + DataFrame after dropping the context raises a clean error instead of + crashing. This locks in the documented ownership contract: the context + must outlive DataFrames that depend on FFI codecs.""" + config = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + ctx = SessionContext(config).with_extensions( + ProviderCodecsExtension(), MyPlannerExtension() + ) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + df = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"') + del ctx + gc.collect() + + with pytest.raises(Exception, match="went out of scope"): + df.collect() + + def test_composed_codecs_with_query_planner(): """A second pair of codecs installed on top of the provider codecs composes with them instead of replacing them. The extra codecs diff --git a/python/datafusion/context.py b/python/datafusion/context.py index ce6bbd53e..19b3191c0 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -1880,6 +1880,11 @@ def with_extensions( registering tables or otherwise mutating the context during binding is not rolled back on failure. + The returned context is the strong owner of the installed components' + task-context providers. Keep it alive for as long as DataFrames or + plans derived from it are in use; FFI operations after the context is + collected raise an error. + Args: extensions: One or more objects implementing ``__datafusion_session_extension__`` (see From 8a75acac9a1f17e2f44672990cf76edc50792676 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 14:20:13 -0400 Subject: [PATCH 09/10] Skip private internal methods in wrapper coverage test Single-underscore methods on internal pyo3 classes (such as SessionContext._install_extensions) are private support methods for the Python wrappers and do not require a public wrapper. Co-Authored-By: Claude Fable 5 --- python/tests/test_wrapper_coverage.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index cf6719ecf..b1afd6832 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -67,6 +67,14 @@ def missing_exports(internal_obj, wrapped_obj) -> None: pytest.fail(f"Missing __repr__: {internal_obj.__name__}") for internal_attr_name in dir(internal_obj): + # Single-underscore names are private support methods for the + # wrappers (e.g. SessionContext._install_extensions) and are not + # part of the public surface that requires a wrapper. + if internal_attr_name.startswith("_") and not internal_attr_name.startswith( + "__" + ): + continue + wrapped_attr_name = internal_attr_name.removeprefix("Raw") assert wrapped_attr_name in dir(wrapped_obj) From 590a4a29a2048d7c3887ad09957c0a09dc5ddd39 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 7 Aug 2026 14:23:14 -0400 Subject: [PATCH 10/10] Test planner rebinding and codec precedence in with_extensions Co-Authored-By: Claude Fable 5 --- .../_test_three_library_query_planner.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) 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 2fde3e2d1..3dd737301 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 @@ -286,6 +286,45 @@ def __datafusion_session_extension__( assert batches[0].column(0).to_pylist() == [0, 1, 2, 3, 4, 5] +def test_with_extensions_rebinds_existing_planner(): + """Codec-only bundles installed on a context that already has an FFI + planner rebind that planner to the new codec chains.""" + config = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + planner = MyQueryPlanner() + ctx = SessionContext(config).with_query_planner(planner) + provider_ext = ProviderCodecsExtension() + ctx = ctx.with_extensions(provider_ext) + ctx.register_table("numbers", MyTableProvider(1, 6, 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 + # The planner only sees these codecs if it was rebound to the chains + # built during with_extensions. + assert provider_ext.logical_codec.table_provider_decode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_decode_calls() > 0 + + +def test_with_extensions_codec_precedence(): + """Extensions are processed left to right and prepend to the codec + chain, so codecs from later extensions are consulted first. Both + bundles' codecs can handle the payload (they share the provider + library's token registry); only the one consulted first is used.""" + config = SessionConfig().with_extension(PlannerConfig(max_rows=2)) + ext_a = ProviderCodecsExtension() + ext_b = ProviderCodecsExtension() + ctx = SessionContext(config).with_extensions(ext_a, ext_b, MyPlannerExtension()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + + assert ext_b.logical_codec.table_provider_encode_calls() > 0 + assert ext_b.logical_codec.table_provider_decode_calls() > 0 + assert ext_a.logical_codec.table_provider_encode_calls() == 0 + assert ext_a.logical_codec.table_provider_decode_calls() == 0 + + def test_dataframe_outliving_context_fails_cleanly(): """A DataFrame does not keep its SessionContext alive. FFI components resolve the task context through a weak reference, so using the