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 common_error::ext::BoxedError;
18use common_function::function::FunctionContext;
19use common_function::function_registry::FUNCTION_REGISTRY;
20use common_query::error::RegisterUdfSnafu;
21use common_query::logical_plan::SubstraitPlanDecoder;
22use datafusion::catalog::CatalogProviderList;
23use datafusion::common::DataFusionError;
24use datafusion::error::Result;
25use datafusion::execution::context::SessionState;
26use datafusion::execution::registry::SerializerRegistry;
27use datafusion::execution::{FunctionRegistry, SessionStateBuilder};
28use datafusion::logical_expr::LogicalPlan;
29use datafusion_expr::UserDefinedLogicalNode;
30use greptime_proto::substrait_extension::MergeScan as PbMergeScan;
31use promql::functions::{
32    AbsentOverTime, AvgOverTime, Changes, CountOverTime, Delta, Deriv, DoubleExponentialSmoothing,
33    IDelta, Increase, LastOverTime, MaxOverTime, MinOverTime, NativeHistogramAbsentOverTime,
34    NativeHistogramAdd, NativeHistogramAggAvg, NativeHistogramAggSum, NativeHistogramAvg,
35    NativeHistogramAvgOverTime, NativeHistogramChanges, NativeHistogramCount,
36    NativeHistogramCountOverTime, NativeHistogramDelta, NativeHistogramDivScalar,
37    NativeHistogramDrop, NativeHistogramEq, NativeHistogramFraction, NativeHistogramIDelta,
38    NativeHistogramIRate, NativeHistogramIncrease, NativeHistogramLastOverTime,
39    NativeHistogramMulScalar, NativeHistogramNeg, NativeHistogramNotEq,
40    NativeHistogramPresentOverTime, NativeHistogramQuantile, NativeHistogramRate,
41    NativeHistogramResets, NativeHistogramScalarMul, NativeHistogramStddev, NativeHistogramStdvar,
42    NativeHistogramSub, NativeHistogramSum, NativeHistogramSumOverTime, NativeHistogramToString,
43    PredictLinear, PresentOverTime, QuantileOverTime, Rate, Resets, Round, StddevOverTime,
44    StdvarOverTime, SumOverTime, quantile_udaf,
45};
46use prost::Message;
47use session::context::QueryContextRef;
48use snafu::ResultExt;
49use substrait::extension_serializer::ExtensionSerializer;
50use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};
51
52use crate::dist_plan::MergeScanLogicalPlan;
53
54/// Extended [`substrait::extension_serializer::ExtensionSerializer`] but supports [`MergeScanLogicalPlan`] serialization.
55#[derive(Debug)]
56pub struct DefaultSerializer;
57
58impl SerializerRegistry for DefaultSerializer {
59    fn serialize_logical_plan(&self, node: &dyn UserDefinedLogicalNode) -> Result<Vec<u8>> {
60        if node.name() == MergeScanLogicalPlan::name() {
61            let merge_scan = node
62                .as_any()
63                .downcast_ref::<MergeScanLogicalPlan>()
64                .expect("Failed to downcast to MergeScanLogicalPlan");
65
66            let input = merge_scan.input();
67            let is_placeholder = merge_scan.is_placeholder();
68            let input = DFLogicalSubstraitConvertor
69                .encode(input, DefaultSerializer)
70                .map_err(|e| DataFusionError::External(Box::new(e)))?
71                .to_vec();
72
73            Ok(PbMergeScan {
74                is_placeholder,
75                input,
76            }
77            .encode_to_vec())
78        } else {
79            ExtensionSerializer.serialize_logical_plan(node)
80        }
81    }
82
83    fn deserialize_logical_plan(
84        &self,
85        name: &str,
86        bytes: &[u8],
87    ) -> Result<Arc<dyn UserDefinedLogicalNode>> {
88        if name == MergeScanLogicalPlan::name() {
89            // TODO(dennis): missing `session_state` to decode the logical plan in `MergeScanLogicalPlan`,
90            // so we only save the unoptimized logical plan for view currently.
91            Err(DataFusionError::Substrait(format!(
92                "Unsupported plan node: {name}"
93            )))
94        } else {
95            ExtensionSerializer.deserialize_logical_plan(name, bytes)
96        }
97    }
98}
99
100/// The datafusion `[LogicalPlan]` decoder.
101pub struct DefaultPlanDecoder {
102    session_state: SessionState,
103    query_ctx: QueryContextRef,
104}
105
106impl DefaultPlanDecoder {
107    pub fn new(
108        session_state: SessionState,
109        query_ctx: &QueryContextRef,
110    ) -> crate::error::Result<Self> {
111        Ok(Self {
112            session_state,
113            query_ctx: query_ctx.clone(),
114        })
115    }
116}
117
118#[async_trait::async_trait]
119impl SubstraitPlanDecoder for DefaultPlanDecoder {
120    async fn decode(
121        &self,
122        message: bytes::Bytes,
123        catalog_list: Arc<dyn CatalogProviderList>,
124        optimize: bool,
125    ) -> common_query::error::Result<LogicalPlan> {
126        // The session_state already has the `DefaultSerialzier` as `SerializerRegistry`.
127        let mut session_state = SessionStateBuilder::new_from_existing(self.session_state.clone())
128            .with_catalog_list(catalog_list)
129            .build();
130        // Substrait decoder will look up the UDFs in SessionState, so we need to register them
131        // Note: the query context must be passed to set the timezone
132        // We MUST register the UDFs after we build the session state, otherwise the UDFs will be lost
133        // if they have the same name as the default UDFs or their alias.
134        // e.g. The default UDF `to_char()` has an alias `date_format()`, if we register a UDF with the name `date_format()`
135        // before we build the session state, the UDF will be lost.
136        for func in FUNCTION_REGISTRY.scalar_functions() {
137            let udf = func.provide(FunctionContext {
138                query_ctx: self.query_ctx.clone(),
139                state: Default::default(),
140            });
141            session_state
142                .register_udf(Arc::new(udf))
143                .context(RegisterUdfSnafu { name: func.name() })?;
144        }
145
146        for func in FUNCTION_REGISTRY.aggregate_functions() {
147            let name = func.name().to_string();
148            session_state
149                .register_udaf(Arc::new(func))
150                .context(RegisterUdfSnafu { name })?;
151        }
152
153        let _ = session_state.register_udaf(quantile_udaf());
154
155        let _ = session_state.register_udf(Arc::new(IDelta::<false>::scalar_udf()));
156        let _ = session_state.register_udf(Arc::new(IDelta::<true>::scalar_udf()));
157        let _ = session_state.register_udf(Arc::new(Rate::scalar_udf()));
158        let _ = session_state.register_udf(Arc::new(Increase::scalar_udf()));
159        let _ = session_state.register_udf(Arc::new(Delta::scalar_udf()));
160        let _ = session_state.register_udf(Arc::new(Resets::scalar_udf()));
161        let _ = session_state.register_udf(Arc::new(Changes::scalar_udf()));
162        let _ = session_state.register_udf(Arc::new(Deriv::scalar_udf()));
163        let _ = session_state.register_udf(Arc::new(Round::scalar_udf()));
164        let _ = session_state.register_udf(Arc::new(AvgOverTime::scalar_udf()));
165        let _ = session_state.register_udf(Arc::new(MinOverTime::scalar_udf()));
166        let _ = session_state.register_udf(Arc::new(MaxOverTime::scalar_udf()));
167        let _ = session_state.register_udf(Arc::new(SumOverTime::scalar_udf()));
168        let _ = session_state.register_udf(Arc::new(CountOverTime::scalar_udf()));
169        let _ = session_state.register_udf(Arc::new(LastOverTime::scalar_udf()));
170        let _ = session_state.register_udf(Arc::new(AbsentOverTime::scalar_udf()));
171        let _ = session_state.register_udf(Arc::new(PresentOverTime::scalar_udf()));
172        let _ = session_state.register_udf(Arc::new(StddevOverTime::scalar_udf()));
173        let _ = session_state.register_udf(Arc::new(StdvarOverTime::scalar_udf()));
174        let _ = session_state.register_udf(Arc::new(QuantileOverTime::scalar_udf()));
175        let _ = session_state.register_udf(Arc::new(PredictLinear::scalar_udf()));
176        let double_exponential_smoothing_udf =
177            DoubleExponentialSmoothing::scalar_udf().with_aliases(["prom_holt_winters"]);
178        let _ = session_state.register_udf(Arc::new(double_exponential_smoothing_udf));
179
180        for udf in [
181            NativeHistogramAbsentOverTime::scalar_udf(),
182            NativeHistogramAdd::scalar_udf(),
183            NativeHistogramAvg::scalar_udf(),
184            NativeHistogramAvgOverTime::scalar_udf(),
185            NativeHistogramChanges::scalar_udf(),
186            NativeHistogramCount::scalar_udf(),
187            NativeHistogramCountOverTime::scalar_udf(),
188            NativeHistogramDelta::scalar_udf(),
189            NativeHistogramDivScalar::scalar_udf(),
190            NativeHistogramDrop::bool_false_udf(String::new(), None),
191            NativeHistogramDrop::float_null_udf(String::new(), None),
192            NativeHistogramEq::scalar_udf(),
193            NativeHistogramFraction::scalar_udf(),
194            NativeHistogramIDelta::scalar_udf(),
195            NativeHistogramIRate::scalar_udf(),
196            NativeHistogramIncrease::scalar_udf(),
197            NativeHistogramLastOverTime::scalar_udf(),
198            NativeHistogramMulScalar::scalar_udf(),
199            NativeHistogramNeg::scalar_udf(),
200            NativeHistogramNotEq::scalar_udf(),
201            NativeHistogramPresentOverTime::scalar_udf(),
202            NativeHistogramQuantile::scalar_udf(),
203            NativeHistogramRate::scalar_udf(),
204            NativeHistogramResets::scalar_udf(),
205            NativeHistogramScalarMul::scalar_udf(),
206            NativeHistogramStddev::scalar_udf(),
207            NativeHistogramStdvar::scalar_udf(),
208            NativeHistogramSub::scalar_udf(),
209            NativeHistogramSum::scalar_udf(),
210            NativeHistogramSumOverTime::scalar_udf(),
211            NativeHistogramToString::scalar_udf(),
212        ] {
213            let _ = session_state.register_udf(Arc::new(udf));
214        }
215        for udaf in [
216            NativeHistogramAggAvg::aggregate_udf(),
217            NativeHistogramAggSum::aggregate_udf(),
218        ] {
219            let _ = session_state.register_udaf(Arc::new(udaf));
220        }
221
222        let logical_plan = DFLogicalSubstraitConvertor
223            .decode(message, session_state)
224            .await
225            .map_err(BoxedError::new)
226            .context(common_query::error::DecodePlanSnafu)?;
227
228        if optimize {
229            self.session_state
230                .optimize(&logical_plan)
231                .map_err(Into::into)
232        } else {
233            Ok(logical_plan)
234        }
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use common_query::native_histogram::native_histogram_value_type;
241    use datafusion::catalog::TableProvider;
242    use datafusion::datasource::MemTable;
243    use datafusion_expr::expr::ScalarFunction;
244    use datafusion_expr::{Expr, LogicalPlanBuilder, LogicalTableSource, col, lit};
245    use datatypes::arrow::datatypes::{Field, Schema, SchemaRef};
246    use datatypes::data_type::DataType;
247    use session::context::QueryContext;
248
249    use super::*;
250    use crate::QueryEngineFactory;
251    use crate::dummy_catalog::DummyCatalogList;
252    use crate::optimizer::test_util::mock_table_provider;
253    use crate::options::QueryOptions;
254
255    fn mock_plan(schema: SchemaRef) -> LogicalPlan {
256        let table_source = LogicalTableSource::new(schema);
257        let projection = None;
258        let builder =
259            LogicalPlanBuilder::scan("devices", Arc::new(table_source), projection).unwrap();
260
261        builder
262            .filter(col("k0").eq(lit("hello")))
263            .unwrap()
264            .build()
265            .unwrap()
266    }
267
268    #[tokio::test]
269    async fn test_serializer_decode_plan() {
270        let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap();
271        let factory = QueryEngineFactory::new(
272            catalog_list,
273            None,
274            None,
275            None,
276            None,
277            false,
278            QueryOptions::default(),
279        );
280
281        let engine = factory.query_engine();
282
283        let table_provider = Arc::new(mock_table_provider(1.into()));
284        let plan = mock_plan(table_provider.schema().clone());
285
286        let bytes = DFLogicalSubstraitConvertor
287            .encode(&plan, DefaultSerializer)
288            .unwrap();
289
290        let plan_decoder = engine
291            .engine_context(QueryContext::arc())
292            .new_plan_decoder()
293            .unwrap();
294        let catalog_list = Arc::new(DummyCatalogList::with_table_provider(table_provider));
295
296        let decode_plan = plan_decoder
297            .decode(bytes, catalog_list, false)
298            .await
299            .unwrap();
300
301        assert_eq!(
302            "Filter: devices.k0 = Utf8(\"hello\")
303  TableScan: devices",
304            decode_plan.to_string(),
305        );
306    }
307
308    #[tokio::test]
309    async fn test_serializer_decode_native_histogram_udf() {
310        let catalog_list = catalog::memory::new_memory_catalog_manager().unwrap();
311        let factory = QueryEngineFactory::new(
312            catalog_list,
313            None,
314            None,
315            None,
316            None,
317            false,
318            QueryOptions::default(),
319        );
320        let engine = factory.query_engine();
321        let schema = Arc::new(Schema::new(vec![Field::new(
322            "histogram",
323            native_histogram_value_type().as_arrow_type(),
324            true,
325        )]));
326        let plan = LogicalPlanBuilder::scan(
327            "devices",
328            Arc::new(LogicalTableSource::new(schema.clone())),
329            None,
330        )
331        .unwrap()
332        .project(vec![
333            Expr::ScalarFunction(ScalarFunction {
334                func: Arc::new(NativeHistogramCount::scalar_udf()),
335                args: vec![col("histogram")],
336            }),
337            Expr::ScalarFunction(ScalarFunction {
338                func: Arc::new(NativeHistogramDrop::bool_false_udf(
339                    "ignored annotation".to_string(),
340                    None,
341                )),
342                args: vec![col("histogram")],
343            }),
344            Expr::ScalarFunction(ScalarFunction {
345                func: Arc::new(NativeHistogramDrop::float_null_udf(
346                    "ignored annotation".to_string(),
347                    None,
348                )),
349                args: vec![col("histogram")],
350            }),
351        ])
352        .unwrap()
353        .build()
354        .unwrap();
355        let bytes = DFLogicalSubstraitConvertor
356            .encode(&plan, DefaultSerializer)
357            .unwrap();
358        let table_provider = Arc::new(MemTable::try_new(schema, vec![vec![]]).unwrap());
359        let plan_decoder = engine
360            .engine_context(QueryContext::arc())
361            .new_plan_decoder()
362            .unwrap();
363
364        let decoded = plan_decoder
365            .decode(
366                bytes,
367                Arc::new(DummyCatalogList::with_table_provider(table_provider)),
368                false,
369            )
370            .await
371            .unwrap();
372
373        let decoded = decoded.to_string();
374        assert!(decoded.contains("prom_native_histogram_count"));
375        assert!(decoded.contains("prom_native_histogram_drop_bool"));
376        assert!(decoded.contains("prom_native_histogram_drop_float"));
377    }
378}
query/query_engine/default_serializer.rs

query/query_engine/
default_serializer.rs
