substrait/
df_substrait.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 async_trait::async_trait;
18use bytes::{Buf, Bytes, BytesMut};
19use common_function::aggrs::aggr_wrapper::fix_order::{
20    FixStateUdafOrderingAnalyzer, UnFixStateUdafOrderingAnalyzer,
21};
22use datafusion::execution::SessionStateBuilder;
23use datafusion::execution::context::SessionState;
24use datafusion::execution::runtime_env::RuntimeEnv;
25use datafusion::optimizer::AnalyzerRule;
26use datafusion::prelude::SessionConfig;
27use datafusion_expr::LogicalPlan;
28use datafusion_substrait::logical_plan::consumer::{
29    SubstraitConsumer, from_substrait_plan, from_substrait_plan_with_consumer,
30};
31use datafusion_substrait::logical_plan::producer::to_substrait_plan;
32use datafusion_substrait::substrait::proto::Plan;
33use prost::Message;
34use snafu::ResultExt;
35
36use crate::error::{DecodeDfPlanSnafu, DecodeRelSnafu, EncodeDfPlanSnafu, EncodeRelSnafu, Error};
37use crate::{SerializerRegistry, SubstraitPlan};
38
39pub struct DFLogicalSubstraitConvertor;
40
41#[async_trait]
42impl SubstraitPlan for DFLogicalSubstraitConvertor {
43    type Error = Error;
44
45    type Plan = LogicalPlan;
46
47    async fn decode<B: Buf + Send>(
48        &self,
49        message: B,
50        state: SessionState,
51    ) -> Result<Self::Plan, Self::Error> {
52        let plan = Plan::decode(message).context(DecodeRelSnafu)?;
53        let df_plan = from_substrait_plan(&state, &plan)
54            .await
55            .context(DecodeDfPlanSnafu)?;
56        let df_plan = FixStateUdafOrderingAnalyzer {}
57            .analyze(df_plan, state.config_options())
58            .context(DecodeDfPlanSnafu)?;
59        Ok(df_plan)
60    }
61
62    fn encode(
63        &self,
64        plan: &Self::Plan,
65        serializer: impl SerializerRegistry + 'static,
66    ) -> Result<Bytes, Self::Error> {
67        let plan = UnFixStateUdafOrderingAnalyzer {}
68            .analyze(plan.clone(), &Default::default())
69            .context(EncodeDfPlanSnafu)?;
70        let mut buf = BytesMut::new();
71        let substrait_plan = self.to_sub_plan(&plan, serializer)?;
72        substrait_plan.encode(&mut buf).context(EncodeRelSnafu)?;
73
74        Ok(buf.freeze())
75    }
76}
77
78impl DFLogicalSubstraitConvertor {
79    /// Decode with a custom consumer, retaining Greptime's state-aggregate ordering correction.
80    pub async fn decode_with_consumer(
81        &self,
82        plan: &Plan,
83        state: &SessionState,
84        consumer: &impl SubstraitConsumer,
85    ) -> Result<LogicalPlan, Error> {
86        let df_plan = from_substrait_plan_with_consumer(consumer, plan)
87            .await
88            .context(DecodeDfPlanSnafu)?;
89        FixStateUdafOrderingAnalyzer {}
90            .analyze(df_plan, state.config_options())
91            .context(DecodeDfPlanSnafu)
92    }
93
94    pub fn to_sub_plan(
95        &self,
96        plan: &LogicalPlan,
97        serializer: impl SerializerRegistry + 'static,
98    ) -> Result<Box<Plan>, Error> {
99        let state = SessionStateBuilder::new()
100            .with_config(SessionConfig::new())
101            .with_runtime_env(Arc::new(RuntimeEnv::default()))
102            .with_default_features()
103            .with_serializer_registry(Arc::new(serializer))
104            .build();
105        to_substrait_plan(plan, &state).context(EncodeDfPlanSnafu)
106    }
107}
substrait/df_substrait.rs

substrait/
df_substrait.rs
