flow/batching_mode/task/
inc.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_telemetry::debug;
19use common_telemetry::tracing::warn;
20use datafusion_expr::{DmlStatement, LogicalPlan};
21use query::options::{
22    FLOW_INCREMENTAL_AFTER_SEQS, FLOW_INCREMENTAL_MODE, FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY,
23    FLOW_SINK_TABLE_ID,
24};
25use snafu::ResultExt;
26use table::metadata::TableId;
27
28use crate::Error;
29use crate::batching_mode::state::CheckpointMode;
30use crate::batching_mode::table_creator::QueryType;
31use crate::batching_mode::task::BatchingTask;
32use crate::batching_mode::utils::{
33    analyze_incremental_aggregate_plan, get_table_info_df_schema,
34    rewrite_incremental_aggregate_with_sink_merge,
35};
36use crate::error::{ExternalSnafu, UnexpectedSnafu};
37
38impl BatchingTask {
39    async fn sink_table_id(&self) -> Result<TableId, Error> {
40        let table = self
41            .config
42            .catalog_manager
43            .table(
44                &self.config.sink_table_name[0],
45                &self.config.sink_table_name[1],
46                &self.config.sink_table_name[2],
47                None,
48            )
49            .await
50            .map_err(BoxedError::new)
51            .context(ExternalSnafu)?
52            .ok_or_else(|| {
53                UnexpectedSnafu {
54                    reason: format!(
55                        "Flow {} cannot build incremental extensions because sink table {:?} was not found",
56                        self.config.flow_id, self.config.sink_table_name
57                    ),
58                }
59                .build()
60            })?;
61        Ok(table.table_info().table_id())
62    }
63
64    /// For incremental-mode SQL queries, attempt to prepare an executable plan
65    /// that is safe for incremental scan extensions.
66    ///
67    /// Returns `Some(plan)` when incremental extensions are safe, and `None`
68    /// when the caller should execute the original plan without incremental
69    /// extensions. The returned plan may be either a rewritten
70    /// delta-LEFT-JOIN-sink merge plan or the original plan. In particular,
71    /// plain GROUP BY queries with no aggregate merge columns are incremental
72    /// safe without a rewrite, so they return `Some(original_plan)`.
73    pub(super) async fn prepare_plan_for_incremental(
74        &self,
75        plan: &LogicalPlan,
76    ) -> Result<Option<LogicalPlan>, Error> {
77        let is_incremental_sql = {
78            let state = self.state.read().unwrap();
79            if state.is_incremental_disabled() {
80                return Ok(None);
81            }
82            state.checkpoint_mode() == CheckpointMode::Incremental
83                && matches!(self.config.query_type, QueryType::Sql)
84        };
85
86        if !is_incremental_sql {
87            return Ok(None);
88        }
89
90        // Extract inner query plan from the DML wrapper.
91        // Non-DML or non-SQL plans bypass the rewrite and keep checkpoint mode;
92        // non-aggregate TQL or non-INSERT plans do not need incremental scan extensions.
93        let inner_plan = match plan {
94            LogicalPlan::Dml(dml) => dml.input.as_ref().clone(),
95            _ => return Ok(None),
96        };
97
98        // Analyze the plan for incremental rewritability.
99        // Incremental reads currently require aggregate / group-by plans that
100        // can be rewritten into a delta-left-join-sink merge. Non-aggregate SQL
101        // (projection, filter, or other non-aggregate shapes) stays full-snapshot
102        // until separately supported, and incremental mode is permanently
103        // disabled for this flow.
104        let Some(analysis) = analyze_incremental_aggregate_plan(&inner_plan)? else {
105            warn!(
106                "Flow {} incremental mode but plan is not an aggregate query; \
107                 permanently disabling incremental for this flow",
108                self.config.flow_id
109            );
110            self.state.write().unwrap().disable_incremental();
111            return Ok(None);
112        };
113
114        if !analysis.unsupported_exprs.is_empty() {
115            warn!(
116                "Flow {} incremental aggregate contains unsupported expressions {:?}; \
117                 permanently disabling incremental for this flow",
118                self.config.flow_id, analysis.unsupported_exprs
119            );
120            self.state.write().unwrap().disable_incremental();
121            return Ok(None);
122        }
123
124        // Plain GROUP BY without aggregate expressions has no values to
125        // merge between delta and sink. The incremental delta scan emits
126        // changed groups, and sink primary-key write semantics make this
127        // idempotent; no explicit left-join rewrite is needed.
128        if analysis.merge_columns.is_empty() {
129            return Ok(Some(plan.clone()));
130        }
131
132        // Fetch sink table for the merge rewrite.
133        // Transient errors (catalog, schema, filter, or rewrite) should not
134        // permanently disable incremental mode. Instead, we fall back to a
135        // full-snapshot plan for this round while keeping incremental retryable.
136        let sink_table = match get_table_info_df_schema(
137            self.config.catalog_manager.clone(),
138            self.config.sink_table_name.clone(),
139        )
140        .await
141        {
142            Ok((table, _)) => table,
143            Err(err) => {
144                warn!(
145                    "Flow {} failed to fetch sink table for incremental rewrite; \
146                     falling back to full snapshot for this round: {:?}",
147                    self.config.flow_id, err
148                );
149                self.state.write().unwrap().mark_full_snapshot();
150                return Ok(None);
151            }
152        };
153        let rewritten_inner = match rewrite_incremental_aggregate_with_sink_merge(
154            &inner_plan,
155            &analysis,
156            sink_table,
157            &self.config.sink_table_name,
158            None,
159        )
160        .await
161        {
162            Ok(plan) => plan,
163            Err(err) => {
164                warn!(
165                    "Flow {} failed to rewrite incremental aggregate with sink merge; \
166                     falling back to full snapshot for this round: {:?}",
167                    self.config.flow_id, err
168                );
169                self.state.write().unwrap().mark_full_snapshot();
170                return Ok(None);
171            }
172        };
173
174        // Reconstruct DML plan with the rewritten inner plan
175        let rewritten = match plan {
176            LogicalPlan::Dml(dml) => LogicalPlan::Dml(DmlStatement::new(
177                dml.table_name.clone(),
178                dml.target.clone(),
179                dml.op.clone(),
180                Arc::new(rewritten_inner),
181            )),
182            _ => unreachable!("already matched Dml above"),
183        };
184
185        debug!(
186            "Flow {} rewrote incremental SQL aggregate query with sink merge",
187            self.config.flow_id
188        );
189
190        Ok(Some(rewritten))
191    }
192
193    pub(super) async fn build_flow_query_extensions(
194        &self,
195        incremental_safe: bool,
196        can_advance_checkpoints: bool,
197    ) -> Result<Vec<(&'static str, String)>, Error> {
198        let mut extensions = vec![("flow.return_region_seq", "true".to_string())];
199
200        let incremental_checkpoints_json = {
201            let state = self.state.read().unwrap();
202            if incremental_safe
203                && can_advance_checkpoints
204                && !state.is_incremental_disabled()
205                && state.checkpoint_mode() == CheckpointMode::Incremental
206                && !state.checkpoints().is_empty()
207            {
208                Some(serde_json::to_string(state.checkpoints()).map_err(|err| {
209                    UnexpectedSnafu {
210                        reason: format!("Failed to serialize checkpoint map: {err}"),
211                    }
212                    .build()
213                })?)
214            } else {
215                None
216            }
217        };
218
219        if let Some(checkpoints_json) = incremental_checkpoints_json {
220            let sink_table_id = self.sink_table_id().await?;
221            extensions.push((FLOW_SINK_TABLE_ID, sink_table_id.to_string()));
222            extensions.push((
223                FLOW_INCREMENTAL_MODE,
224                FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY.to_string(),
225            ));
226            extensions.push((FLOW_INCREMENTAL_AFTER_SEQS, checkpoints_json));
227        }
228
229        Ok(extensions)
230    }
231}
flow/batching_mode/task/inc.rs

flow/batching_mode/task/
inc.rs
