datafusion_physical_expr/
scalar_subquery.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Physical expression for uncorrelated scalar subqueries.
19
20use std::fmt;
21use std::hash::Hash;
22use std::sync::Arc;
23
24use arrow::datatypes::{DataType, Field, FieldRef, Schema};
25use arrow::record_batch::RecordBatch;
26use datafusion_common::{Result, internal_datafusion_err};
27use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex};
28use datafusion_expr_common::columnar_value::ColumnarValue;
29use datafusion_expr_common::sort_properties::{ExprProperties, SortProperties};
30use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
31
32/// A physical expression whose value is provided by a scalar subquery.
33///
34/// Subquery execution is handled by `ScalarSubqueryExec`, which stores the
35/// result in a shared [`ScalarSubqueryResults`] container. This expression
36/// simply reads from that container at the appropriate index.
37#[derive(Debug)]
38pub struct ScalarSubqueryExpr {
39    data_type: DataType,
40    nullable: bool,
41    /// Index of this subquery in the shared results container.
42    index: SubqueryIndex,
43    /// Shared results container populated by `ScalarSubqueryExec`.
44    results: ScalarSubqueryResults,
45}
46
47impl ScalarSubqueryExpr {
48    pub fn new(
49        data_type: DataType,
50        nullable: bool,
51        index: SubqueryIndex,
52        results: ScalarSubqueryResults,
53    ) -> Self {
54        Self {
55            data_type,
56            nullable,
57            index,
58            results,
59        }
60    }
61
62    pub fn results(&self) -> &ScalarSubqueryResults {
63        &self.results
64    }
65
66    #[deprecated(
67        since = "55.0.0",
68        note = "was only used for proto serialization, which no longer needs it; use `return_field` for type/nullability. It will be removed in 61.0.0 or 6 months after 55.0.0 is released, whichever is longer."
69    )]
70    pub fn data_type(&self) -> &DataType {
71        &self.data_type
72    }
73
74    #[deprecated(
75        since = "55.0.0",
76        note = "was only used for proto serialization, which no longer needs it; use `return_field` for type/nullability. It will be removed in 61.0.0 or 6 months after 55.0.0 is released, whichever is longer."
77    )]
78    pub fn nullable(&self) -> bool {
79        self.nullable
80    }
81
82    /// Returns the index of this subquery in the shared results container.
83    #[deprecated(
84        since = "55.0.0",
85        note = "was only used for proto serialization, which no longer needs it. It will be removed in 61.0.0 or 6 months after 55.0.0 is released, whichever is longer."
86    )]
87    pub fn index(&self) -> SubqueryIndex {
88        self.index
89    }
90}
91
92impl fmt::Display for ScalarSubqueryExpr {
93    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
94        match self.results.get(self.index) {
95            Some(v) => write!(f, "scalar_subquery({v})"),
96            None => write!(f, "scalar_subquery(<pending>)"),
97        }
98    }
99}
100
101// Two ScalarSubqueryExprs are considered the "same" if they refer to the
102// same underlying shared results container and the same index within it.
103impl Hash for ScalarSubqueryExpr {
104    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
105        self.results.hash(state);
106        self.index.hash(state);
107    }
108}
109
110impl PartialEq for ScalarSubqueryExpr {
111    fn eq(&self, other: &Self) -> bool {
112        self.results == other.results && self.index == other.index
113    }
114}
115
116impl Eq for ScalarSubqueryExpr {}
117
118impl PhysicalExpr for ScalarSubqueryExpr {
119    fn return_field(&self, _input_schema: &Schema) -> Result<FieldRef> {
120        Ok(Arc::new(Field::new(
121            "scalar_subquery",
122            self.data_type.clone(),
123            self.nullable,
124        )))
125    }
126
127    fn evaluate(&self, _batch: &RecordBatch) -> Result<ColumnarValue> {
128        let value = self.results.get(self.index).ok_or_else(|| {
129            internal_datafusion_err!(
130                "ScalarSubqueryExpr evaluated before the subquery was executed"
131            )
132        })?;
133        Ok(ColumnarValue::Scalar(value))
134    }
135
136    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
137        vec![]
138    }
139
140    fn with_new_children(
141        self: Arc<Self>,
142        _children: Vec<Arc<dyn PhysicalExpr>>,
143    ) -> Result<Arc<dyn PhysicalExpr>> {
144        Ok(self)
145    }
146
147    fn get_properties(&self, _children: &[ExprProperties]) -> Result<ExprProperties> {
148        Ok(ExprProperties::new_unknown().with_order(SortProperties::Singleton))
149    }
150
151    fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
152        write!(f, "(scalar subquery)")
153    }
154
155    #[cfg(feature = "proto")]
156    fn try_to_proto(
157        &self,
158        _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
159    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
160        use datafusion_proto_models::protobuf;
161        Ok(Some(protobuf::PhysicalExprNode {
162            expr_id: None,
163            expr_type: Some(protobuf::physical_expr_node::ExprType::ScalarSubquery(
164                protobuf::PhysicalScalarSubqueryExprNode {
165                    data_type: Some((&self.data_type).try_into()?),
166                    nullable: self.nullable,
167                    index: u32::try_from(self.index.as_usize()).map_err(|_| {
168                        internal_datafusion_err!(
169                            "scalar subquery index {} does not fit in u32",
170                            self.index.as_usize()
171                        )
172                    })?,
173                },
174            )),
175        }))
176    }
177}
178
179#[cfg(feature = "proto")]
180impl ScalarSubqueryExpr {
181    /// Reconstruct a [`ScalarSubqueryExpr`] from its protobuf representation.
182    ///
183    /// Unlike other expressions, this takes a third argument: the shared
184    /// [`ScalarSubqueryResults`] container. That container is a runtime-only
185    /// `Arc` shared with the surrounding `ScalarSubqueryExec` and is not part of
186    /// the wire format, so it cannot be reconstructed here or carried on the
187    /// decode context (which lives in a crate that cannot depend on
188    /// `datafusion-expr`). The match arm in `from_proto.rs` fetches it from the
189    /// plan-level decode context and passes it in.
190    pub fn try_from_proto(
191        node: &datafusion_proto_models::protobuf::PhysicalExprNode,
192        _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
193        results: &ScalarSubqueryResults,
194    ) -> Result<Arc<dyn PhysicalExpr>> {
195        use datafusion_physical_expr_common::expect_expr_variant;
196        use datafusion_physical_expr_common::physical_expr::proto_decode::require_proto_field;
197        use datafusion_proto_models::protobuf;
198
199        let sq = expect_expr_variant!(
200            node,
201            protobuf::physical_expr_node::ExprType::ScalarSubquery,
202            "ScalarSubqueryExpr",
203        );
204        let data_type = require_proto_field(
205            sq.data_type.as_ref(),
206            "ScalarSubqueryExpr",
207            "data_type",
208        )?
209        .try_into()?;
210        Ok(Arc::new(ScalarSubqueryExpr::new(
211            data_type,
212            sq.nullable,
213            SubqueryIndex::new(sq.index as usize),
214            results.clone(),
215        )))
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222
223    use arrow::array::Int32Array;
224    use arrow::datatypes::Field;
225    use datafusion_common::ScalarValue;
226
227    fn make_results(values: Vec<Option<ScalarValue>>) -> ScalarSubqueryResults {
228        let results = ScalarSubqueryResults::new(values.len());
229        for (index, value) in values.into_iter().enumerate() {
230            if let Some(value) = value {
231                results.set(SubqueryIndex::new(index), value).unwrap();
232            }
233        }
234        results
235    }
236
237    #[test]
238    fn test_evaluate_with_value() -> Result<()> {
239        let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
240        let a = Int32Array::from(vec![1, 2, 3]);
241        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)])?;
242
243        let results = make_results(vec![Some(ScalarValue::Int32(Some(42)))]);
244        let expr = ScalarSubqueryExpr::new(
245            DataType::Int32,
246            false,
247            SubqueryIndex::new(0),
248            results,
249        );
250
251        let result = expr.evaluate(&batch)?;
252        match result {
253            ColumnarValue::Scalar(ScalarValue::Int32(Some(42))) => {}
254            other => panic!("Expected Scalar(Int32(42)), got {other:?}"),
255        }
256        Ok(())
257    }
258
259    #[test]
260    fn test_evaluate_before_populated() {
261        let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
262        let a = Int32Array::from(vec![1]);
263        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
264
265        let results = ScalarSubqueryResults::new(1);
266        let expr = ScalarSubqueryExpr::new(
267            DataType::Int32,
268            false,
269            SubqueryIndex::new(0),
270            results,
271        );
272
273        let result = expr.evaluate(&batch);
274        assert!(result.is_err());
275    }
276
277    #[test]
278    fn test_identity_equality() {
279        let results = make_results(vec![None, None]);
280
281        let e1a = ScalarSubqueryExpr::new(
282            DataType::Int32,
283            false,
284            SubqueryIndex::new(0),
285            results.clone(),
286        );
287        let e1b = ScalarSubqueryExpr::new(
288            DataType::Int32,
289            false,
290            SubqueryIndex::new(0),
291            results.clone(),
292        );
293        let e2 = ScalarSubqueryExpr::new(
294            DataType::Int32,
295            false,
296            SubqueryIndex::new(1),
297            results.clone(),
298        );
299
300        // Same container + same index → equal
301        assert_eq!(e1a, e1b);
302        // Same container, different index → not equal
303        assert_ne!(e1a, e2);
304
305        // Different container, same index → not equal
306        let other_results = make_results(vec![None]);
307        let e3 = ScalarSubqueryExpr::new(
308            DataType::Int32,
309            false,
310            SubqueryIndex::new(0),
311            other_results,
312        );
313        assert_ne!(e1a, e3);
314    }
315}
316
317/// Tests for the `try_to_proto` / `try_from_proto` hooks.
318#[cfg(all(test, feature = "proto"))]
319mod proto_tests {
320    use super::*;
321    use crate::proto_test_util::{StubEncoder, UnreachableDecoder, column_node};
322    use datafusion_common::DataFusionError;
323    use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
324    use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
325    use datafusion_proto_models::protobuf::{
326        PhysicalExprNode, PhysicalScalarSubqueryExprNode, physical_expr_node,
327    };
328
329    /// Build a `ScalarSubquery` proto node directly, with control over each
330    /// field, so the decode error paths can be exercised independently.
331    fn proto_scalar_subquery_node(
332        data_type: Option<datafusion_proto_models::datafusion_common::ArrowType>,
333        nullable: bool,
334        index: u32,
335    ) -> PhysicalExprNode {
336        PhysicalExprNode {
337            expr_id: None,
338            expr_type: Some(physical_expr_node::ExprType::ScalarSubquery(
339                PhysicalScalarSubqueryExprNode {
340                    data_type,
341                    nullable,
342                    index,
343                },
344            )),
345        }
346    }
347
348    #[test]
349    fn round_trips_through_proto() {
350        // A three-slot results container so index 2 is meaningful.
351        let results = ScalarSubqueryResults::new(3);
352        let expr = ScalarSubqueryExpr::new(
353            DataType::Int32,
354            true,
355            SubqueryIndex::new(2),
356            results.clone(),
357        );
358
359        // Encode: the expression serializes itself via try_to_proto.
360        let encoder = StubEncoder::ok();
361        let enc_ctx = PhysicalExprEncodeCtx::new(&encoder);
362        let node = expr
363            .try_to_proto(&enc_ctx)
364            .unwrap()
365            .expect("ScalarSubqueryExpr should encode to Some(node)");
366
367        assert!(node.expr_id.is_none());
368        let sq = match &node.expr_type {
369            Some(physical_expr_node::ExprType::ScalarSubquery(sq)) => sq,
370            other => panic!("expected a ScalarSubquery node, got {other:?}"),
371        };
372        assert!(sq.nullable);
373        assert_eq!(sq.index, 2);
374        let encoded_type: DataType = sq
375            .data_type
376            .as_ref()
377            .expect("data_type encoded")
378            .try_into()
379            .unwrap();
380        assert_eq!(encoded_type, DataType::Int32);
381
382        // Decode: reconstruct from the proto node, threading in the shared
383        // results container the surrounding exec would provide.
384        let decoder = UnreachableDecoder;
385        let schema = Schema::empty();
386        let dec_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
387        let decoded =
388            ScalarSubqueryExpr::try_from_proto(&node, &dec_ctx, &results).unwrap();
389        let decoded = decoded
390            .downcast_ref::<ScalarSubqueryExpr>()
391            .expect("decoded expr should be a ScalarSubqueryExpr");
392
393        // data_type + nullable survive the round-trip (observed via return_field).
394        let field = decoded.return_field(&Schema::empty()).unwrap();
395        assert_eq!(field.data_type(), &DataType::Int32);
396        assert!(field.is_nullable());
397
398        // Same shared container + same index → equal to the original.
399        assert_eq!(decoded, &expr);
400    }
401
402    #[test]
403    fn rejects_non_scalar_subquery_node() {
404        let node = column_node("a");
405        let results = ScalarSubqueryResults::new(1);
406        let decoder = UnreachableDecoder;
407        let schema = Schema::empty();
408        let dec_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
409
410        let err =
411            ScalarSubqueryExpr::try_from_proto(&node, &dec_ctx, &results).unwrap_err();
412        assert!(matches!(
413            err,
414            DataFusionError::Internal(msg)
415                if msg.contains("PhysicalExprNode is not a ScalarSubqueryExpr")
416        ));
417    }
418
419    #[test]
420    fn rejects_missing_data_type() {
421        let node = proto_scalar_subquery_node(None, false, 0);
422        let results = ScalarSubqueryResults::new(1);
423        let decoder = UnreachableDecoder;
424        let schema = Schema::empty();
425        let dec_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
426
427        let err =
428            ScalarSubqueryExpr::try_from_proto(&node, &dec_ctx, &results).unwrap_err();
429        assert!(matches!(
430            err,
431            DataFusionError::Internal(msg)
432                if msg.contains("ScalarSubqueryExpr is missing required field 'data_type'")
433        ));
434    }
435}
datafusion_physical_expr/scalar_subquery.rs

datafusion_physical_expr/
scalar_subquery.rs
