fix: resolve a planner fallback when it is installed, not constructed

MyQueryPlanner::new imported its fallback immediately, with no session
to pass, so the fallback's getter was called with no arguments. That
works for a SessionContext, whose getter takes the session optionally,
and for a raw capsule, which has no getter at all. It fails for another
foreign planner, which implements the same protocol this type does and
requires the argument -- and layering on another planner is the case a
distributed engine actually needs. The docstring claimed fallback "takes
anything exporting __datafusion_query_planner__", which was not true.

Holds the Python object instead and imports it in
__datafusion_query_planner__, where the session is in hand and can be
forwarded. All three fallback kinds now work.

Deferring also removes a footgun rather than adding one. Passing a
SessionContext now delegates to whichever planner it holds at install
time, and since with_query_planner calls the getter before installing,
the context still reports its previous planner, so wrapping a context in
a planner installed on that same context does not recurse.

Arc<Py<PyAny>> rather than Py<PyAny> because pyo3 0.29 gates Py: Clone
behind the py-clone feature, and this type derives Clone. Matches how
PythonTableFunctionCallable holds its callable.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: timsaucer

examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py

@@ -358,6 +358,46 @@ def test_planner_layers_on_the_session_planner():
     assert physical_codec.execution_plan_decode_calls() > 0
 
 
+def test_a_planner_can_fall_back_to_another_planner_library():
+    """A fallback may be another foreign planner, not only a session.
+
+    The fallback is imported when this planner is installed rather than when
+    it is constructed, so its own getter receives the session. Importing it at
+    construction time would mean calling that getter with no session, which
+    only a ``SessionContext`` or a raw capsule tolerates -- and layering on
+    another planner is the case a distributed engine actually needs.
+    """
+    ctx, logical_codec, physical_codec = configured_context(max_rows=3)
+    inner = MyQueryPlanner()
+    outer = MyQueryPlanner(fallback=inner)
+    ctx = ctx.with_query_planner(outer)
+
+    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
+    assert batches[0].column(0).to_pylist() == [0, 1, 2]
+    assert outer.plan_calls() > 0
+    assert outer.used_fallback()
+    # The delegation reached the inner planner rather than stopping at the
+    # default physical planner.
+    assert inner.plan_calls() > 0
+    assert logical_codec.table_provider_decode_calls() > 0
+    assert physical_codec.execution_plan_decode_calls() > 0
+
+
+def test_a_session_fallback_delegates_to_its_installed_planner():
+    """Passing a SessionContext delegates to whatever planner it holds."""
+    ctx, _logical_codec, _physical_codec = configured_context(max_rows=3)
+    first = MyQueryPlanner()
+    ctx = ctx.with_query_planner(first)
+
+    second = MyQueryPlanner(fallback=ctx)
+    ctx = ctx.with_query_planner(second)
+
+    batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect()
+    assert batches[0].column(0).to_pylist() == [0, 1, 2]
+    assert second.used_fallback()
+    assert first.plan_calls() > 0
+
+
 def test_second_planner_replaces_the_first():
     """A session holds exactly one planner, so installing another replaces it."""
     ctx, _logical_codec, _physical_codec = configured_context(max_rows=2)

examples/datafusion-ffi-query-planner-example/src/planner.rs

@@ -203,30 +203,38 @@ impl QueryPlanner for DistributedQueryPlanner {
 #[derive(Debug, Default, Clone)]
 pub(crate) struct MyQueryPlanner {
     observations: Arc<PlannerObservations>,
-    fallback: Option<Arc<dyn QueryPlanner + Send + Sync>>,
+    /// Held as the Python object rather than an imported planner, and resolved
+    /// in `__datafusion_query_planner__` where a session is in hand.
+    ///
+    /// Importing it here would mean calling its getter with no session, which
+    /// only a `SessionContext` or a raw capsule accepts. Another foreign
+    /// planner -- the case that matters, since layering is the whole point of
+    /// a fallback -- implements the same protocol this type does and requires
+    /// the argument.
+    fallback: Option<Arc<Py<PyAny>>>,
 }
 
 #[pymethods]
 impl MyQueryPlanner {
     /// Build a planner, optionally layered on top of an existing one.
     ///
-    /// `fallback` takes anything exporting `__datafusion_query_planner__`,
-    /// including a `SessionContext`. Capture it *before* installing this
-    /// planner on that context, or the capsule will describe this planner and
-    /// planning will recurse.
+    /// `fallback` takes anything exporting `__datafusion_query_planner__`:
+    /// another planner library, a `SessionContext`, or a raw capsule. It is
+    /// imported when this planner is installed, not here, so that the session
+    /// can be handed to its getter.
+    ///
+    /// Passing a `SessionContext` delegates to whichever planner that context
+    /// holds at install time. If you instead capture a capsule with
+    /// `ctx.__datafusion_query_planner__()`, capture it *before* installing
+    /// this planner on that context, or the capsule will describe this planner
+    /// and planning will recurse.
     #[new]
     #[pyo3(signature = (fallback=None))]
-    fn new(fallback: Option<Bound<'_, PyAny>>) -> PyResult<Self> {
-        let fallback = fallback
-            .map(|planner| {
-                ffi_query_planner_from_pycapsule(&planner, None)
-                    .map(|ffi| -> Arc<dyn QueryPlanner + Send + Sync> { (&ffi).into() })
-            })
-            .transpose()?;
-        Ok(Self {
-            fallback,
+    fn new(fallback: Option<Bound<'_, PyAny>>) -> Self {
+        Self {
+            fallback: fallback.map(|obj| Arc::new(obj.unbind())),
             ..Self::default()
-        })
+        }
     }
 
     fn used_fallback(&self) -> bool {
@@ -264,9 +272,21 @@ impl MyQueryPlanner {
         py: Python<'py>,
         session: Bound<'py, PyAny>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
+        // Resolved here rather than in `new` so the fallback's own getter
+        // receives the session, which is what the protocol requires of every
+        // implementation other than a `SessionContext`.
+        let fallback = self
+            .fallback
+            .as_ref()
+            .map(|planner| {
+                ffi_query_planner_from_pycapsule(planner.bind(py), Some(&session))
+                    .map(|ffi| -> Arc<dyn QueryPlanner + Send + Sync> { (&ffi).into() })
+            })
+            .transpose()?;
+
         let planner: Arc<dyn QueryPlanner + Send + Sync> = Arc::new(DistributedQueryPlanner {
             observations: Arc::clone(&self.observations),
-            fallback: self.fallback.clone(),
+            fallback,
         });
         let logical_codec = ffi_logical_codec_from_pycapsule(session.clone(), None)?;
         let physical_codec = ffi_physical_codec_from_pycapsule(session, None)?;