query/query_engine/
default_serializer.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use bytes::Bytes;
18use catalog::CatalogManagerRef;
19use common_error::ext::BoxedError;
20use common_function::function::FunctionContext;
21use common_function::function_registry::FUNCTION_REGISTRY;
22use common_query::error::RegisterUdfSnafu;
23use common_query::logical_plan::SubstraitPlanDecoder;
24use datafusion::catalog::{CatalogProviderList, TableProvider};
25use datafusion::common::{DFSchema, DataFusionError, TableReference, not_impl_err, substrait_err};
26use datafusion::error::Result;
27use datafusion::execution::context::SessionState;
28use datafusion::execution::registry::SerializerRegistry;
29use datafusion::execution::{FunctionRegistry, SessionStateBuilder};
30use datafusion::logical_expr::{Extension, LogicalPlan};
31use datafusion_expr::{Expr, UserDefinedLogicalNode};
32use greptime_proto::substrait_extension::MergeScan as PbMergeScan;
33use promql::functions::register_promql_functions;
34use prost::Message;
35use session::context::{QueryContext, QueryContextRef};
36use snafu::ResultExt;
37use substrait::df_logical_plan::consumer::{DefaultSubstraitConsumer, SubstraitConsumer};
38use substrait::error::{DecodeDfPlanSnafu, DecodeRelSnafu};
39use substrait::extension_serializer::ExtensionSerializer;
40use substrait::substrait_proto_df::proto::{
41    ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, Plan, Type,
42};
43use substrait::{DFLogicalSubstraitConvertor, Extensions, SubstraitPlan};
44
45use crate::dist_plan::MergeScanLogicalPlan;
46use crate::query_engine::QueryEngineState;
47
48/// Extended [`substrait::extension_serializer::ExtensionSerializer`] but supports [`MergeScanLogicalPlan`] serialization.
49#[derive(Debug)]
50pub struct DefaultSerializer;
51
52impl SerializerRegistry for DefaultSerializer {
53    fn serialize_logical_plan(&self, node: &dyn UserDefinedLogicalNode) -> Result<Vec<u8>> {
54        if node.name() == MergeScanLogicalPlan::name() {
55            let merge_scan = node
56                .as_any()
57                .downcast_ref::<MergeScanLogicalPlan>()
58                .expect("Failed to downcast to MergeScanLogicalPlan");
59
60            let input = merge_scan.input();
61            let is_placeholder = merge_scan.is_placeholder();
62            let input = DFLogicalSubstraitConvertor
63                .encode(input, DefaultSerializer)
64                .map_err(|e| DataFusionError::External(Box::new(e)))?
65                .to_vec();
66
67            Ok(PbMergeScan {
68                is_placeholder,
69                input,
70            }
71            .encode_to_vec())
72        } else {
73            ExtensionSerializer.serialize_logical_plan(node)
74        }
75    }
76
77    fn deserialize_logical_plan(
78        &self,
79        name: &str,
80        bytes: &[u8],
81    ) -> Result<Arc<dyn UserDefinedLogicalNode>> {
82        if name == MergeScanLogicalPlan::name() {
83            // This registry has no session state; the query engine decodes `MergeScan` payloads.
84            Err(DataFusionError::Substrait(format!(
85                "Unsupported plan node: {name}"
86            )))
87        } else {
88            ExtensionSerializer.deserialize_logical_plan(name, bytes)
89        }
90    }
91}
92
93/// Async `MergeScan` payload decoding over DataFusion's default Substrait consumer.
94struct MergeScanSubstraitConsumer<'a> {
95    inner: DefaultSubstraitConsumer<'a>,
96    session_state: &'a SessionState,
97    catalog_manager: Option<CatalogManagerRef>,
98    payload: bool,
99}
100
101impl MergeScanSubstraitConsumer<'_> {
102    fn query_ctx(&self) -> QueryContextRef {
103        self.session_state
104            .config()
105            .get_extension()
106            .unwrap_or_else(QueryContext::arc)
107    }
108
109    async fn decode_payload(&self, payload: Vec<u8>) -> Result<LogicalPlan> {
110        decode_plan(Bytes::from(payload), self.session_state, true)
111            .await
112            .map_err(|e| DataFusionError::External(Box::new(e)))
113    }
114}
115
116#[async_trait::async_trait]
117impl SubstraitConsumer for MergeScanSubstraitConsumer<'_> {
118    async fn resolve_table_ref(
119        &self,
120        table_ref: &TableReference,
121    ) -> Result<Option<Arc<dyn TableProvider>>> {
122        if !self.payload {
123            return self.inner.resolve_table_ref(table_ref).await;
124        }
125
126        let query_ctx = self.query_ctx();
127        let resolved = table_ref
128            .clone()
129            .resolve(query_ctx.current_catalog(), &query_ctx.current_schema());
130        // Manual states retain DataFusion's own catalog and information-schema behavior.
131        let Some(manager) = &self.catalog_manager else {
132            return self
133                .inner
134                .resolve_table_ref(&TableReference::from(resolved))
135                .await;
136        };
137        let catalog_list: Arc<dyn CatalogProviderList> = Arc::new(
138            catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx(
139                manager.clone(),
140                query_ctx,
141            ),
142        );
143        let schema_name = resolved.schema.as_ref();
144        let catalog_name = resolved.catalog.as_ref();
145        let schema = catalog_list
146            .catalog(catalog_name)
147            .ok_or_else(|| {
148                DataFusionError::Plan(format!("failed to resolve catalog: {catalog_name}"))
149            })?
150            .schema(schema_name)
151            .ok_or_else(|| {
152                DataFusionError::Plan(format!("failed to resolve schema: {schema_name}"))
153            })?;
154        schema.table(&resolved.table).await
155    }
156
157    fn get_extensions(&self) -> &Extensions {
158        self.inner.get_extensions()
159    }
160
161    fn get_function_registry(&self) -> &impl FunctionRegistry {
162        self.inner.get_function_registry()
163    }
164
165    fn push_outer_schema(&self, schema: Arc<DFSchema>) {
166        self.inner.push_outer_schema(schema)
167    }
168
169    fn pop_outer_schema(&self) {
170        self.inner.pop_outer_schema()
171    }
172
173    fn get_outer_schema(&self, steps_out: usize) -> Option<Arc<DFSchema>> {
174        self.inner.get_outer_schema(steps_out)
175    }
176
177    fn push_lambda_parameters(
178        &self,
179        lambda_parameters: &[Type],
180        input_schema: &DFSchema,
181    ) -> Result<Vec<String>> {
182        self.inner
183            .push_lambda_parameters(lambda_parameters, input_schema)
184    }
185
186    fn pop_lambda_parameters(&self) {
187        self.inner.pop_lambda_parameters()
188    }
189
190    fn lambda_variable(&self, steps_out: usize, field_idx: usize) -> Result<Expr> {
191        self.inner.lambda_variable(steps_out, field_idx)
192    }
193
194    async fn consume_extension_leaf(&self, rel: &ExtensionLeafRel) -> Result<LogicalPlan> {
195        // Unknown and detail-less leaves keep the default consumer's errors.
196        let Some(detail) = rel.detail.as_ref() else {
197            return self.inner.consume_extension_leaf(rel).await;
198        };
199        if detail.type_url != MergeScanLogicalPlan::name() {
200            return self.inner.consume_extension_leaf(rel).await;
201        }
202
203        let merge_scan = PbMergeScan::decode(detail.value.as_ref()).map_err(|e| {
204            DataFusionError::Substrait(format!("Failed to decode the MergeScan plan node: {e}"))
205        })?;
206
207        let input = self.decode_payload(merge_scan.input).await?;
208
209        Ok(MergeScanLogicalPlan::new(input, merge_scan.is_placeholder).into_logical_plan())
210    }
211
212    // Keep child dispatch on this consumer; delegating the parent bypasses async MergeScan decoding.
213    async fn consume_extension_single(&self, rel: &ExtensionSingleRel) -> Result<LogicalPlan> {
214        let Some(detail) = &rel.detail else {
215            return substrait_err!("Unexpected empty detail in ExtensionSingleRel");
216        };
217        let plan = self
218            .session_state
219            .serializer_registry()
220            .deserialize_logical_plan(&detail.type_url, &detail.value)?;
221        let Some(input_rel) = &rel.input else {
222            return substrait_err!(
223                "ExtensionSingleRel missing input rel, try using ExtensionLeafRel instead"
224            );
225        };
226        let input_plan = self.consume_rel(input_rel).await?;
227        let plan = plan.with_exprs_and_inputs(plan.expressions(), vec![input_plan])?;
228        Ok(LogicalPlan::Extension(Extension { node: plan }))
229    }
230
231    // Multi-input extension children likewise need this consumer's dispatch.
232    async fn consume_extension_multi(&self, rel: &ExtensionMultiRel) -> Result<LogicalPlan> {
233        let Some(detail) = &rel.detail else {
234            return substrait_err!("Unexpected empty detail in ExtensionMultiRel");
235        };
236        let plan = self
237            .session_state
238            .serializer_registry()
239            .deserialize_logical_plan(&detail.type_url, &detail.value)?;
240        let mut inputs = Vec::with_capacity(rel.inputs.len());
241        for input in &rel.inputs {
242            let input_plan = self.consume_rel(input).await?;
243            inputs.push(input_plan);
244        }
245        let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?;
246        Ok(LogicalPlan::Extension(Extension { node: plan }))
247    }
248}
249
250/// Decode with async MergeScan handling, preserving the default consumer's extension checks.
251async fn decode_plan(
252    bytes: Bytes,
253    state: &SessionState,
254    payload: bool,
255) -> substrait::error::Result<LogicalPlan> {
256    let plan = Plan::decode(bytes).context(DecodeRelSnafu)?;
257    let extensions = Extensions::try_from(&plan.extensions).context(DecodeDfPlanSnafu)?;
258    if !extensions.type_variations.is_empty() {
259        return not_impl_err!("Type variation extensions are not supported")
260            .context(DecodeDfPlanSnafu);
261    }
262
263    let catalog_manager = state
264        .config()
265        .get_extension::<QueryEngineState>()
266        .map(|engine_state| engine_state.catalog_manager().clone());
267    let consumer = MergeScanSubstraitConsumer {
268        inner: DefaultSubstraitConsumer::new(&extensions, state),
269        session_state: state,
270        catalog_manager,
271        payload,
272    };
273
274    DFLogicalSubstraitConvertor
275        .decode_with_consumer(&plan, state, &consumer)
276        .await
277}
278
279/// The datafusion `[LogicalPlan]` decoder.
280pub struct DefaultPlanDecoder {
281    session_state: SessionState,
282    query_ctx: QueryContextRef,
283}
284
285impl DefaultPlanDecoder {
286    pub fn new(
287        session_state: SessionState,
288        query_ctx: &QueryContextRef,
289    ) -> crate::error::Result<Self> {
290        Ok(Self {
291            session_state,
292            query_ctx: query_ctx.clone(),
293        })
294    }
295}
296
297/// Registers Greptime functions after building the request decoder state.
298///
299/// DataFusion aliases can rebind Greptime `date_format` to built-in `to_char`; registration must
300/// happen after `SessionStateBuilder::build` on the state that decodes the request plan.
301fn register_greptime_functions(
302    session_state: &mut SessionState,
303    query_ctx: &QueryContextRef,
304) -> common_query::error::Result<()> {
305    for func in FUNCTION_REGISTRY.scalar_functions() {
306        let udf = func.provide(FunctionContext {
307            query_ctx: query_ctx.clone(),
308            state: Default::default(),
309        });
310        session_state
311            .register_udf(Arc::new(udf))
312            .context(RegisterUdfSnafu { name: func.name() })?;
313    }
314
315    for func in FUNCTION_REGISTRY.aggregate_functions() {
316        let name = func.name().to_string();
317        session_state
318            .register_udaf(Arc::new(func))
319            .context(RegisterUdfSnafu { name })?;
320    }
321
322    register_promql_functions(session_state);
323
324    Ok(())
325}
326
327#[async_trait::async_trait]
328impl SubstraitPlanDecoder for DefaultPlanDecoder {
329    async fn decode(
330        &self,
331        message: bytes::Bytes,
332        catalog_list: Arc<dyn CatalogProviderList>,
333        optimize: bool,
334    ) -> common_query::error::Result<LogicalPlan> {
335        let mut session_state = SessionStateBuilder::new_from_existing(self.session_state.clone())
336            .with_serializer_registry(Arc::new(DefaultSerializer))
337            .with_catalog_list(catalog_list)
338            .build();
339        // Re-register after the build to avoid Greptime UDF alias collisions.
340        register_greptime_functions(&mut session_state, &self.query_ctx)?;
341
342        let logical_plan = decode_plan(message, &session_state, false)
343            .await
344            .map_err(BoxedError::new)
345            .context(common_query::error::DecodePlanSnafu)?;
346
347        if optimize {
348            self.session_state
349                .optimize(&logical_plan)
350                .map_err(Into::into)
351        } else {
352            Ok(logical_plan)
353        }
354    }
355}
356
357#[cfg(test)]
358mod tests {
359    use std::any::Any;
360    use std::sync::atomic::{AtomicBool, Ordering};
361
362    use catalog::error::Result as CatalogResult;
363    use catalog::{CatalogManager, RegisterTableRequest};
364    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, NUMBERS_TABLE_ID};
365    use common_function::aggrs::aggr_wrapper::StateWrapper;
366    use common_function::aggrs::aggr_wrapper::fix_order::FixStateUdafOrderingAnalyzer;
367    use common_query::logical_plan::SubstraitPlanDecoderRef;
368    use common_query::native_histogram::native_histogram_value_type;
369    use common_time::Timezone;
370    use datafusion::catalog::TableProvider;
371    use datafusion::datasource::MemTable;
372    use datafusion::logical_expr::Extension;
373    use datafusion::optimizer::AnalyzerRule;
374    use datafusion_expr::expr::{AggregateFunction, Cast, InSubquery, ScalarFunction, Sort};
375    use datafusion_expr::{
376        AggregateUDF, Expr, LogicalPlanBuilder, LogicalTableSource, ScalarUDF, Subquery, col, lit,
377    };
378    use datatypes::arrow::datatypes::{
379        DataType as ArrowDataType, Field, Schema, SchemaRef, TimeUnit,
380    };
381    use datatypes::data_type::DataType;
382    use futures::stream::BoxStream;
383    use promql::extension_plan::{RangeManipulate, SeriesNormalize, UnionDistinctOn};
384    use promql::functions::{
385        MixedRange, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramCount,
386        NativeHistogramDrop, PromqlFloatToString,
387    };
388    use session::context::QueryContext;
389    use table::TableRef;
390    use table::metadata::{TableId, TableInfoRef};
391    use table::table::numbers::{NUMBERS_TABLE_NAME, NumbersTable};
392
393    use super::*;
394    use crate::QueryEngineFactory;
395    use crate::dummy_catalog::DummyCatalogList;
396    use crate::optimizer::test_util::mock_table_provider;
397    use crate::options::QueryOptions;
398
399    fn mock_plan(schema: SchemaRef) -> LogicalPlan {
400        let table_source = LogicalTableSource::new(schema);
401        let projection = None;
402        let builder =
403            LogicalPlanBuilder::scan("devices", Arc::new(table_source), projection).unwrap();
404
405        builder
406            .filter(col("k0").eq(lit("hello")))
407            .unwrap()
408            .build()
409            .unwrap()
410    }
411
412    #[tokio::test]
413    async fn test_serializer_decode_plan() {
414        let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap();
415        let factory = QueryEngineFactory::new(
416            catalog_list,
417            None,
418            None,
419            None,
420            None,
421            false,
422            QueryOptions::default(),
423        );
424
425        let engine = factory.query_engine();
426
427        let table_provider = Arc::new(mock_table_provider(1.into()));
428        let plan = mock_plan(table_provider.schema().clone());
429
430        let bytes = DFLogicalSubstraitConvertor
431            .encode(&plan, DefaultSerializer)
432            .unwrap();
433
434        let plan_decoder = engine
435            .engine_context(QueryContext::arc())
436            .new_plan_decoder()
437            .unwrap();
438        let catalog_list = Arc::new(DummyCatalogList::with_table_provider(table_provider));
439
440        let decode_plan = plan_decoder
441            .decode(bytes, catalog_list, false)
442            .await
443            .unwrap();
444
445        assert_eq!(
446            "Filter: devices.k0 = Utf8(\"hello\")
447  TableScan: devices",
448            decode_plan.to_string(),
449        );
450    }
451
452    #[tokio::test]
453    async fn test_serializer_decode_native_histogram_udf() {
454        let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap();
455        let factory = QueryEngineFactory::new(
456            catalog_list,
457            None,
458            None,
459            None,
460            None,
461            false,
462            QueryOptions::default(),
463        );
464        let engine = factory.query_engine();
465        let schema = Arc::new(Schema::new(vec![Field::new(
466            "histogram",
467            native_histogram_value_type().as_arrow_type(),
468            true,
469        )]));
470        let plan = LogicalPlanBuilder::scan(
471            "devices",
472            Arc::new(LogicalTableSource::new(schema.clone())),
473            None,
474        )
475        .unwrap()
476        .aggregate(
477            Vec::<Expr>::new(),
478            vec![
479                Arc::new(NativeHistogramAggSum::aggregate_udf())
480                    .call(vec![col("histogram")])
481                    .alias("sum"),
482                Arc::new(NativeHistogramAggAvg::aggregate_udf())
483                    .call(vec![col("histogram")])
484                    .alias("avg"),
485            ],
486        )
487        .unwrap()
488        .project(vec![
489            Expr::ScalarFunction(ScalarFunction {
490                func: Arc::new(NativeHistogramCount::scalar_udf()),
491                args: vec![col("sum")],
492            }),
493            Expr::ScalarFunction(ScalarFunction {
494                func: Arc::new(NativeHistogramDrop::bool_false_udf(
495                    "ignored annotation".to_string(),
496                    None,
497                )),
498                args: vec![col("sum")],
499            }),
500            Expr::ScalarFunction(ScalarFunction {
501                func: Arc::new(NativeHistogramDrop::bool_true_udf(
502                    "ignored annotation".to_string(),
503                    None,
504                )),
505                args: vec![col("sum")],
506            }),
507            Expr::ScalarFunction(ScalarFunction {
508                func: Arc::new(NativeHistogramDrop::float_null_udf(
509                    "ignored annotation".to_string(),
510                    None,
511                )),
512                args: vec![col("sum")],
513            }),
514            Expr::ScalarFunction(ScalarFunction {
515                func: Arc::new(PromqlFloatToString::scalar_udf()),
516                args: vec![lit(2.0)],
517            }),
518        ])
519        .unwrap()
520        .build()
521        .unwrap();
522        let bytes = DFLogicalSubstraitConvertor
523            .encode(&plan, DefaultSerializer)
524            .unwrap();
525        let table_provider = Arc::new(MemTable::try_new(schema, vec![vec![]]).unwrap());
526        let plan_decoder = engine
527            .engine_context(QueryContext::arc())
528            .new_plan_decoder()
529            .unwrap();
530
531        let decoded = plan_decoder
532            .decode(
533                bytes,
534                Arc::new(DummyCatalogList::with_table_provider(table_provider)),
535                false,
536            )
537            .await
538            .unwrap();
539
540        let decoded = decoded.to_string();
541        assert!(decoded.contains("prom_native_histogram_count"));
542        assert!(decoded.contains("prom_native_histogram_drop_bool"));
543        assert!(decoded.contains("prom_native_histogram_keep_bool"));
544        assert!(decoded.contains("prom_native_histogram_drop_float"));
545        assert!(decoded.contains(NativeHistogramAggSum::name()));
546        assert!(decoded.contains(NativeHistogramAggAvg::name()));
547        assert!(decoded.contains(PromqlFloatToString::name()));
548
549        let schema = Arc::new(Schema::new(vec![
550            Field::new(
551                "timestamp",
552                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
553                false,
554            ),
555            Field::new("float", ArrowDataType::Float64, true),
556            Field::new(
557                "histogram",
558                native_histogram_value_type().as_arrow_type(),
559                true,
560            ),
561        ]));
562        let input = LogicalPlanBuilder::scan(
563            "devices",
564            Arc::new(LogicalTableSource::new(schema.clone())),
565            None,
566        )
567        .unwrap()
568        .build()
569        .unwrap();
570        let input = LogicalPlan::Extension(Extension {
571            node: Arc::new(
572                RangeManipulate::new(
573                    0,
574                    1000,
575                    1000,
576                    0,
577                    1000,
578                    "timestamp".to_string(),
579                    vec!["float".to_string(), "histogram".to_string()],
580                    input,
581                )
582                .unwrap(),
583            ),
584        });
585        let plan = LogicalPlanBuilder::from(input)
586            .project(vec![
587                Expr::ScalarFunction(ScalarFunction {
588                    func: Arc::new(MixedRange::float_udf(None)),
589                    args: vec![
590                        lit("last_over_time"),
591                        col("timestamp_range"),
592                        col("float"),
593                        col("histogram"),
594                    ],
595                })
596                .alias("mixed_float"),
597                Expr::ScalarFunction(ScalarFunction {
598                    func: Arc::new(MixedRange::histogram_udf(None)),
599                    args: vec![
600                        lit("last_over_time"),
601                        col("timestamp_range"),
602                        col("float"),
603                        col("histogram"),
604                    ],
605                })
606                .alias("mixed_histogram"),
607            ])
608            .unwrap()
609            .build()
610            .unwrap();
611        let bytes = DFLogicalSubstraitConvertor
612            .encode(&plan, DefaultSerializer)
613            .unwrap();
614        let table_provider = Arc::new(MemTable::try_new(schema, vec![vec![]]).unwrap());
615        let decoded = plan_decoder
616            .decode(
617                bytes,
618                Arc::new(DummyCatalogList::with_table_provider(table_provider)),
619                false,
620            )
621            .await
622            .unwrap()
623            .to_string();
624        assert!(decoded.contains("prom_mixed_range_float"));
625        assert!(decoded.contains("prom_mixed_range_histogram"));
626    }
627
628    #[tokio::test]
629    async fn test_serializer_decode_merge_scan() {
630        let engine = QueryEngineFactory::new(
631            numbers_catalog_manager(),
632            None,
633            None,
634            None,
635            None,
636            false,
637            QueryOptions::default(),
638        )
639        .query_engine();
640        let plan = merge_scan(numbers_scan(), true);
641        let decoded = plan_decoder(engine.as_ref())
642            .decode(encode_plan(&plan), request_catalog_list(), false)
643            .await
644            .unwrap();
645
646        let LogicalPlan::Extension(extension) = &decoded else {
647            panic!("Expect a MergeScan plan, got: {decoded}");
648        };
649        let merge_scan = extension
650            .node
651            .as_any()
652            .downcast_ref::<MergeScanLogicalPlan>()
653            .expect("Expect a MergeScan plan node");
654        assert!(merge_scan.is_placeholder());
655        assert_eq!(merge_scan.input().to_string(), numbers_scan().to_string());
656    }
657
658    /// Payload table binding uses the engine catalog, not the request catalog.
659    #[tokio::test]
660    async fn test_serializer_decode_merge_scan_with_engine_catalog() {
661        let engine = QueryEngineFactory::new(
662            numbers_catalog_manager(),
663            None,
664            None,
665            None,
666            None,
667            false,
668            QueryOptions::default(),
669        )
670        .query_engine();
671        let decoded = plan_decoder(engine.as_ref())
672            .decode(
673                encode_plan(&merge_scan(numbers_scan(), false)),
674                request_catalog_list(),
675                false,
676            )
677            .await
678            .unwrap();
679
680        let LogicalPlan::Extension(extension) = &decoded else {
681            panic!("Expect a MergeScan plan, got: {decoded}");
682        };
683        let merge_scan = extension
684            .node
685            .as_any()
686            .downcast_ref::<MergeScanLogicalPlan>()
687            .expect("Expect a MergeScan plan node");
688        let LogicalPlan::TableScan(scan) = merge_scan.input() else {
689            panic!("Expect a table scan, got: {}", merge_scan.input());
690        };
691        assert_eq!(scan.table_name.table(), NUMBERS_TABLE_NAME);
692        assert_eq!(scan.source.schema().fields().len(), 1);
693        assert_eq!(scan.source.schema().field(0).name(), "number");
694    }
695
696    /// Nested `MergeScan` payloads decode recursively.
697    #[tokio::test]
698    async fn test_serializer_decode_nested_merge_scan() {
699        let engine = QueryEngineFactory::new(
700            numbers_catalog_manager(),
701            None,
702            None,
703            None,
704            None,
705            false,
706            QueryOptions::default(),
707        )
708        .query_engine();
709        let inner = merge_scan(numbers_scan(), false);
710        let input = LogicalPlanBuilder::from(inner)
711            .filter(col("number").lt(lit(10u32)))
712            .unwrap()
713            .build()
714            .unwrap();
715        let plan = merge_scan(input, false);
716
717        let decoded = plan_decoder(engine.as_ref())
718            .decode(encode_plan(&plan), request_catalog_list(), false)
719            .await
720            .unwrap();
721
722        assert_eq!(decoded.to_string(), plan.to_string());
723        assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2);
724    }
725
726    fn greptime_date_format_udf(query_ctx: &QueryContextRef) -> Arc<ScalarUDF> {
727        FUNCTION_REGISTRY
728            .get_function("date_format")
729            .expect("`date_format` is a function of GreptimeDB")
730            .provide(FunctionContext {
731                query_ctx: query_ctx.clone(),
732                state: Default::default(),
733            })
734            .into()
735    }
736
737    /// Uses a non-default timezone to distinguish the Greptime implementation.
738    fn date_format_query_ctx() -> QueryContextRef {
739        let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
740        query_ctx.set_timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap());
741        Arc::new(query_ctx)
742    }
743
744    fn greptime_date_format_expr(query_ctx: &QueryContextRef, timestamp: Expr) -> Expr {
745        Expr::ScalarFunction(ScalarFunction {
746            func: greptime_date_format_udf(query_ctx),
747            args: vec![timestamp, lit("%Y-%m-%d %H:%i:%S")],
748        })
749    }
750
751    /// Includes `MergeScan` payload expressions, which are not node inputs.
752    fn scalar_functions_of(plan: &LogicalPlan) -> Vec<Arc<ScalarUDF>> {
753        use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
754
755        let mut udfs = Vec::new();
756        for expr in plan.expressions() {
757            expr.apply(|expr| {
758                if let Expr::ScalarFunction(func) = expr {
759                    udfs.push(Arc::clone(&func.func));
760                }
761                Ok(TreeNodeRecursion::Continue)
762            })
763            .unwrap();
764        }
765
766        if let LogicalPlan::Extension(extension) = plan
767            && let Some(merge_scan) = extension
768                .node
769                .as_any()
770                .downcast_ref::<MergeScanLogicalPlan>()
771        {
772            udfs.extend(scalar_functions_of(merge_scan.input()));
773        }
774
775        for input in plan.inputs() {
776            udfs.extend(scalar_functions_of(input));
777        }
778
779        udfs
780    }
781
782    /// Asserts the Greptime `date_format` binding rather than DataFusion's `to_char` alias.
783    fn assert_only_function_is_greptime_date_format(
784        plan: &LogicalPlan,
785        query_ctx: &QueryContextRef,
786    ) {
787        let udfs = scalar_functions_of(plan);
788        assert_eq!(udfs.len(), 1, "expected one function in: {plan}");
789
790        let udf = &udfs[0];
791        assert_eq!(udf.name(), "date_format", "got: {plan}");
792        assert_eq!(
793            udf.as_ref(),
794            greptime_date_format_udf(query_ctx).as_ref(),
795            "`date_format` is not bound to the implementation of GreptimeDB: {plan}"
796        );
797    }
798
799    /// Request state registration keeps Greptime `date_format` ahead of DataFusion aliases.
800    #[tokio::test]
801    async fn test_serializer_decode_keeps_greptime_date_format() {
802        // Repeat to expose nondeterministic function-registration order.
803        for _ in 0..8 {
804            let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap();
805            let factory = QueryEngineFactory::new(
806                catalog_list,
807                None,
808                None,
809                None,
810                None,
811                false,
812                QueryOptions::default(),
813            );
814            let engine = factory.query_engine();
815
816            let table_provider = Arc::new(mock_table_provider(1.into()));
817            let query_ctx = date_format_query_ctx();
818            let plan = LogicalPlanBuilder::scan(
819                "devices",
820                Arc::new(LogicalTableSource::new(table_provider.schema().clone())),
821                None,
822            )
823            .unwrap()
824            .project(vec![greptime_date_format_expr(&query_ctx, col("ts"))])
825            .unwrap()
826            .build()
827            .unwrap();
828
829            let bytes = DFLogicalSubstraitConvertor
830                .encode(&plan, DefaultSerializer)
831                .unwrap();
832            let plan_decoder = engine
833                .engine_context(query_ctx.clone())
834                .new_plan_decoder()
835                .unwrap();
836            let decoded = plan_decoder
837                .decode(
838                    bytes,
839                    Arc::new(DummyCatalogList::with_table_provider(table_provider)),
840                    false,
841                )
842                .await
843                .unwrap();
844
845            assert_only_function_is_greptime_date_format(&decoded, &query_ctx);
846        }
847    }
848
849    /// Nested payload decoding retains the Greptime `date_format` binding.
850    #[tokio::test]
851    async fn test_serializer_decode_merge_scan_payload_keeps_greptime_date_format() {
852        let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap();
853        catalog_manager
854            .register_table_sync(RegisterTableRequest {
855                catalog: DEFAULT_CATALOG_NAME.to_string(),
856                schema: DEFAULT_SCHEMA_NAME.to_string(),
857                table_name: NUMBERS_TABLE_NAME.to_string(),
858                table_id: NUMBERS_TABLE_ID,
859                table: NumbersTable::table(NUMBERS_TABLE_ID),
860            })
861            .unwrap();
862
863        for _ in 0..8 {
864            let factory = QueryEngineFactory::new(
865                catalog_manager.clone(),
866                None,
867                None,
868                None,
869                None,
870                false,
871                QueryOptions::default(),
872            );
873            let engine = factory.query_engine();
874            let query_ctx = date_format_query_ctx();
875
876            let input = LogicalPlanBuilder::scan(
877                NUMBERS_TABLE_NAME,
878                Arc::new(LogicalTableSource::new(
879                    NumbersTable::schema().arrow_schema().clone(),
880                )),
881                None,
882            )
883            .unwrap()
884            .project(vec![greptime_date_format_expr(
885                &query_ctx,
886                Expr::Cast(Cast::new(
887                    Box::new(col("number")),
888                    ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
889                )),
890            )])
891            .unwrap()
892            .build()
893            .unwrap();
894            let plan = MergeScanLogicalPlan::new(input, false).into_logical_plan();
895
896            let bytes = DFLogicalSubstraitConvertor
897                .encode(&plan, DefaultSerializer)
898                .unwrap();
899            let plan_decoder = engine
900                .engine_context(query_ctx.clone())
901                .new_plan_decoder()
902                .unwrap();
903            let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new(
904                mock_table_provider(1.into()),
905            )));
906
907            let decoded = plan_decoder
908                .decode(bytes, catalog_list, false)
909                .await
910                .unwrap();
911
912            let LogicalPlan::Extension(extension) = &decoded else {
913                panic!("Expect a MergeScan plan, got: {decoded}");
914            };
915            let merge_scan = extension
916                .node
917                .as_any()
918                .downcast_ref::<MergeScanLogicalPlan>()
919                .expect("Expect a MergeScan plan node");
920
921            assert_only_function_is_greptime_date_format(merge_scan.input(), &query_ctx);
922        }
923    }
924
925    /// Registers a `numbers` table in a memory catalog manager.
926    fn numbers_catalog_manager() -> CatalogManagerRef {
927        let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap();
928        catalog_manager
929            .register_table_sync(RegisterTableRequest {
930                catalog: DEFAULT_CATALOG_NAME.to_string(),
931                schema: DEFAULT_SCHEMA_NAME.to_string(),
932                table_name: NUMBERS_TABLE_NAME.to_string(),
933                table_id: NUMBERS_TABLE_ID,
934                table: NumbersTable::table(NUMBERS_TABLE_ID),
935            })
936            .unwrap();
937        catalog_manager
938    }
939
940    /// Scan plan over the `numbers` schema shared by the fixtures below.
941    fn numbers_scan() -> LogicalPlan {
942        LogicalPlanBuilder::scan(
943            NUMBERS_TABLE_NAME,
944            Arc::new(LogicalTableSource::new(
945                NumbersTable::schema().arrow_schema().clone(),
946            )),
947            None,
948        )
949        .unwrap()
950        .build()
951        .unwrap()
952    }
953
954    fn merge_scan(input: LogicalPlan, is_placeholder: bool) -> LogicalPlan {
955        MergeScanLogicalPlan::new(input, is_placeholder).into_logical_plan()
956    }
957
958    fn encode_plan(plan: &LogicalPlan) -> Bytes {
959        DFLogicalSubstraitConvertor
960            .encode(plan, DefaultSerializer)
961            .unwrap()
962    }
963
964    fn plan_decoder(engine: &dyn crate::QueryEngine) -> SubstraitPlanDecoderRef {
965        engine
966            .engine_context(QueryContext::arc())
967            .new_plan_decoder()
968            .unwrap()
969    }
970
971    /// Request catalog that resolves any table to a provider payload resolution must not use.
972    fn request_catalog_list() -> Arc<dyn CatalogProviderList> {
973        Arc::new(DummyCatalogList::with_table_provider(Arc::new(
974            mock_table_provider(1.into()),
975        )))
976    }
977
978    /// Payload resolution preserves context and explicit catalog/schema components.
979    #[tokio::test]
980    async fn test_serializer_decode_payload_observes_catalog_and_query_context() {
981        let observed = Arc::new(std::sync::Mutex::new(Vec::new()));
982        let inner = catalog::memory::new_memory_catalog_manager().unwrap();
983        let manager: CatalogManagerRef = Arc::new(InterceptingCatalogManager {
984            inner,
985            resolution: TableResolution::ObserveContext(observed.clone()),
986        });
987        let mut resolver_ctx = QueryContext::arc();
988        Arc::make_mut(&mut resolver_ctx).set_current_catalog("request_cat");
989        Arc::make_mut(&mut resolver_ctx).set_current_schema("request_schema");
990        let engine = QueryEngineFactory::new(
991            manager,
992            None,
993            None,
994            None,
995            None,
996            false,
997            QueryOptions::default(),
998        )
999        .query_engine();
1000        let mut state = engine.engine_context(resolver_ctx.clone()).state().clone();
1001        state.config_mut().set_extension(resolver_ctx);
1002        let extensions = Extensions::default();
1003        let consumer = MergeScanSubstraitConsumer {
1004            inner: DefaultSubstraitConsumer::new(&extensions, &state),
1005            session_state: &state,
1006            catalog_manager: state
1007                .config()
1008                .get_extension::<QueryEngineState>()
1009                .map(|s| s.catalog_manager().clone()),
1010            payload: true,
1011        };
1012        for reference in [
1013            TableReference::bare(NUMBERS_TABLE_NAME),
1014            TableReference::partial("explicit_schema", NUMBERS_TABLE_NAME),
1015            TableReference::full(
1016                DEFAULT_CATALOG_NAME,
1017                DEFAULT_SCHEMA_NAME,
1018                NUMBERS_TABLE_NAME,
1019            ),
1020        ] {
1021            consumer
1022                .resolve_table_ref(&reference)
1023                .await
1024                .unwrap()
1025                .unwrap();
1026        }
1027        assert_eq!(
1028            *observed.lock().unwrap(),
1029            vec![
1030                [
1031                    "request_cat".to_string(),
1032                    "request_schema".to_string(),
1033                    NUMBERS_TABLE_NAME.to_string(),
1034                    "request_cat".to_string(),
1035                    "request_schema".to_string(),
1036                ],
1037                [
1038                    "request_cat".to_string(),
1039                    "explicit_schema".to_string(),
1040                    NUMBERS_TABLE_NAME.to_string(),
1041                    "request_cat".to_string(),
1042                    "request_schema".to_string(),
1043                ],
1044                [
1045                    DEFAULT_CATALOG_NAME.to_string(),
1046                    DEFAULT_SCHEMA_NAME.to_string(),
1047                    NUMBERS_TABLE_NAME.to_string(),
1048                    "request_cat".to_string(),
1049                    "request_schema".to_string(),
1050                ],
1051            ],
1052        );
1053    }
1054
1055    /// How [`InterceptingCatalogManager`] resolves payload tables.
1056    enum TableResolution {
1057        /// Records catalog lookup and request context values.
1058        ObserveContext(Arc<std::sync::Mutex<Vec<[String; 5]>>>),
1059        /// Resolves once another task on the caller runtime signals.
1060        AwaitSignal {
1061            notify: Arc<tokio::sync::Notify>,
1062            entered: Arc<AtomicBool>,
1063        },
1064    }
1065
1066    /// Engine catalog manager that takes over `table` resolution while delegating the rest.
1067    struct InterceptingCatalogManager {
1068        inner: CatalogManagerRef,
1069        resolution: TableResolution,
1070    }
1071
1072    #[async_trait::async_trait]
1073    impl CatalogManager for InterceptingCatalogManager {
1074        fn as_any(&self) -> &dyn Any {
1075            self.inner.as_any()
1076        }
1077
1078        async fn catalog_names(&self) -> CatalogResult<Vec<String>> {
1079            self.inner.catalog_names().await
1080        }
1081
1082        async fn schema_names(
1083            &self,
1084            catalog: &str,
1085            query_ctx: Option<&QueryContext>,
1086        ) -> CatalogResult<Vec<String>> {
1087            self.inner.schema_names(catalog, query_ctx).await
1088        }
1089
1090        async fn table_names(
1091            &self,
1092            catalog: &str,
1093            schema: &str,
1094            query_ctx: Option<&QueryContext>,
1095        ) -> CatalogResult<Vec<String>> {
1096            self.inner.table_names(catalog, schema, query_ctx).await
1097        }
1098
1099        async fn catalog_exists(&self, catalog: &str) -> CatalogResult<bool> {
1100            self.inner.catalog_exists(catalog).await
1101        }
1102
1103        async fn schema_exists(
1104            &self,
1105            catalog: &str,
1106            schema: &str,
1107            query_ctx: Option<&QueryContext>,
1108        ) -> CatalogResult<bool> {
1109            self.inner.schema_exists(catalog, schema, query_ctx).await
1110        }
1111
1112        async fn table_exists(
1113            &self,
1114            catalog: &str,
1115            schema: &str,
1116            table: &str,
1117            query_ctx: Option<&QueryContext>,
1118        ) -> CatalogResult<bool> {
1119            self.inner
1120                .table_exists(catalog, schema, table, query_ctx)
1121                .await
1122        }
1123
1124        async fn table(
1125            &self,
1126            catalog: &str,
1127            schema: &str,
1128            table_name: &str,
1129            query_ctx: Option<&QueryContext>,
1130        ) -> CatalogResult<Option<TableRef>> {
1131            match &self.resolution {
1132                TableResolution::ObserveContext(observed) => {
1133                    let query_ctx = query_ctx.unwrap();
1134                    observed.lock().unwrap().push([
1135                        catalog.to_string(),
1136                        schema.to_string(),
1137                        table_name.to_string(),
1138                        query_ctx.current_catalog().to_string(),
1139                        query_ctx.current_schema(),
1140                    ]);
1141                    Ok(Some(NumbersTable::table(NUMBERS_TABLE_ID)))
1142                }
1143
1144                TableResolution::AwaitSignal { notify, entered } => {
1145                    entered.store(true, Ordering::SeqCst);
1146                    notify.notified().await;
1147                    self.inner
1148                        .table(catalog, schema, table_name, query_ctx)
1149                        .await
1150                }
1151            }
1152        }
1153
1154        async fn table_info_by_id(&self, table_id: TableId) -> CatalogResult<Option<TableInfoRef>> {
1155            self.inner.table_info_by_id(table_id).await
1156        }
1157
1158        async fn tables_by_ids(
1159            &self,
1160            catalog: &str,
1161            schema: &str,
1162            table_ids: &[TableId],
1163        ) -> CatalogResult<Vec<TableRef>> {
1164            self.inner.tables_by_ids(catalog, schema, table_ids).await
1165        }
1166
1167        fn tables<'a>(
1168            &'a self,
1169            catalog: &'a str,
1170            schema: &'a str,
1171            query_ctx: Option<&'a QueryContext>,
1172        ) -> BoxStream<'a, CatalogResult<TableRef>> {
1173            self.inner.tables(catalog, schema, query_ctx)
1174        }
1175    }
1176
1177    fn engine_with_intercepting_catalog(
1178        resolution: TableResolution,
1179    ) -> Arc<dyn crate::QueryEngine> {
1180        let catalog_manager: CatalogManagerRef = Arc::new(InterceptingCatalogManager {
1181            inner: numbers_catalog_manager(),
1182            resolution,
1183        });
1184        QueryEngineFactory::new(
1185            catalog_manager,
1186            None,
1187            None,
1188            None,
1189            None,
1190            false,
1191            QueryOptions::default(),
1192        )
1193        .query_engine()
1194    }
1195
1196    /// A payload lookup parked on a caller-runtime task must keep the decode async.
1197    ///
1198    /// The engine catalog resolves only once a task spawned on this runtime signals it, so a
1199    /// blocking bridge would deadlock the test instead of decoding.
1200    #[tokio::test(flavor = "current_thread")]
1201    async fn test_serializer_decode_merge_scan_awaits_caller_runtime_task() {
1202        let notify = Arc::new(tokio::sync::Notify::new());
1203        let entered = Arc::new(AtomicBool::new(false));
1204        let engine = engine_with_intercepting_catalog(TableResolution::AwaitSignal {
1205            notify: notify.clone(),
1206            entered: entered.clone(),
1207        });
1208        let plan = merge_scan(numbers_scan(), false);
1209        let bytes = encode_plan(&plan);
1210
1211        let decoder = plan_decoder(engine.as_ref());
1212        let mut decode = Box::pin(decoder.decode(bytes, request_catalog_list(), false));
1213        // First poll reaches the payload table lookup, which waits for the signal below.
1214        assert!(futures::poll!(decode.as_mut()).is_pending());
1215        assert!(entered.load(Ordering::SeqCst));
1216
1217        let signal = tokio::spawn(async move {
1218            notify.notify_waiters();
1219        });
1220        tokio::task::yield_now().await;
1221        let decoded = decode.await.unwrap();
1222        signal.await.unwrap();
1223
1224        assert_eq!(decoded.to_string(), plan.to_string());
1225    }
1226
1227    /// A missing engine-catalog table does not fall back to the usable request catalog.
1228    #[tokio::test]
1229    async fn test_serializer_decode_missing_engine_catalog_table() {
1230        let missing_manager: CatalogManagerRef =
1231            catalog::memory::new_memory_catalog_manager().unwrap();
1232        let engine = QueryEngineFactory::new(
1233            missing_manager,
1234            None,
1235            None,
1236            None,
1237            None,
1238            false,
1239            QueryOptions::default(),
1240        )
1241        .query_engine();
1242        let bytes = encode_plan(&merge_scan(numbers_scan(), false));
1243
1244        let err = plan_decoder(engine.as_ref())
1245            .decode(bytes, request_catalog_list(), false)
1246            .await
1247            .unwrap_err();
1248
1249        let err = format!("{err:?}");
1250        assert!(
1251            err.contains("Table not found: greptime.public.numbers"),
1252            "unexpected error: {err}"
1253        );
1254    }
1255
1256    /// Manual states use their own catalog and query-context defaults for payloads.
1257    #[tokio::test]
1258    async fn test_serializer_decode_manual_state_keeps_extension_registry() {
1259        let manager = catalog::memory::new_memory_catalog_manager().unwrap();
1260        manager.register_catalog_sync("manual_cat").unwrap();
1261        manager
1262            .register_schema_sync(catalog::RegisterSchemaRequest {
1263                catalog: "manual_cat".to_string(),
1264                schema: "manual_schema".to_string(),
1265            })
1266            .unwrap();
1267        manager
1268            .register_table_sync(RegisterTableRequest {
1269                catalog: "manual_cat".to_string(),
1270                schema: "manual_schema".to_string(),
1271                table_name: NUMBERS_TABLE_NAME.to_string(),
1272                table_id: NUMBERS_TABLE_ID,
1273                table: NumbersTable::table(NUMBERS_TABLE_ID),
1274            })
1275            .unwrap();
1276        let query_ctx = Arc::new(QueryContext::with("manual_cat", "manual_schema"));
1277        let mut config = datafusion::execution::context::SessionConfig::new()
1278            .with_information_schema(true)
1279            .with_default_catalog_and_schema("manual_cat", "manual_schema");
1280        config.set_extension(query_ctx.clone());
1281        let state = SessionStateBuilder::new()
1282            .with_default_features()
1283            .with_config(config)
1284            .build();
1285        let decoder = DefaultPlanDecoder::new(state, &query_ctx).unwrap();
1286        let plan = LogicalPlan::Extension(Extension {
1287            node: Arc::new(SeriesNormalize::new(
1288                0,
1289                "number",
1290                false,
1291                Vec::new(),
1292                merge_scan(numbers_scan(), false),
1293            )),
1294        });
1295        let request_catalog = Arc::new(
1296            catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx(
1297                manager, query_ctx,
1298            ),
1299        );
1300        let decoded = decoder
1301            .decode(encode_plan(&plan), request_catalog, false)
1302            .await
1303            .unwrap();
1304
1305        assert_eq!(decoded.to_string(), plan.to_string());
1306        assert_eq!(decoded.schema().as_arrow(), plan.schema().as_arrow());
1307    }
1308
1309    /// `MergeScan` decodes under an ordinary single-input extension parent.
1310    #[tokio::test]
1311    async fn test_serializer_decode_merge_scan_under_series_normalize() {
1312        let outer = merge_scan(numbers_scan(), false);
1313        let plan = LogicalPlan::Extension(Extension {
1314            node: Arc::new(SeriesNormalize::new(
1315                0,
1316                "number",
1317                false,
1318                Vec::new(),
1319                outer.clone(),
1320            )),
1321        });
1322        let bytes = encode_plan(&plan);
1323
1324        let decoded = plan_decoder(
1325            QueryEngineFactory::new(
1326                numbers_catalog_manager(),
1327                None,
1328                None,
1329                None,
1330                None,
1331                false,
1332                QueryOptions::default(),
1333            )
1334            .query_engine()
1335            .as_ref(),
1336        )
1337        .decode(bytes, request_catalog_list(), false)
1338        .await
1339        .unwrap();
1340
1341        let LogicalPlan::Extension(extension) = &decoded else {
1342            panic!("Expect a SeriesNormalize plan, got: {decoded}");
1343        };
1344        let normalize = extension
1345            .node
1346            .as_any()
1347            .downcast_ref::<SeriesNormalize>()
1348            .expect("Expect a SeriesNormalize plan node");
1349        assert_eq!(normalize.inputs()[0].to_string(), outer.to_string());
1350    }
1351
1352    /// `MergeScan` decodes under an ordinary multi-input extension parent.
1353    #[tokio::test]
1354    async fn test_serializer_decode_merge_scan_under_union_distinct_on() {
1355        let left = merge_scan(numbers_scan(), false);
1356        let right = merge_scan(numbers_scan(), true);
1357        let plan = LogicalPlan::Extension(Extension {
1358            node: Arc::new(
1359                UnionDistinctOn::try_new(left.clone(), right.clone(), vec![], 0).unwrap(),
1360            ),
1361        });
1362        let bytes = encode_plan(&plan);
1363
1364        let decoded = plan_decoder(
1365            QueryEngineFactory::new(
1366                numbers_catalog_manager(),
1367                None,
1368                None,
1369                None,
1370                None,
1371                false,
1372                QueryOptions::default(),
1373            )
1374            .query_engine()
1375            .as_ref(),
1376        )
1377        .decode(bytes, request_catalog_list(), false)
1378        .await
1379        .unwrap();
1380
1381        let LogicalPlan::Extension(extension) = &decoded else {
1382            panic!("Expect a UnionDistinctOn plan, got: {decoded}");
1383        };
1384        let union = extension
1385            .node
1386            .as_any()
1387            .downcast_ref::<UnionDistinctOn>()
1388            .expect("Expect a UnionDistinctOn plan node");
1389        assert_eq!(union.inputs()[0].to_string(), left.to_string());
1390        assert_eq!(union.inputs()[1].to_string(), right.to_string());
1391    }
1392
1393    /// `MergeScan` inside a subquery still decodes through the wrapper consumer.
1394    ///
1395    /// Correlated outer references cannot be covered here: the in-repo Substrait producer
1396    /// rejects `Expr::OuterReferenceColumn`, so no supported helper can build such a plan.
1397    #[tokio::test]
1398    async fn test_serializer_decode_merge_scan_inside_subquery() {
1399        let inner = merge_scan(numbers_scan(), false);
1400        let payload = LogicalPlanBuilder::from(numbers_scan())
1401            .filter(Expr::InSubquery(InSubquery::new(
1402                Box::new(col("number")),
1403                Subquery {
1404                    subquery: Arc::new(inner.clone()),
1405                    outer_ref_columns: vec![],
1406                    spans: Default::default(),
1407                },
1408                false,
1409            )))
1410            .unwrap()
1411            .build()
1412            .unwrap();
1413        let plan = merge_scan(payload.clone(), false);
1414        let bytes = encode_plan(&plan);
1415
1416        let decoded = plan_decoder(
1417            QueryEngineFactory::new(
1418                numbers_catalog_manager(),
1419                None,
1420                None,
1421                None,
1422                None,
1423                false,
1424                QueryOptions::default(),
1425            )
1426            .query_engine()
1427            .as_ref(),
1428        )
1429        .decode(bytes, request_catalog_list(), false)
1430        .await
1431        .unwrap();
1432
1433        assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2);
1434        assert!(
1435            decoded.to_string().contains("IN (<subquery>)"),
1436            "expect the decoded subquery filter, got: {decoded}"
1437        );
1438    }
1439
1440    /// Recursive payload decoding must restore ordered state aggregates, including their schema.
1441    #[tokio::test]
1442    async fn test_serializer_decode_payload_keeps_state_udaf_ordering() {
1443        let last_value = (*datafusion::functions_aggregate::first_last::last_value_udaf()).clone();
1444        let state_udaf = AggregateUDF::new_from_impl(StateWrapper::new(last_value).unwrap());
1445        let aggr_expr = Expr::AggregateFunction(AggregateFunction::new_udf(
1446            Arc::new(state_udaf),
1447            vec![col("number")],
1448            false,
1449            None,
1450            vec![Sort::new(col("number"), true, true)],
1451            None,
1452        ));
1453        let payload = LogicalPlanBuilder::from(numbers_scan())
1454            .aggregate(Vec::<Expr>::new(), vec![aggr_expr])
1455            .unwrap()
1456            .build()
1457            .unwrap();
1458        let bytes = encode_plan(&payload);
1459        let encoded = Plan::decode(bytes.clone()).unwrap();
1460        let extensions = Extensions::try_from(&encoded.extensions).unwrap();
1461        let catalog_manager = numbers_catalog_manager();
1462        let engine = QueryEngineFactory::new(
1463            catalog_manager.clone(),
1464            None,
1465            None,
1466            None,
1467            None,
1468            false,
1469            QueryOptions::default(),
1470        )
1471        .query_engine();
1472        let mut state = SessionStateBuilder::new_from_existing(
1473            engine.engine_context(QueryContext::arc()).state().clone(),
1474        )
1475        .with_serializer_registry(Arc::new(DefaultSerializer))
1476        .build();
1477        register_greptime_functions(&mut state, &QueryContext::arc()).unwrap();
1478        let consumer = MergeScanSubstraitConsumer {
1479            inner: DefaultSubstraitConsumer::new(&extensions, &state),
1480            session_state: &state,
1481            catalog_manager: Some(catalog_manager),
1482            payload: false,
1483        };
1484        let decoded = consumer.decode_payload(bytes.to_vec()).await.unwrap();
1485
1486        let fixed_payload = FixStateUdafOrderingAnalyzer {}
1487            .analyze(payload.clone(), &Default::default())
1488            .unwrap();
1489        assert_ne!(
1490            fixed_payload.schema().as_arrow(),
1491            payload.schema().as_arrow()
1492        );
1493        assert_eq!(
1494            decoded.schema().as_arrow(),
1495            fixed_payload.schema().as_arrow()
1496        );
1497        let LogicalPlan::Aggregate(aggregate) = &decoded else {
1498            panic!("Expected an aggregate payload, got: {decoded}");
1499        };
1500        let Expr::AggregateFunction(function) = &aggregate.aggr_expr[0] else {
1501            panic!("Expected a state aggregate, got: {:?}", aggregate.aggr_expr);
1502        };
1503        assert!(function.func.inner().is::<StateWrapper>());
1504        assert_eq!(function.params.order_by.len(), 1);
1505        let ArrowDataType::Struct(fields) = decoded.schema().field(0).data_type() else {
1506            panic!("Expected the ordered aggregate state type");
1507        };
1508        assert_eq!(fields.len(), 3);
1509        assert_eq!(fields[1].name(), "number");
1510        assert!(fields[1].is_nullable());
1511    }
1512
1513    /// Type variation extensions keep being rejected while decoding.
1514    #[tokio::test]
1515    async fn test_serializer_decode_rejects_type_variations() {
1516        use substrait::substrait_proto_df::proto::extensions::SimpleExtensionDeclaration;
1517        use substrait::substrait_proto_df::proto::extensions::simple_extension_declaration::{
1518            ExtensionTypeVariation, MappingType,
1519        };
1520
1521        let engine = QueryEngineFactory::new(
1522            numbers_catalog_manager(),
1523            None,
1524            None,
1525            None,
1526            None,
1527            false,
1528            QueryOptions::default(),
1529        )
1530        .query_engine();
1531        let encoded = encode_plan(&merge_scan(numbers_scan(), false));
1532        let mut substrait_plan = Plan::decode(encoded).unwrap();
1533        substrait_plan.extensions.push(SimpleExtensionDeclaration {
1534            mapping_type: Some(MappingType::ExtensionTypeVariation(
1535                ExtensionTypeVariation {
1536                    extension_urn_reference: u32::MAX,
1537                    type_variation_anchor: 1,
1538                    name: "u!variation".to_string(),
1539                },
1540            )),
1541        });
1542        let bytes = Bytes::from(substrait_plan.encode_to_vec());
1543
1544        let err = plan_decoder(engine.as_ref())
1545            .decode(bytes, request_catalog_list(), false)
1546            .await
1547            .unwrap_err();
1548
1549        let err = format!("{err:?}");
1550        assert!(
1551            err.contains("Type variation extensions are not supported"),
1552            "unexpected error: {err}"
1553        );
1554    }
1555}
query/query_engine/default_serializer.rs

query/query_engine/
default_serializer.rs
