fix: Hybrid planning when the root plan has local preference

Fixes: #2104

Signed-off-by: Vaibhav <vrongmeal@gmail.com>

crates/datafusion_ext/src/transform.rs

@@ -15,6 +15,16 @@ pub trait TreeNodeExt: TreeNode {
         let new_node = op(after_op_children)?.into();
         Ok(new_node)
     }
+
+    /// Like `transform_down`, but allows providing a closure with mutable access
+    /// to the environment.
+    fn transform_down_mut<F>(self, op: &mut F) -> Result<Self>
+    where
+        F: FnMut(Self) -> Result<Transformed<Self>>,
+    {
+        let after_op = op(self)?.into();
+        after_op.map_children(|node| node.transform_down_mut(op))
+    }
 }
 
 impl<T: DynTreeNode + ?Sized> TreeNodeExt for Arc<T> {}

crates/sqlexec/src/planner/physical_plan/remote_scan.rs

@@ -14,7 +14,9 @@ use datafusion::physical_plan::{
 };
 use datafusion::prelude::Expr;
 use datafusion_ext::metrics::AggregateMetricsStreamAdapter;
+use datafusion_ext::runtime::runtime_group::RuntimeGroupExec;
 use futures::{stream, TryStreamExt};
+use protogen::metastore::types::catalog::RuntimePreference;
 use std::any::Any;
 use std::fmt;
 use std::hash::Hash;
@@ -78,21 +80,27 @@ pub struct RemoteScanExec {
 }
 
 impl RemoteScanExec {
+    /// Returns a remote scan exec wrapped inside a [`RuntimeGroupExec`]
+    /// asserting that this can only be run on the remote side.
+    #[allow(clippy::new_ret_no_self)]
     pub fn new(
         provider: ProviderReference,
         projected_schema: Arc<Schema>,
         projection: Option<Vec<usize>>,
         filters: Vec<Expr>,
         limit: Option<usize>,
-    ) -> Self {
-        Self {
-            provider,
-            projected_schema,
-            projection,
-            filters,
-            limit,
-            metrics: ExecutionPlanMetricsSet::new(),
-        }
+    ) -> RuntimeGroupExec {
+        RuntimeGroupExec::new(
+            RuntimePreference::Remote,
+            Arc::new(Self {
+                provider,
+                projected_schema,
+                projection,
+                filters,
+                limit,
+                metrics: ExecutionPlanMetricsSet::new(),
+            }),
+        )
     }
 }