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

query/query_engine/
default_serializer.rs
