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            .unwrap()
437            .new_plan_decoder()
438            .unwrap();
439        let catalog_list = Arc::new(DummyCatalogList::with_table_provider(table_provider));
440
441        let decode_plan = plan_decoder
442            .decode(bytes, catalog_list, false)
443            .await
444            .unwrap();
445
446        assert_eq!(
447            "Filter: devices.k0 = Utf8(\"hello\")
448  TableScan: devices",
449            decode_plan.to_string(),
450        );
451    }
452
453    #[tokio::test]
454    async fn test_serializer_decode_native_histogram_udf() {
455        let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap();
456        let factory = QueryEngineFactory::new(
457            catalog_list,
458            None,
459            None,
460            None,
461            None,
462            false,
463            QueryOptions::default(),
464        );
465        let engine = factory.query_engine();
466        let schema = Arc::new(Schema::new(vec![Field::new(
467            "histogram",
468            native_histogram_value_type().as_arrow_type(),
469            true,
470        )]));
471        let plan = LogicalPlanBuilder::scan(
472            "devices",
473            Arc::new(LogicalTableSource::new(schema.clone())),
474            None,
475        )
476        .unwrap()
477        .aggregate(
478            Vec::<Expr>::new(),
479            vec![
480                Arc::new(NativeHistogramAggSum::aggregate_udf())
481                    .call(vec![col("histogram")])
482                    .alias("sum"),
483                Arc::new(NativeHistogramAggAvg::aggregate_udf())
484                    .call(vec![col("histogram")])
485                    .alias("avg"),
486            ],
487        )
488        .unwrap()
489        .project(vec![
490            Expr::ScalarFunction(ScalarFunction {
491                func: Arc::new(NativeHistogramCount::scalar_udf()),
492                args: vec![col("sum")],
493            }),
494            Expr::ScalarFunction(ScalarFunction {
495                func: Arc::new(NativeHistogramDrop::bool_false_udf(
496                    "ignored annotation".to_string(),
497                    None,
498                )),
499                args: vec![col("sum")],
500            }),
501            Expr::ScalarFunction(ScalarFunction {
502                func: Arc::new(NativeHistogramDrop::bool_true_udf(
503                    "ignored annotation".to_string(),
504                    None,
505                )),
506                args: vec![col("sum")],
507            }),
508            Expr::ScalarFunction(ScalarFunction {
509                func: Arc::new(NativeHistogramDrop::float_null_udf(
510                    "ignored annotation".to_string(),
511                    None,
512                )),
513                args: vec![col("sum")],
514            }),
515            Expr::ScalarFunction(ScalarFunction {
516                func: Arc::new(PromqlFloatToString::scalar_udf()),
517                args: vec![lit(2.0)],
518            }),
519        ])
520        .unwrap()
521        .build()
522        .unwrap();
523        let bytes = DFLogicalSubstraitConvertor
524            .encode(&plan, DefaultSerializer)
525            .unwrap();
526        let table_provider = Arc::new(MemTable::try_new(schema, vec![vec![]]).unwrap());
527        let plan_decoder = engine
528            .engine_context(QueryContext::arc())
529            .unwrap()
530            .new_plan_decoder()
531            .unwrap();
532
533        let decoded = plan_decoder
534            .decode(
535                bytes,
536                Arc::new(DummyCatalogList::with_table_provider(table_provider)),
537                false,
538            )
539            .await
540            .unwrap();
541
542        let decoded = decoded.to_string();
543        assert!(decoded.contains("prom_native_histogram_count"));
544        assert!(decoded.contains("prom_native_histogram_drop_bool"));
545        assert!(decoded.contains("prom_native_histogram_keep_bool"));
546        assert!(decoded.contains("prom_native_histogram_drop_float"));
547        assert!(decoded.contains(NativeHistogramAggSum::name()));
548        assert!(decoded.contains(NativeHistogramAggAvg::name()));
549        assert!(decoded.contains(PromqlFloatToString::name()));
550
551        let schema = Arc::new(Schema::new(vec![
552            Field::new(
553                "timestamp",
554                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
555                false,
556            ),
557            Field::new("float", ArrowDataType::Float64, true),
558            Field::new(
559                "histogram",
560                native_histogram_value_type().as_arrow_type(),
561                true,
562            ),
563        ]));
564        let input = LogicalPlanBuilder::scan(
565            "devices",
566            Arc::new(LogicalTableSource::new(schema.clone())),
567            None,
568        )
569        .unwrap()
570        .build()
571        .unwrap();
572        let input = LogicalPlan::Extension(Extension {
573            node: Arc::new(
574                RangeManipulate::new(
575                    0,
576                    1000,
577                    1000,
578                    0,
579                    1000,
580                    "timestamp".to_string(),
581                    vec!["float".to_string(), "histogram".to_string()],
582                    input,
583                )
584                .unwrap(),
585            ),
586        });
587        let plan = LogicalPlanBuilder::from(input)
588            .project(vec![
589                Expr::ScalarFunction(ScalarFunction {
590                    func: Arc::new(MixedRange::float_udf(None)),
591                    args: vec![
592                        lit("last_over_time"),
593                        col("timestamp_range"),
594                        col("float"),
595                        col("histogram"),
596                    ],
597                })
598                .alias("mixed_float"),
599                Expr::ScalarFunction(ScalarFunction {
600                    func: Arc::new(MixedRange::histogram_udf(None)),
601                    args: vec![
602                        lit("last_over_time"),
603                        col("timestamp_range"),
604                        col("float"),
605                        col("histogram"),
606                    ],
607                })
608                .alias("mixed_histogram"),
609            ])
610            .unwrap()
611            .build()
612            .unwrap();
613        let bytes = DFLogicalSubstraitConvertor
614            .encode(&plan, DefaultSerializer)
615            .unwrap();
616        let table_provider = Arc::new(MemTable::try_new(schema, vec![vec![]]).unwrap());
617        let decoded = plan_decoder
618            .decode(
619                bytes,
620                Arc::new(DummyCatalogList::with_table_provider(table_provider)),
621                false,
622            )
623            .await
624            .unwrap()
625            .to_string();
626        assert!(decoded.contains("prom_mixed_range_float"));
627        assert!(decoded.contains("prom_mixed_range_histogram"));
628    }
629
630    #[tokio::test]
631    async fn test_serializer_decode_merge_scan() {
632        let engine = QueryEngineFactory::new(
633            numbers_catalog_manager(),
634            None,
635            None,
636            None,
637            None,
638            false,
639            QueryOptions::default(),
640        )
641        .query_engine();
642        let plan = merge_scan(numbers_scan(), true);
643        let decoded = plan_decoder(engine.as_ref())
644            .decode(encode_plan(&plan), request_catalog_list(), false)
645            .await
646            .unwrap();
647
648        let LogicalPlan::Extension(extension) = &decoded else {
649            panic!("Expect a MergeScan plan, got: {decoded}");
650        };
651        let merge_scan = extension
652            .node
653            .as_any()
654            .downcast_ref::<MergeScanLogicalPlan>()
655            .expect("Expect a MergeScan plan node");
656        assert!(merge_scan.is_placeholder());
657        assert_eq!(merge_scan.input().to_string(), numbers_scan().to_string());
658    }
659
660    /// Payload table binding uses the engine catalog, not the request catalog.
661    #[tokio::test]
662    async fn test_serializer_decode_merge_scan_with_engine_catalog() {
663        let engine = QueryEngineFactory::new(
664            numbers_catalog_manager(),
665            None,
666            None,
667            None,
668            None,
669            false,
670            QueryOptions::default(),
671        )
672        .query_engine();
673        let decoded = plan_decoder(engine.as_ref())
674            .decode(
675                encode_plan(&merge_scan(numbers_scan(), false)),
676                request_catalog_list(),
677                false,
678            )
679            .await
680            .unwrap();
681
682        let LogicalPlan::Extension(extension) = &decoded else {
683            panic!("Expect a MergeScan plan, got: {decoded}");
684        };
685        let merge_scan = extension
686            .node
687            .as_any()
688            .downcast_ref::<MergeScanLogicalPlan>()
689            .expect("Expect a MergeScan plan node");
690        let LogicalPlan::TableScan(scan) = merge_scan.input() else {
691            panic!("Expect a table scan, got: {}", merge_scan.input());
692        };
693        assert_eq!(scan.table_name.table(), NUMBERS_TABLE_NAME);
694        assert_eq!(scan.source.schema().fields().len(), 1);
695        assert_eq!(scan.source.schema().field(0).name(), "number");
696    }
697
698    /// Nested `MergeScan` payloads decode recursively.
699    #[tokio::test]
700    async fn test_serializer_decode_nested_merge_scan() {
701        let engine = QueryEngineFactory::new(
702            numbers_catalog_manager(),
703            None,
704            None,
705            None,
706            None,
707            false,
708            QueryOptions::default(),
709        )
710        .query_engine();
711        let inner = merge_scan(numbers_scan(), false);
712        let input = LogicalPlanBuilder::from(inner)
713            .filter(col("number").lt(lit(10u32)))
714            .unwrap()
715            .build()
716            .unwrap();
717        let plan = merge_scan(input, false);
718
719        let decoded = plan_decoder(engine.as_ref())
720            .decode(encode_plan(&plan), request_catalog_list(), false)
721            .await
722            .unwrap();
723
724        assert_eq!(decoded.to_string(), plan.to_string());
725        assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2);
726    }
727
728    fn greptime_date_format_udf(query_ctx: &QueryContextRef) -> Arc<ScalarUDF> {
729        FUNCTION_REGISTRY
730            .get_function("date_format")
731            .expect("`date_format` is a function of GreptimeDB")
732            .provide(FunctionContext {
733                query_ctx: query_ctx.clone(),
734                state: Default::default(),
735            })
736            .into()
737    }
738
739    /// Uses a non-default timezone to distinguish the Greptime implementation.
740    fn date_format_query_ctx() -> QueryContextRef {
741        let query_ctx = QueryContext::with(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME);
742        query_ctx.set_timezone(Timezone::from_tz_string("Asia/Shanghai").unwrap());
743        Arc::new(query_ctx)
744    }
745
746    fn greptime_date_format_expr(query_ctx: &QueryContextRef, timestamp: Expr) -> Expr {
747        Expr::ScalarFunction(ScalarFunction {
748            func: greptime_date_format_udf(query_ctx),
749            args: vec![timestamp, lit("%Y-%m-%d %H:%i:%S")],
750        })
751    }
752
753    /// Includes `MergeScan` payload expressions, which are not node inputs.
754    fn scalar_functions_of(plan: &LogicalPlan) -> Vec<Arc<ScalarUDF>> {
755        use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
756
757        let mut udfs = Vec::new();
758        for expr in plan.expressions() {
759            expr.apply(|expr| {
760                if let Expr::ScalarFunction(func) = expr {
761                    udfs.push(Arc::clone(&func.func));
762                }
763                Ok(TreeNodeRecursion::Continue)
764            })
765            .unwrap();
766        }
767
768        if let LogicalPlan::Extension(extension) = plan
769            && let Some(merge_scan) = extension
770                .node
771                .as_any()
772                .downcast_ref::<MergeScanLogicalPlan>()
773        {
774            udfs.extend(scalar_functions_of(merge_scan.input()));
775        }
776
777        for input in plan.inputs() {
778            udfs.extend(scalar_functions_of(input));
779        }
780
781        udfs
782    }
783
784    /// Asserts the Greptime `date_format` binding rather than DataFusion's `to_char` alias.
785    fn assert_only_function_is_greptime_date_format(
786        plan: &LogicalPlan,
787        query_ctx: &QueryContextRef,
788    ) {
789        let udfs = scalar_functions_of(plan);
790        assert_eq!(udfs.len(), 1, "expected one function in: {plan}");
791
792        let udf = &udfs[0];
793        assert_eq!(udf.name(), "date_format", "got: {plan}");
794        assert_eq!(
795            udf.as_ref(),
796            greptime_date_format_udf(query_ctx).as_ref(),
797            "`date_format` is not bound to the implementation of GreptimeDB: {plan}"
798        );
799    }
800
801    /// Request state registration keeps Greptime `date_format` ahead of DataFusion aliases.
802    #[tokio::test]
803    async fn test_serializer_decode_keeps_greptime_date_format() {
804        // Repeat to expose nondeterministic function-registration order.
805        for _ in 0..8 {
806            let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap();
807            let factory = QueryEngineFactory::new(
808                catalog_list,
809                None,
810                None,
811                None,
812                None,
813                false,
814                QueryOptions::default(),
815            );
816            let engine = factory.query_engine();
817
818            let table_provider = Arc::new(mock_table_provider(1.into()));
819            let query_ctx = date_format_query_ctx();
820            let plan = LogicalPlanBuilder::scan(
821                "devices",
822                Arc::new(LogicalTableSource::new(table_provider.schema().clone())),
823                None,
824            )
825            .unwrap()
826            .project(vec![greptime_date_format_expr(&query_ctx, col("ts"))])
827            .unwrap()
828            .build()
829            .unwrap();
830
831            let bytes = DFLogicalSubstraitConvertor
832                .encode(&plan, DefaultSerializer)
833                .unwrap();
834            let plan_decoder = engine
835                .engine_context(query_ctx.clone())
836                .unwrap()
837                .new_plan_decoder()
838                .unwrap();
839            let decoded = plan_decoder
840                .decode(
841                    bytes,
842                    Arc::new(DummyCatalogList::with_table_provider(table_provider)),
843                    false,
844                )
845                .await
846                .unwrap();
847
848            assert_only_function_is_greptime_date_format(&decoded, &query_ctx);
849        }
850    }
851
852    /// Nested payload decoding retains the Greptime `date_format` binding.
853    #[tokio::test]
854    async fn test_serializer_decode_merge_scan_payload_keeps_greptime_date_format() {
855        let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap();
856        catalog_manager
857            .register_table_sync(RegisterTableRequest {
858                catalog: DEFAULT_CATALOG_NAME.to_string(),
859                schema: DEFAULT_SCHEMA_NAME.to_string(),
860                table_name: NUMBERS_TABLE_NAME.to_string(),
861                table_id: NUMBERS_TABLE_ID,
862                table: NumbersTable::table(NUMBERS_TABLE_ID),
863            })
864            .unwrap();
865
866        for _ in 0..8 {
867            let factory = QueryEngineFactory::new(
868                catalog_manager.clone(),
869                None,
870                None,
871                None,
872                None,
873                false,
874                QueryOptions::default(),
875            );
876            let engine = factory.query_engine();
877            let query_ctx = date_format_query_ctx();
878
879            let input = LogicalPlanBuilder::scan(
880                NUMBERS_TABLE_NAME,
881                Arc::new(LogicalTableSource::new(
882                    NumbersTable::schema().arrow_schema().clone(),
883                )),
884                None,
885            )
886            .unwrap()
887            .project(vec![greptime_date_format_expr(
888                &query_ctx,
889                Expr::Cast(Cast::new(
890                    Box::new(col("number")),
891                    ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
892                )),
893            )])
894            .unwrap()
895            .build()
896            .unwrap();
897            let plan = MergeScanLogicalPlan::new(input, false).into_logical_plan();
898
899            let bytes = DFLogicalSubstraitConvertor
900                .encode(&plan, DefaultSerializer)
901                .unwrap();
902            let plan_decoder = engine
903                .engine_context(query_ctx.clone())
904                .unwrap()
905                .new_plan_decoder()
906                .unwrap();
907            let catalog_list = Arc::new(DummyCatalogList::with_table_provider(Arc::new(
908                mock_table_provider(1.into()),
909            )));
910
911            let decoded = plan_decoder
912                .decode(bytes, catalog_list, false)
913                .await
914                .unwrap();
915
916            let LogicalPlan::Extension(extension) = &decoded else {
917                panic!("Expect a MergeScan plan, got: {decoded}");
918            };
919            let merge_scan = extension
920                .node
921                .as_any()
922                .downcast_ref::<MergeScanLogicalPlan>()
923                .expect("Expect a MergeScan plan node");
924
925            assert_only_function_is_greptime_date_format(merge_scan.input(), &query_ctx);
926        }
927    }
928
929    /// Registers a `numbers` table in a memory catalog manager.
930    fn numbers_catalog_manager() -> CatalogManagerRef {
931        let catalog_manager = catalog::memory::new_memory_catalog_manager().unwrap();
932        catalog_manager
933            .register_table_sync(RegisterTableRequest {
934                catalog: DEFAULT_CATALOG_NAME.to_string(),
935                schema: DEFAULT_SCHEMA_NAME.to_string(),
936                table_name: NUMBERS_TABLE_NAME.to_string(),
937                table_id: NUMBERS_TABLE_ID,
938                table: NumbersTable::table(NUMBERS_TABLE_ID),
939            })
940            .unwrap();
941        catalog_manager
942    }
943
944    /// Scan plan over the `numbers` schema shared by the fixtures below.
945    fn numbers_scan() -> LogicalPlan {
946        LogicalPlanBuilder::scan(
947            NUMBERS_TABLE_NAME,
948            Arc::new(LogicalTableSource::new(
949                NumbersTable::schema().arrow_schema().clone(),
950            )),
951            None,
952        )
953        .unwrap()
954        .build()
955        .unwrap()
956    }
957
958    fn merge_scan(input: LogicalPlan, is_placeholder: bool) -> LogicalPlan {
959        MergeScanLogicalPlan::new(input, is_placeholder).into_logical_plan()
960    }
961
962    fn encode_plan(plan: &LogicalPlan) -> Bytes {
963        DFLogicalSubstraitConvertor
964            .encode(plan, DefaultSerializer)
965            .unwrap()
966    }
967
968    fn plan_decoder(engine: &dyn crate::QueryEngine) -> SubstraitPlanDecoderRef {
969        engine
970            .engine_context(QueryContext::arc())
971            .unwrap()
972            .new_plan_decoder()
973            .unwrap()
974    }
975
976    /// Request catalog that resolves any table to a provider payload resolution must not use.
977    fn request_catalog_list() -> Arc<dyn CatalogProviderList> {
978        Arc::new(DummyCatalogList::with_table_provider(Arc::new(
979            mock_table_provider(1.into()),
980        )))
981    }
982
983    /// Payload resolution preserves context and explicit catalog/schema components.
984    #[tokio::test]
985    async fn test_serializer_decode_payload_observes_catalog_and_query_context() {
986        let observed = Arc::new(std::sync::Mutex::new(Vec::new()));
987        let inner = catalog::memory::new_memory_catalog_manager().unwrap();
988        let manager: CatalogManagerRef = Arc::new(InterceptingCatalogManager {
989            inner,
990            resolution: TableResolution::ObserveContext(observed.clone()),
991        });
992        let mut resolver_ctx = QueryContext::arc();
993        Arc::make_mut(&mut resolver_ctx).set_current_catalog("request_cat");
994        Arc::make_mut(&mut resolver_ctx).set_current_schema("request_schema");
995        let engine = QueryEngineFactory::new(
996            manager,
997            None,
998            None,
999            None,
1000            None,
1001            false,
1002            QueryOptions::default(),
1003        )
1004        .query_engine();
1005        let mut state = engine
1006            .engine_context(resolver_ctx.clone())
1007            .unwrap()
1008            .state()
1009            .clone();
1010        state.config_mut().set_extension(resolver_ctx);
1011        let extensions = Extensions::default();
1012        let consumer = MergeScanSubstraitConsumer {
1013            inner: DefaultSubstraitConsumer::new(&extensions, &state),
1014            session_state: &state,
1015            catalog_manager: state
1016                .config()
1017                .get_extension::<QueryEngineState>()
1018                .map(|s| s.catalog_manager().clone()),
1019            payload: true,
1020        };
1021        for reference in [
1022            TableReference::bare(NUMBERS_TABLE_NAME),
1023            TableReference::partial("explicit_schema", NUMBERS_TABLE_NAME),
1024            TableReference::full(
1025                DEFAULT_CATALOG_NAME,
1026                DEFAULT_SCHEMA_NAME,
1027                NUMBERS_TABLE_NAME,
1028            ),
1029        ] {
1030            consumer
1031                .resolve_table_ref(&reference)
1032                .await
1033                .unwrap()
1034                .unwrap();
1035        }
1036        assert_eq!(
1037            *observed.lock().unwrap(),
1038            vec![
1039                [
1040                    "request_cat".to_string(),
1041                    "request_schema".to_string(),
1042                    NUMBERS_TABLE_NAME.to_string(),
1043                    "request_cat".to_string(),
1044                    "request_schema".to_string(),
1045                ],
1046                [
1047                    "request_cat".to_string(),
1048                    "explicit_schema".to_string(),
1049                    NUMBERS_TABLE_NAME.to_string(),
1050                    "request_cat".to_string(),
1051                    "request_schema".to_string(),
1052                ],
1053                [
1054                    DEFAULT_CATALOG_NAME.to_string(),
1055                    DEFAULT_SCHEMA_NAME.to_string(),
1056                    NUMBERS_TABLE_NAME.to_string(),
1057                    "request_cat".to_string(),
1058                    "request_schema".to_string(),
1059                ],
1060            ],
1061        );
1062    }
1063
1064    /// How [`InterceptingCatalogManager`] resolves payload tables.
1065    enum TableResolution {
1066        /// Records catalog lookup and request context values.
1067        ObserveContext(Arc<std::sync::Mutex<Vec<[String; 5]>>>),
1068        /// Resolves once another task on the caller runtime signals.
1069        AwaitSignal {
1070            notify: Arc<tokio::sync::Notify>,
1071            entered: Arc<AtomicBool>,
1072        },
1073    }
1074
1075    /// Engine catalog manager that takes over `table` resolution while delegating the rest.
1076    struct InterceptingCatalogManager {
1077        inner: CatalogManagerRef,
1078        resolution: TableResolution,
1079    }
1080
1081    #[async_trait::async_trait]
1082    impl CatalogManager for InterceptingCatalogManager {
1083        fn as_any(&self) -> &dyn Any {
1084            self.inner.as_any()
1085        }
1086
1087        async fn catalog_names(&self) -> CatalogResult<Vec<String>> {
1088            self.inner.catalog_names().await
1089        }
1090
1091        async fn schema_names(
1092            &self,
1093            catalog: &str,
1094            query_ctx: Option<&QueryContext>,
1095        ) -> CatalogResult<Vec<String>> {
1096            self.inner.schema_names(catalog, query_ctx).await
1097        }
1098
1099        async fn table_names(
1100            &self,
1101            catalog: &str,
1102            schema: &str,
1103            query_ctx: Option<&QueryContext>,
1104        ) -> CatalogResult<Vec<String>> {
1105            self.inner.table_names(catalog, schema, query_ctx).await
1106        }
1107
1108        async fn catalog_exists(&self, catalog: &str) -> CatalogResult<bool> {
1109            self.inner.catalog_exists(catalog).await
1110        }
1111
1112        async fn schema_exists(
1113            &self,
1114            catalog: &str,
1115            schema: &str,
1116            query_ctx: Option<&QueryContext>,
1117        ) -> CatalogResult<bool> {
1118            self.inner.schema_exists(catalog, schema, query_ctx).await
1119        }
1120
1121        async fn table_exists(
1122            &self,
1123            catalog: &str,
1124            schema: &str,
1125            table: &str,
1126            query_ctx: Option<&QueryContext>,
1127        ) -> CatalogResult<bool> {
1128            self.inner
1129                .table_exists(catalog, schema, table, query_ctx)
1130                .await
1131        }
1132
1133        async fn table(
1134            &self,
1135            catalog: &str,
1136            schema: &str,
1137            table_name: &str,
1138            query_ctx: Option<&QueryContext>,
1139        ) -> CatalogResult<Option<TableRef>> {
1140            match &self.resolution {
1141                TableResolution::ObserveContext(observed) => {
1142                    let query_ctx = query_ctx.unwrap();
1143                    observed.lock().unwrap().push([
1144                        catalog.to_string(),
1145                        schema.to_string(),
1146                        table_name.to_string(),
1147                        query_ctx.current_catalog().to_string(),
1148                        query_ctx.current_schema(),
1149                    ]);
1150                    Ok(Some(NumbersTable::table(NUMBERS_TABLE_ID)))
1151                }
1152
1153                TableResolution::AwaitSignal { notify, entered } => {
1154                    entered.store(true, Ordering::SeqCst);
1155                    notify.notified().await;
1156                    self.inner
1157                        .table(catalog, schema, table_name, query_ctx)
1158                        .await
1159                }
1160            }
1161        }
1162
1163        async fn table_info_by_id(&self, table_id: TableId) -> CatalogResult<Option<TableInfoRef>> {
1164            self.inner.table_info_by_id(table_id).await
1165        }
1166
1167        async fn tables_by_ids(
1168            &self,
1169            catalog: &str,
1170            schema: &str,
1171            table_ids: &[TableId],
1172        ) -> CatalogResult<Vec<TableRef>> {
1173            self.inner.tables_by_ids(catalog, schema, table_ids).await
1174        }
1175
1176        fn tables<'a>(
1177            &'a self,
1178            catalog: &'a str,
1179            schema: &'a str,
1180            query_ctx: Option<&'a QueryContext>,
1181        ) -> BoxStream<'a, CatalogResult<TableRef>> {
1182            self.inner.tables(catalog, schema, query_ctx)
1183        }
1184    }
1185
1186    fn engine_with_intercepting_catalog(
1187        resolution: TableResolution,
1188    ) -> Arc<dyn crate::QueryEngine> {
1189        let catalog_manager: CatalogManagerRef = Arc::new(InterceptingCatalogManager {
1190            inner: numbers_catalog_manager(),
1191            resolution,
1192        });
1193        QueryEngineFactory::new(
1194            catalog_manager,
1195            None,
1196            None,
1197            None,
1198            None,
1199            false,
1200            QueryOptions::default(),
1201        )
1202        .query_engine()
1203    }
1204
1205    /// A payload lookup parked on a caller-runtime task must keep the decode async.
1206    ///
1207    /// The engine catalog resolves only once a task spawned on this runtime signals it, so a
1208    /// blocking bridge would deadlock the test instead of decoding.
1209    #[tokio::test(flavor = "current_thread")]
1210    async fn test_serializer_decode_merge_scan_awaits_caller_runtime_task() {
1211        let notify = Arc::new(tokio::sync::Notify::new());
1212        let entered = Arc::new(AtomicBool::new(false));
1213        let engine = engine_with_intercepting_catalog(TableResolution::AwaitSignal {
1214            notify: notify.clone(),
1215            entered: entered.clone(),
1216        });
1217        let plan = merge_scan(numbers_scan(), false);
1218        let bytes = encode_plan(&plan);
1219
1220        let decoder = plan_decoder(engine.as_ref());
1221        let mut decode = Box::pin(decoder.decode(bytes, request_catalog_list(), false));
1222        // First poll reaches the payload table lookup, which waits for the signal below.
1223        assert!(futures::poll!(decode.as_mut()).is_pending());
1224        assert!(entered.load(Ordering::SeqCst));
1225
1226        let signal = tokio::spawn(async move {
1227            notify.notify_waiters();
1228        });
1229        tokio::task::yield_now().await;
1230        let decoded = decode.await.unwrap();
1231        signal.await.unwrap();
1232
1233        assert_eq!(decoded.to_string(), plan.to_string());
1234    }
1235
1236    /// A missing engine-catalog table does not fall back to the usable request catalog.
1237    #[tokio::test]
1238    async fn test_serializer_decode_missing_engine_catalog_table() {
1239        let missing_manager: CatalogManagerRef =
1240            catalog::memory::new_memory_catalog_manager().unwrap();
1241        let engine = QueryEngineFactory::new(
1242            missing_manager,
1243            None,
1244            None,
1245            None,
1246            None,
1247            false,
1248            QueryOptions::default(),
1249        )
1250        .query_engine();
1251        let bytes = encode_plan(&merge_scan(numbers_scan(), false));
1252
1253        let err = plan_decoder(engine.as_ref())
1254            .decode(bytes, request_catalog_list(), false)
1255            .await
1256            .unwrap_err();
1257
1258        let err = format!("{err:?}");
1259        assert!(
1260            err.contains("Table not found: greptime.public.numbers"),
1261            "unexpected error: {err}"
1262        );
1263    }
1264
1265    /// Manual states use their own catalog and query-context defaults for payloads.
1266    #[tokio::test]
1267    async fn test_serializer_decode_manual_state_keeps_extension_registry() {
1268        let manager = catalog::memory::new_memory_catalog_manager().unwrap();
1269        manager.register_catalog_sync("manual_cat").unwrap();
1270        manager
1271            .register_schema_sync(catalog::RegisterSchemaRequest {
1272                catalog: "manual_cat".to_string(),
1273                schema: "manual_schema".to_string(),
1274            })
1275            .unwrap();
1276        manager
1277            .register_table_sync(RegisterTableRequest {
1278                catalog: "manual_cat".to_string(),
1279                schema: "manual_schema".to_string(),
1280                table_name: NUMBERS_TABLE_NAME.to_string(),
1281                table_id: NUMBERS_TABLE_ID,
1282                table: NumbersTable::table(NUMBERS_TABLE_ID),
1283            })
1284            .unwrap();
1285        let query_ctx = Arc::new(QueryContext::with("manual_cat", "manual_schema"));
1286        let mut config = datafusion::execution::context::SessionConfig::new()
1287            .with_information_schema(true)
1288            .with_default_catalog_and_schema("manual_cat", "manual_schema");
1289        config.set_extension(query_ctx.clone());
1290        let state = SessionStateBuilder::new()
1291            .with_default_features()
1292            .with_config(config)
1293            .build();
1294        let decoder = DefaultPlanDecoder::new(state, &query_ctx).unwrap();
1295        let plan = LogicalPlan::Extension(Extension {
1296            node: Arc::new(SeriesNormalize::new(
1297                0,
1298                "number",
1299                false,
1300                Vec::new(),
1301                merge_scan(numbers_scan(), false),
1302            )),
1303        });
1304        let request_catalog = Arc::new(
1305            catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx(
1306                manager, query_ctx,
1307            ),
1308        );
1309        let decoded = decoder
1310            .decode(encode_plan(&plan), request_catalog, false)
1311            .await
1312            .unwrap();
1313
1314        assert_eq!(decoded.to_string(), plan.to_string());
1315        assert_eq!(decoded.schema().as_arrow(), plan.schema().as_arrow());
1316    }
1317
1318    /// `MergeScan` decodes under an ordinary single-input extension parent.
1319    #[tokio::test]
1320    async fn test_serializer_decode_merge_scan_under_series_normalize() {
1321        let outer = merge_scan(numbers_scan(), false);
1322        let plan = LogicalPlan::Extension(Extension {
1323            node: Arc::new(SeriesNormalize::new(
1324                0,
1325                "number",
1326                false,
1327                Vec::new(),
1328                outer.clone(),
1329            )),
1330        });
1331        let bytes = encode_plan(&plan);
1332
1333        let decoded = plan_decoder(
1334            QueryEngineFactory::new(
1335                numbers_catalog_manager(),
1336                None,
1337                None,
1338                None,
1339                None,
1340                false,
1341                QueryOptions::default(),
1342            )
1343            .query_engine()
1344            .as_ref(),
1345        )
1346        .decode(bytes, request_catalog_list(), false)
1347        .await
1348        .unwrap();
1349
1350        let LogicalPlan::Extension(extension) = &decoded else {
1351            panic!("Expect a SeriesNormalize plan, got: {decoded}");
1352        };
1353        let normalize = extension
1354            .node
1355            .as_any()
1356            .downcast_ref::<SeriesNormalize>()
1357            .expect("Expect a SeriesNormalize plan node");
1358        assert_eq!(normalize.inputs()[0].to_string(), outer.to_string());
1359    }
1360
1361    /// `MergeScan` decodes under an ordinary multi-input extension parent.
1362    #[tokio::test]
1363    async fn test_serializer_decode_merge_scan_under_union_distinct_on() {
1364        let left = merge_scan(numbers_scan(), false);
1365        let right = merge_scan(numbers_scan(), true);
1366        let plan = LogicalPlan::Extension(Extension {
1367            node: Arc::new(
1368                UnionDistinctOn::try_new(left.clone(), right.clone(), vec![], 0).unwrap(),
1369            ),
1370        });
1371        let bytes = encode_plan(&plan);
1372
1373        let decoded = plan_decoder(
1374            QueryEngineFactory::new(
1375                numbers_catalog_manager(),
1376                None,
1377                None,
1378                None,
1379                None,
1380                false,
1381                QueryOptions::default(),
1382            )
1383            .query_engine()
1384            .as_ref(),
1385        )
1386        .decode(bytes, request_catalog_list(), false)
1387        .await
1388        .unwrap();
1389
1390        let LogicalPlan::Extension(extension) = &decoded else {
1391            panic!("Expect a UnionDistinctOn plan, got: {decoded}");
1392        };
1393        let union = extension
1394            .node
1395            .as_any()
1396            .downcast_ref::<UnionDistinctOn>()
1397            .expect("Expect a UnionDistinctOn plan node");
1398        assert_eq!(union.inputs()[0].to_string(), left.to_string());
1399        assert_eq!(union.inputs()[1].to_string(), right.to_string());
1400    }
1401
1402    /// `MergeScan` inside a subquery still decodes through the wrapper consumer.
1403    ///
1404    /// Correlated outer references cannot be covered here: the in-repo Substrait producer
1405    /// rejects `Expr::OuterReferenceColumn`, so no supported helper can build such a plan.
1406    #[tokio::test]
1407    async fn test_serializer_decode_merge_scan_inside_subquery() {
1408        let inner = merge_scan(numbers_scan(), false);
1409        let payload = LogicalPlanBuilder::from(numbers_scan())
1410            .filter(Expr::InSubquery(InSubquery::new(
1411                Box::new(col("number")),
1412                Subquery {
1413                    subquery: Arc::new(inner.clone()),
1414                    outer_ref_columns: vec![],
1415                    spans: Default::default(),
1416                },
1417                false,
1418            )))
1419            .unwrap()
1420            .build()
1421            .unwrap();
1422        let plan = merge_scan(payload.clone(), false);
1423        let bytes = encode_plan(&plan);
1424
1425        let decoded = plan_decoder(
1426            QueryEngineFactory::new(
1427                numbers_catalog_manager(),
1428                None,
1429                None,
1430                None,
1431                None,
1432                false,
1433                QueryOptions::default(),
1434            )
1435            .query_engine()
1436            .as_ref(),
1437        )
1438        .decode(bytes, request_catalog_list(), false)
1439        .await
1440        .unwrap();
1441
1442        assert_eq!(decoded.to_string().matches("MergeScan [").count(), 2);
1443        assert!(
1444            decoded.to_string().contains("IN (<subquery>)"),
1445            "expect the decoded subquery filter, got: {decoded}"
1446        );
1447    }
1448
1449    /// Recursive payload decoding must restore ordered state aggregates, including their schema.
1450    #[tokio::test]
1451    async fn test_serializer_decode_payload_keeps_state_udaf_ordering() {
1452        let last_value = (*datafusion::functions_aggregate::first_last::last_value_udaf()).clone();
1453        let state_udaf = AggregateUDF::new_from_impl(StateWrapper::new(last_value).unwrap());
1454        let aggr_expr = Expr::AggregateFunction(AggregateFunction::new_udf(
1455            Arc::new(state_udaf),
1456            vec![col("number")],
1457            false,
1458            None,
1459            vec![Sort::new(col("number"), true, true)],
1460            None,
1461        ));
1462        let payload = LogicalPlanBuilder::from(numbers_scan())
1463            .aggregate(Vec::<Expr>::new(), vec![aggr_expr])
1464            .unwrap()
1465            .build()
1466            .unwrap();
1467        let bytes = encode_plan(&payload);
1468        let encoded = Plan::decode(bytes.clone()).unwrap();
1469        let extensions = Extensions::try_from(&encoded.extensions).unwrap();
1470        let catalog_manager = numbers_catalog_manager();
1471        let engine = QueryEngineFactory::new(
1472            catalog_manager.clone(),
1473            None,
1474            None,
1475            None,
1476            None,
1477            false,
1478            QueryOptions::default(),
1479        )
1480        .query_engine();
1481        let mut state = SessionStateBuilder::new_from_existing(
1482            engine
1483                .engine_context(QueryContext::arc())
1484                .unwrap()
1485                .state()
1486                .clone(),
1487        )
1488        .with_serializer_registry(Arc::new(DefaultSerializer))
1489        .build();
1490        register_greptime_functions(&mut state, &QueryContext::arc()).unwrap();
1491        let consumer = MergeScanSubstraitConsumer {
1492            inner: DefaultSubstraitConsumer::new(&extensions, &state),
1493            session_state: &state,
1494            catalog_manager: Some(catalog_manager),
1495            payload: false,
1496        };
1497        let decoded = consumer.decode_payload(bytes.to_vec()).await.unwrap();
1498
1499        let fixed_payload = FixStateUdafOrderingAnalyzer {}
1500            .analyze(payload.clone(), &Default::default())
1501            .unwrap();
1502        assert_ne!(
1503            fixed_payload.schema().as_arrow(),
1504            payload.schema().as_arrow()
1505        );
1506        assert_eq!(
1507            decoded.schema().as_arrow(),
1508            fixed_payload.schema().as_arrow()
1509        );
1510        let LogicalPlan::Aggregate(aggregate) = &decoded else {
1511            panic!("Expected an aggregate payload, got: {decoded}");
1512        };
1513        let Expr::AggregateFunction(function) = &aggregate.aggr_expr[0] else {
1514            panic!("Expected a state aggregate, got: {:?}", aggregate.aggr_expr);
1515        };
1516        assert!(function.func.inner().is::<StateWrapper>());
1517        assert_eq!(function.params.order_by.len(), 1);
1518        let ArrowDataType::Struct(fields) = decoded.schema().field(0).data_type() else {
1519            panic!("Expected the ordered aggregate state type");
1520        };
1521        assert_eq!(fields.len(), 3);
1522        assert_eq!(fields[1].name(), "number");
1523        assert!(fields[1].is_nullable());
1524    }
1525
1526    /// Type variation extensions keep being rejected while decoding.
1527    #[tokio::test]
1528    async fn test_serializer_decode_rejects_type_variations() {
1529        use substrait::substrait_proto_df::proto::extensions::SimpleExtensionDeclaration;
1530        use substrait::substrait_proto_df::proto::extensions::simple_extension_declaration::{
1531            ExtensionTypeVariation, MappingType,
1532        };
1533
1534        let engine = QueryEngineFactory::new(
1535            numbers_catalog_manager(),
1536            None,
1537            None,
1538            None,
1539            None,
1540            false,
1541            QueryOptions::default(),
1542        )
1543        .query_engine();
1544        let encoded = encode_plan(&merge_scan(numbers_scan(), false));
1545        let mut substrait_plan = Plan::decode(encoded).unwrap();
1546        substrait_plan.extensions.push(SimpleExtensionDeclaration {
1547            mapping_type: Some(MappingType::ExtensionTypeVariation(
1548                ExtensionTypeVariation {
1549                    extension_urn_reference: u32::MAX,
1550                    type_variation_anchor: 1,
1551                    name: "u!variation".to_string(),
1552                },
1553            )),
1554        });
1555        let bytes = Bytes::from(substrait_plan.encode_to_vec());
1556
1557        let err = plan_decoder(engine.as_ref())
1558            .decode(bytes, request_catalog_list(), false)
1559            .await
1560            .unwrap_err();
1561
1562        let err = format!("{err:?}");
1563        assert!(
1564            err.contains("Type variation extensions are not supported"),
1565            "unexpected error: {err}"
1566        );
1567    }
1568}
query/query_engine/default_serializer.rs

query/query_engine/
default_serializer.rs
