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..047b35039 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -93,11 +93,15 @@ jobs: uv venv --python "${{ steps.setup-python.outputs.python-path }}" VENV_PY="$PWD/.venv/bin/python" uv sync --python "$VENV_PY" --dev --no-install-package datafusion + # Search recursively: the FFI artifact bundles more than one + # project, so upload-artifact keeps a `/dist/` prefix + # and the wheels are not all at the top of wheels/. WHEELS=$(find wheels/ -name "*.whl") if [ -n "$WHEELS" ]; then echo "Installing wheels:" echo "$WHEELS" - uv pip install --python "$VENV_PY" wheels/*.whl + # shellcheck disable=SC2086 # intentional split on newlines + uv pip install --python "$VENV_PY" $WHEELS else echo "ERROR: No wheels found!" exit 1 @@ -121,6 +125,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 d34862ac7..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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" [[package]] name = "datafusion-execution" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +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/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-schema", @@ -1193,10 +1197,27 @@ dependencies = [ "pyo3-log", ] +[[package]] +name = "datafusion-ffi-query-planner-example" +version = "54.0.0" +dependencies = [ + "async-trait", + "datafusion", + "datafusion-catalog", + "datafusion-common", + "datafusion-ffi", + "datafusion-proto", + "datafusion-python-util", + "datafusion-session", + "pyo3", + "pyo3-build-config", + "pyo3-log", +] + [[package]] name = "datafusion-functions" version = "54.1.0" -source = "git+https://github.com/apache/datafusion?rev=dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48#dbcb5c0f729e9ef6b0ab4c79253fe3b657929f48" +source = "git+https://github.com/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-buffer", @@ -1227,7 +1248,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1247,7 +1268,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1258,7 +1279,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-ord", @@ -1282,7 +1303,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "async-trait", @@ -1297,7 +1318,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1313,7 +1334,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1322,7 +1343,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "datafusion-doc", "quote", @@ -1332,7 +1353,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "chrono", @@ -1351,7 +1372,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1373,7 +1394,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1387,7 +1408,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "chrono", @@ -1404,7 +1425,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1423,7 +1444,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "arrow-data", @@ -1459,7 +1480,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "chrono", @@ -1486,7 +1507,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1496,7 +1517,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "datafusion-proto-common", "prost", @@ -1505,7 +1526,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "datafusion-common", @@ -1530,6 +1551,7 @@ dependencies = [ "datafusion-ffi", "datafusion-proto", "datafusion-python-util", + "datafusion-session", "datafusion-spark", "datafusion-substrait", "futures", @@ -1565,7 +1587,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow-schema", "async-trait", @@ -1579,7 +1601,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "bigdecimal", @@ -1608,7 +1630,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "arrow", "bigdecimal", @@ -1627,7 +1649,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/apache/datafusion?rev=e08aed1e5de41dcf81d529140dae07723b942a5e#e08aed1e5de41dcf81d529140dae07723b942a5e" dependencies = [ "async-recursion", "async-trait", @@ -2446,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", ] @@ -2525,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" @@ -2639,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", @@ -2650,7 +2682,7 @@ dependencies = [ "arrow-ipc", "arrow-schema", "arrow-select", - "base64 0.22.1", + "base64 0.23.0", "brotli", "bytes", "chrono", @@ -2659,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", @@ -2672,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 362159913..809ef7a98 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] @@ -50,6 +55,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 +78,14 @@ codegen-units = 2 # We cannot publish to crates.io with any patches in the below section. Developers # must remove any entries in this section before creating a release candidate. [patch.crates-io] -datafusion = { git = "https://github.com/apache/datafusion", rev = "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/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/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..a1dc0169f 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,42 @@ impl PySessionConfig { } } +/// Adapts an FFI planner to the Tokio runtime owned by datafusion-python. +/// +/// Upstream's `ForeignQueryPlanner` cannot recover the runtime handle from the +/// `QueryPlanner` trait, so embedders that own the runtime must call +/// `create_physical_plan_with_session_runtime` directly. +#[derive(Debug, Clone)] +struct RuntimeAwareQueryPlanner { + planner: FFI_QueryPlanner, +} + +impl RuntimeAwareQueryPlanner { + fn with_ffi_codecs( + &self, + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, + ) -> Self { + let planner: Arc = (&self.planner).into(); + let planner = FFI_QueryPlanner::new_with_ffi_codecs(planner, logical_codec, physical_codec); + Self { planner } + } +} + +#[async_trait] +impl QueryPlanner for RuntimeAwareQueryPlanner { + async fn create_physical_plan( + &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 +1253,32 @@ impl PySessionContext { Ok(()) } + pub fn with_query_planner(&self, planner: Bound<'_, PyAny>) -> PyDataFusionResult { + let planner = ffi_query_planner_from_pycapsule(&planner)?; + + // Build the codecs against the derived context, then update that same + // context in place. FFI codecs keep a weak task-context provider, so + // rebuilding the context after creating them would leave a stale link. + let ctx = Arc::new(SessionContext::new_with_state(self.ctx.state())); + let planner: Arc = (&planner).into(); + let planner = FFI_QueryPlanner::new_with_ffi_codecs( + planner, + Self::ffi_logical_codec_for(&ctx, &self.logical_codec), + Self::ffi_physical_codec_for(&ctx, &self.physical_codec), + ); + let planner = Arc::new(RuntimeAwareQueryPlanner { planner }); + let state = SessionStateBuilder::new_from_existing(ctx.state()) + .with_query_planner(planner) + .build(); + *ctx.state_ref().write() = state; + + 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 +1453,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>, @@ -1393,10 +1474,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, }) } @@ -1415,9 +1499,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, }) } @@ -1431,8 +1518,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, } @@ -1440,6 +1529,34 @@ 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::() + .cloned() + else { + return Arc::clone(&self.ctx); + }; + + // Preserve the context identity captured by the replacement codecs. + let ctx = Arc::new(SessionContext::new_with_state(state)); + let planner = Arc::new(planner.with_ffi_codecs( + Self::ffi_logical_codec_for(&ctx, &logical_codec), + Self::ffi_physical_codec_for(&ctx, &physical_codec), + )); + let state = SessionStateBuilder::new_from_existing(ctx.state()) + .with_query_planner(planner) + .build(); + *ctx.state_ref().write() = state; + ctx + } + async fn _table(&self, name: &str) -> datafusion::common::Result { self.ctx.table(name).await } @@ -1501,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/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/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/README.md b/examples/datafusion-ffi-example/README.md new file mode 100644 index 000000000..c897c067a --- /dev/null +++ b/examples/datafusion-ffi-example/README.md @@ -0,0 +1,50 @@ + + +# 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/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-query-planner-example/Cargo.toml b/examples/datafusion-ffi-query-planner-example/Cargo.toml new file mode 100644 index 000000000..263f034b8 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/Cargo.toml @@ -0,0 +1,50 @@ +# 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-proto = { 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..53e2bddc5 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -0,0 +1,60 @@ + + +# 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 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/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-query-planner-example/pyproject.toml b/examples/datafusion-ffi-query-planner-example/pyproject.toml new file mode 100644 index 000000000..9e34b4cd4 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/pyproject.toml @@ -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. + +[build-system] +requires = ["maturin>=1.6,<2.0"] +build-backend = "maturin" + +[project] +name = "datafusion_ffi_query_planner_example" +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", +] +dynamic = ["version"] + +[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..cb767ffa5 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -0,0 +1,206 @@ +// 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_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::*; +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, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::new(DefaultPhysicalExtensionCodec {}), + ); + PyCapsule::new_with_value(py, ffi, cr"datafusion_query_planner") + } +} diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 94b2bb1c6..c7b73c5da 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,36 @@ 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. Codec changes made on + a derived context are rebound to the planner before planning. + + 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 +2222,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]])