perf: Implement physical execution of uncorrelated scalar subqueries (#21240)

## Which issue does this PR close?

- Closes #3781.
- Closes #18181.

## Rationale for this change

Previously, DataFusion evaluated uncorrelated scalar subqueries by
transforming them into joins. This has three shortcomings:

1. Scalar subqueries that return > 1 row were allowed, producing
incorrect query results. Such queries should instead result in a runtime
error.
2. Performance. Evaluating scalar subqueries as a join requires going
through the join machinery. More importantly, it means that UDFs that
have specialized handling of scalar inputs cannot use those code paths
for scalar subqueries, which often results in significantly slower query
execution (e.g., #18181). It also makes filter pushdown for scalar
subquery filters more difficult (#21324)
3. Uncorrelated scalar subqueries previously did not work in `ORDER BY`
or `JOIN ON`, or as arguments to an aggregate function. Those cases are
now supported.

This PR introduces physical execution of uncorrelated scalar subqueries:

* Uncorrelated subqueries are left in the plan by the optimizer, not
rewritten into joins
* The physical planner collects uncorrelated scalar subqueries and plans
them recursively (supporting nested subqueries). We add a
`ScalarSubqueryExec` plan node to the top of any physical plan with
uncorrelated subqueries: it has N+1 children, N subqueries and its
"main" input, which is the rest of the query plan. The subquery
expression in the parent plan is replaced with a `ScalarSubqueryExpr`.
* `ScalarSubqueryExec` manages the execution of the subqueries. Subquery
evaluation is done in parallel (for a given query level), but at present
it happens strictly before evaluation of the parent query. This might be
improved in the future (#21591).
* `ScalarSubqueryExpr` reads its value from a shared slot that
`ScalarSubqueryExec` populates when the subquery finishes; the physical
planner assigns each subquery its slot index via `ExecutionProps`.

This architecture makes it easy to avoid the shortcomings described
above. Performance seems roughly unchanged (benchmarks added in this
PR), but in situations like #18181, we can now leverage scalar
fast-paths; in the case of #18181 specifically, this improves
performance from ~800 ms to ~30 ms.

## What changes are included in this PR?

* Modify subquery rewriter to not transform subqueries -> joins
* Collect and plan uncorrelated scalar subqueries in the physical
planner, and wire up `ScalarSubqueryExpr`
* Support for subqueries in physical plan serialization/deserialization
using `PhysicalProtoConverterExtension` to wire up `ScalarSubqueryExpr`
correctly
* Support for subqueries in logical plan serialization/deserialization
* Add various SLT tests and update expected plan shapes for some tests

## Are these changes tested?

Yes. New SLT coverage for cardinality errors, `ORDER BY` / `JOIN ON` /
aggregate-arg contexts, nested uncorrelated subqueries,
duplicate-subquery deduplication, and partition-pruning filters; new
roundtrip tests for logical and physical plan serialization.

## Are there any user-facing changes?

SQL:
* Uncorrelated scalar subqueries that return more than one row now
result in a runtime error, instead of silently producing incorrect
results.
* Uncorrelated scalar subqueries now work in `ORDER BY`, `JOIN ON`, and
as aggregate function arguments.

Rust APIs:
* In `datafusion-proto`, breaking changes to
`Serializeable::from_bytes_with_registry` (renamed to
`from_bytes_with_ctx`), `parse_expr` / `parse_sorts` / `parse_exprs`,
and the `PhysicalProtoConverterExtension` trait.

Plan shape:
* `LogicalPlan::Subquery` nodes will now be preserved in the logical
plan
* Physical plans can now contain `ScalarSubqueryExec` plan node and
`ScalarSubqueryExpr` expressions

The wire format has also changed to include scalar subqueries.

---------

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: neilconway

.github/workflows/large_files.yml

@@ -34,9 +34,9 @@ jobs:
           fetch-depth: 0
       - name: Check size of new Git objects
         env:
-          # 1 MB ought to be enough for anybody.
+          # 1.5 MB ought to be enough for anybody.
           # TODO in case we may want to consciously commit a bigger file to the repo without using Git LFS we may disable the check e.g. with a label
-          MAX_FILE_SIZE_BYTES: 1048576
+          MAX_FILE_SIZE_BYTES: 1572864
         shell: bash
         run: |
           if [ "${{ github.event_name }}" = "merge_group" ]; then

Cargo.lock

@@ -17,9 +17,9 @@
 
 //! See `main.rs` for how to run it.
 //!
-//! This example demonstrates how to use the `PhysicalExtensionCodec` trait's
-//! interception methods (`serialize_physical_plan` and `deserialize_physical_plan`)
-//! to implement custom serialization logic.
+//! This example demonstrates how to use the `PhysicalProtoConverterExtension`
+//! trait's interception methods (`execution_plan_to_proto` and
+//! `proto_to_execution_plan`) to implement custom serialization logic.
 //!
 //! The key insight is that `FileScanConfig::expr_adapter_factory` is NOT serialized by
 //! default. This example shows how to:
@@ -28,9 +28,10 @@
 //! 3. Store the inner DataSourceExec (without adapter) as a child in the extension's inputs field
 //! 4. Unwrap and restore the adapter during deserialization
 //!
-//! This demonstrates nested serialization (protobuf outer, JSON inner) and the power
-//! of the `PhysicalExtensionCodec` interception pattern. Both plan and expression
-//! serialization route through the codec, enabling interception at every node in the tree.
+//! This demonstrates nested serialization (protobuf outer, JSON inner) and the
+//! power of `PhysicalProtoConverterExtension`. Both plan and expression
+//! serialization route through converter hooks, enabling interception at every
+//! node in the tree.
 
 use std::fmt::Debug;
 use std::sync::Arc;
@@ -61,7 +62,7 @@ use datafusion_proto::bytes::{
 use datafusion_proto::physical_plan::from_proto::parse_physical_expr_with_converter;
 use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter;
 use datafusion_proto::physical_plan::{
-    PhysicalExtensionCodec, PhysicalProtoConverterExtension,
+    PhysicalExtensionCodec, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension,
 };
 use datafusion_proto::protobuf::physical_plan_node::PhysicalPlanType;
 use datafusion_proto::protobuf::{
@@ -177,7 +178,7 @@ pub async fn adapter_serialization() -> Result<()> {
     println!("\n=== Example Complete! ===");
     println!("Key takeaways:");
     println!(
-        "  1. PhysicalExtensionCodec provides serialize_physical_plan/deserialize_physical_plan hooks"
+        "  1. PhysicalProtoConverterExtension provides execution_plan_to_proto/proto_to_execution_plan hooks"
     );
     println!("  2. Custom metadata can be wrapped as PhysicalExtensionNode");
     println!("  3. Nested serialization (protobuf + JSON) works seamlessly");
@@ -303,9 +304,10 @@ impl PhysicalExtensionCodec for AdapterPreservingCodec {
         _node: Arc<dyn ExecutionPlan>,
         _buf: &mut Vec<u8>,
     ) -> Result<()> {
-        // We don't need this for the example - we use serialize_physical_plan instead
+        // We don't need this for the example - adapter wrapping happens in
+        // `execution_plan_to_proto` instead.
         not_impl_err!(
-            "try_encode not used - adapter wrapping happens in serialize_physical_plan"
+            "try_encode not used - adapter wrapping happens in execution_plan_to_proto"
         )
     }
 }
@@ -371,9 +373,8 @@ impl PhysicalProtoConverterExtension for AdapterPreservingCodec {
     // Interception point: override deserialization to unwrap adapters
     fn proto_to_execution_plan(
         &self,
-        ctx: &TaskContext,
-        extension_codec: &dyn PhysicalExtensionCodec,
         proto: &PhysicalPlanNode,
+        ctx: &PhysicalPlanDecodeContext<'_>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         // Check if this is our custom extension wrapper
         if let Some(PhysicalPlanType::Extension(extension)) = &proto.physical_plan_type
@@ -395,11 +396,7 @@ impl PhysicalProtoConverterExtension for AdapterPreservingCodec {
             let inner_proto = &extension.inputs[0];
 
             // Deserialize the inner plan
-            let inner_plan = inner_proto.try_into_physical_plan_with_converter(
-                ctx,
-                extension_codec,
-                self,
-            )?;
+            let inner_plan = self.default_proto_to_execution_plan(inner_proto, ctx)?;
 
             // Recreate the adapter factory
             let adapter_factory = create_adapter_factory(&payload.adapter_metadata.tag);
@@ -409,17 +406,16 @@ impl PhysicalProtoConverterExtension for AdapterPreservingCodec {
         }
 
         // Not our extension - use default deserialization
-        proto.try_into_physical_plan_with_converter(ctx, extension_codec, self)
+        self.default_proto_to_execution_plan(proto, ctx)
     }
 
     fn proto_to_physical_expr(
         &self,
         proto: &PhysicalExprNode,
-        ctx: &TaskContext,
         input_schema: &Schema,
-        codec: &dyn PhysicalExtensionCodec,
+        ctx: &PhysicalPlanDecodeContext<'_>,
     ) -> Result<Arc<dyn PhysicalExpr>> {
-        parse_physical_expr_with_converter(proto, ctx, input_schema, codec, self)
+        parse_physical_expr_with_converter(proto, input_schema, ctx, self)
     }
 
     fn physical_expr_to_proto(

datafusion-examples/examples/custom_data_source/adapter_serialization.rs

@@ -17,8 +17,9 @@
 
 //! See `main.rs` for how to run it.
 //!
-//! This example demonstrates how to use the `PhysicalExtensionCodec` trait's
-//! interception methods to implement expression deduplication during deserialization.
+//! This example demonstrates how to use the
+//! `PhysicalProtoConverterExtension` trait's interception methods to
+//! implement expression deduplication during deserialization.
 //!
 //! This pattern is inspired by PR #18192, which introduces expression caching
 //! to reduce memory usage when deserializing plans with duplicate expressions.
@@ -29,8 +30,9 @@
 //! 2. Reduce memory allocation during deserialization
 //! 3. Enable downstream optimizations that rely on Arc pointer equality
 //!
-//! This demonstrates the decorator pattern enabled by the `PhysicalExtensionCodec` trait,
-//! where all expression serialization/deserialization routes through the codec methods.
+//! This demonstrates the decorator pattern enabled by
+//! `PhysicalProtoConverterExtension`, where physical-expression
+//! serialization and deserialization route through converter hooks.
 
 use std::collections::HashMap;
 use std::fmt::Debug;
@@ -49,7 +51,7 @@ use datafusion::prelude::SessionContext;
 use datafusion_proto::physical_plan::from_proto::parse_physical_expr_with_converter;
 use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter;
 use datafusion_proto::physical_plan::{
-    DefaultPhysicalExtensionCodec, PhysicalExtensionCodec,
+    DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalPlanDecodeContext,
     PhysicalProtoConverterExtension,
 };
 use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode};
@@ -202,11 +204,10 @@ impl PhysicalExtensionCodec for CachingCodec {
 impl PhysicalProtoConverterExtension for CachingCodec {
     fn proto_to_execution_plan(
         &self,
-        ctx: &TaskContext,
-        extension_codec: &dyn PhysicalExtensionCodec,
         proto: &PhysicalPlanNode,
+        ctx: &PhysicalPlanDecodeContext<'_>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        proto.try_into_physical_plan_with_converter(ctx, extension_codec, self)
+        self.default_proto_to_execution_plan(proto, ctx)
     }
 
     fn execution_plan_to_proto(
@@ -225,9 +226,8 @@ impl PhysicalProtoConverterExtension for CachingCodec {
     fn proto_to_physical_expr(
         &self,
         proto: &PhysicalExprNode,
-        ctx: &TaskContext,
         input_schema: &Schema,
-        codec: &dyn PhysicalExtensionCodec,
+        ctx: &PhysicalPlanDecodeContext<'_>,
     ) -> Result<Arc<dyn PhysicalExpr>> {
         // Create cache key from protobuf bytes
         let mut key = Vec::new();
@@ -249,8 +249,7 @@ impl PhysicalProtoConverterExtension for CachingCodec {
         }
 
         // Cache miss - deserialize and store
-        let expr =
-            parse_physical_expr_with_converter(proto, ctx, input_schema, codec, self)?;
+        let expr = parse_physical_expr_with_converter(proto, input_schema, ctx, self)?;
 
         // Store in cache
         {

datafusion-examples/examples/proto/expression_deduplication.rs

@@ -144,6 +144,7 @@ datafusion-session = { workspace = true }
 datafusion-sql = { workspace = true, optional = true }
 flate2 = { workspace = true, optional = true }
 futures = { workspace = true }
+indexmap = { workspace = true }
 itertools = { workspace = true }
 liblzma = { workspace = true, optional = true }
 log = { workspace = true }