flow/adapter/
stateless.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
15//! Stateless DataFusion execution for streaming flows.
16
17use std::collections::HashSet;
18use std::sync::Arc;
19
20use api::helper::{to_grpc_value, vectors_to_rows};
21use api::v1::greptime_request::Request;
22use api::v1::{RowInsertRequest, RowInsertRequests, Rows};
23use common_error::ext::BoxedError;
24use common_query::OutputData;
25use common_recordbatch::{RecordBatch, RecordBatches, map_dictionary_to_values_data_type};
26use common_time::Timestamp;
27use datafusion::catalog::MemTable;
28use datafusion::datasource::{TableProvider, provider_as_source, source_as_provider};
29use datafusion_common::tree_node::{Transformed, TreeNode};
30use datafusion_common::{Column, DFSchema, TableReference};
31use datafusion_expr::logical_plan::{Distinct, Projection};
32use datafusion_expr::{Expr, LogicalPlan};
33use datatypes::schema::{ColumnSchema, SchemaRef};
34use datatypes::value::Value;
35use query::QueryEngine;
36use session::context::QueryContextRef;
37use snafu::{OptionExt, ResultExt, ensure};
38use table::metadata::TableId;
39use table::table::adapter::DfTableProviderAdapter;
40
41use crate::TableName;
42use crate::adapter::util::column_schemas_to_proto;
43use crate::batching_mode::frontend_client::FrontendClient;
44use crate::error::{DatafusionSnafu, Error, ExternalSnafu, InvalidQuerySnafu, UnexpectedSnafu};
45use crate::repr::DiffRow;
46
47/// The validated, immutable part of one streaming flow.
48#[derive(Clone)]
49pub(crate) struct StatelessFlow {
50    pub(crate) source_table_id: TableId,
51    pub(crate) source_table_name: TableName,
52    pub(crate) source_schema: SchemaRef,
53    pub(crate) source_schema_version: u32,
54    pub(crate) sink_table_name: TableName,
55    pub(crate) sink_schema: Vec<ColumnSchema>,
56    pub(crate) sink_primary_keys: Vec<String>,
57    /// The exact trailing columns resolved when the flow was created.
58    pub(crate) auto_columns: Vec<ColumnSchema>,
59    pub(crate) plan: LogicalPlan,
60    pub(crate) query_ctx: QueryContextRef,
61    pub(crate) create_args: crate::CreateFlowArgs,
62}
63
64/// Per-request input provider. It owns no catalog or storage state.
65#[cfg(test)]
66fn test_source_plan(table_name: TableReference, provider: Arc<dyn TableProvider>) -> LogicalPlan {
67    datafusion_expr::LogicalPlanBuilder::scan(table_name, provider_as_source(provider), None)
68        .unwrap()
69        .filter(datafusion_expr::col("number").gt(datafusion_expr::lit(1)))
70        .unwrap()
71        .project(vec![datafusion_expr::col("number")])
72        .unwrap()
73        .build()
74        .unwrap()
75}
76
77fn input_provider(batch: &RecordBatch) -> Result<Arc<dyn TableProvider>, Error> {
78    let arrow_batch = batch.df_record_batch().clone();
79    let provider = MemTable::try_new(arrow_batch.schema(), vec![vec![arrow_batch]]).context(
80        DatafusionSnafu {
81            context: "Failed to create transient flow input provider",
82        },
83    )?;
84    Ok(Arc::new(provider))
85}
86
87/// Adds the source timestamp to every supported plan node that has to carry it through a
88/// filter or projection. The expression is appended only after the visible expressions, so the
89/// sink contract remains positional.
90pub(crate) fn rewrite_source_timestamp(
91    plan: LogicalPlan,
92    source_name: &TableReference,
93    source_timestamp_name: &str,
94) -> Result<LogicalPlan, Error> {
95    let mut names = HashSet::new();
96    plan.apply(|node| {
97        names.extend(
98            node.schema()
99                .fields()
100                .iter()
101                .map(|field| field.name().clone()),
102        );
103        Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue)
104    })
105    .context(DatafusionSnafu {
106        context: "Failed to inspect streaming flow plan schema",
107    })?;
108    let mut hidden_name = "__flow_source_timestamp".to_string();
109    let mut suffix = 0;
110    while !names.insert(hidden_name.clone()) {
111        suffix += 1;
112        hidden_name = format!("__flow_source_timestamp_{suffix}");
113    }
114    let visible_count = plan.schema().fields().len();
115    let source_timestamp = Column::from_name(source_timestamp_name);
116    let mut plan = plan
117        .transform_up_with_subqueries(|node| match node {
118            LogicalPlan::TableScan(mut scan) => {
119                if scan.table_name.resolved_eq(source_name)
120                    && let Some(projection) = &mut scan.projection
121                {
122                    let timestamp_index = scan
123                        .source
124                        .schema()
125                        .index_of(source_timestamp_name)
126                        .map_err(|_| {
127                            datafusion::error::DataFusionError::Plan(
128                                "Source timestamp is absent from source scan".into(),
129                            )
130                        })?;
131                    if !projection.contains(&timestamp_index) {
132                        projection.push(timestamp_index);
133                        let schema = scan.source.schema();
134                        scan.projected_schema = Arc::new(DFSchema::new_with_metadata(
135                            projection
136                                .iter()
137                                .map(|index| {
138                                    (
139                                        Some(scan.table_name.clone()),
140                                        Arc::new(schema.field(*index).clone()),
141                                    )
142                                })
143                                .collect(),
144                            schema.metadata().clone(),
145                        )?);
146                    }
147                }
148                Ok(Transformed::yes(LogicalPlan::TableScan(scan)))
149            }
150            LogicalPlan::Projection(mut projection) => {
151                let hidden_expr = if projection
152                    .input
153                    .schema()
154                    .fields()
155                    .iter()
156                    .any(|field| field.name() == &hidden_name)
157                {
158                    Expr::Column(Column::from_name(hidden_name.clone()))
159                } else {
160                    Expr::Column(source_timestamp.clone())
161                };
162                projection.expr.push(hidden_expr.alias(hidden_name.clone()));
163                let projection = Projection::try_new(projection.expr, projection.input)?;
164                Ok(Transformed::yes(LogicalPlan::Projection(projection)))
165            }
166            _ => Ok(Transformed::no(node)),
167        })
168        .context(DatafusionSnafu {
169            context: "Failed to add source timestamp to streaming flow plan",
170        })?
171        .data;
172
173    // A plan ending at a scan or filter has no projection at which to give the carried value its
174    // hidden name. Add one only in that case; a projection below a filter already carries it.
175    if !plan
176        .schema()
177        .fields()
178        .iter()
179        .any(|field| field.name() == &hidden_name)
180    {
181        let expressions = plan
182            .schema()
183            .fields()
184            .iter()
185            .take(visible_count)
186            .map(|field| Expr::Column(Column::from_name(field.name())))
187            .chain(std::iter::once(
188                Expr::Column(source_timestamp).alias(hidden_name),
189            ))
190            .collect::<Vec<_>>();
191        plan = Projection::try_new(expressions, Arc::new(plan))
192            .map(LogicalPlan::Projection)
193            .context(DatafusionSnafu {
194                context: "Failed to finalize source timestamp in streaming flow plan",
195            })?;
196    }
197    Ok(plan)
198}
199
200fn replace_source(
201    plan: LogicalPlan,
202    source_name: &TableReference,
203    provider: Arc<dyn TableProvider>,
204) -> Result<LogicalPlan, Error> {
205    let mut scan_count = 0;
206    let mut replaced_scan_count = 0;
207    let plan = plan
208        .transform_up_with_subqueries(|node| match node {
209            LogicalPlan::TableScan(mut scan) => {
210                scan_count += 1;
211                if scan.table_name.resolved_eq(source_name) {
212                    replaced_scan_count += 1;
213                    scan.source = provider_as_source(provider.clone());
214                    let schema = scan.source.schema();
215                    scan.projected_schema = if let Some(projection) = &scan.projection {
216                        Arc::new(DFSchema::new_with_metadata(
217                            projection
218                                .iter()
219                                .map(|index| {
220                                    (
221                                        Some(scan.table_name.clone()),
222                                        Arc::new(schema.field(*index).clone()),
223                                    )
224                                })
225                                .collect(),
226                            schema.metadata().clone(),
227                        )?)
228                    } else {
229                        Arc::new(DFSchema::try_from_qualified_schema(
230                            scan.table_name.clone(),
231                            &schema,
232                        )?)
233                    };
234                    Ok(Transformed::yes(LogicalPlan::TableScan(scan)))
235                } else {
236                    Ok(Transformed::no(LogicalPlan::TableScan(scan)))
237                }
238            }
239            LogicalPlan::Join(_) | LogicalPlan::Aggregate(_) => {
240                Err(datafusion::error::DataFusionError::Plan(
241                    "Streaming flow supports only a single projection/filter source".into(),
242                ))
243            }
244            _ => Ok(Transformed::no(node)),
245        })
246        .context(DatafusionSnafu {
247            context: "Failed to substitute transient flow input provider",
248        })?
249        .data;
250    ensure!(
251        scan_count == 1 && replaced_scan_count == 1,
252        InvalidQuerySnafu {
253            reason: format!(
254                "Expected one source scan matching {:?}, found {scan_count} scans and {replaced_scan_count} matches",
255                source_name
256            )
257        }
258    );
259    Ok(plan)
260}
261
262/// Validates the deliberately small stateless streaming SQL subset.
263pub(crate) fn validate_plan(plan: &LogicalPlan) -> Result<(), Error> {
264    let mut scans = 0;
265    plan.apply(|node| {
266        match node {
267            LogicalPlan::TableScan(_) => scans += 1,
268            LogicalPlan::Projection(_) | LogicalPlan::Filter(_) => {}
269            LogicalPlan::Aggregate(_) => {
270                return Err(datafusion::error::DataFusionError::Plan(
271                    "Aggregation is unsupported in streaming flows. Recreate the flow to select batching mode. A source table with TTL=instant must use persisted retention first. Aggregation SQL without a time window requires EVAL INTERVAL.".into(),
272                ));
273            }
274            // DISTINCT is evaluated against this request's transient input only. DISTINCT ON
275            // has ordering/selection semantics beyond the supported stateless subset.
276            LogicalPlan::Distinct(Distinct::All(_)) => {}
277            _ => {
278                return Err(datafusion::error::DataFusionError::Plan(
279                    "Streaming flow supports only projection and filter over one source scan"
280                        .into(),
281                ));
282            }
283        }
284        Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue)
285    })
286    .context(DatafusionSnafu {
287        context: "Failed to validate streaming flow plan",
288    })?;
289    ensure!(
290        scans == 1,
291        InvalidQuerySnafu {
292            reason: format!("Expected one source scan, found {scans}")
293        }
294    );
295    Ok(())
296}
297
298/// Ensures the retained scan was planned against the source metadata captured for this flow.
299pub(crate) fn validate_source_scan(
300    plan: &LogicalPlan,
301    source_table_id: TableId,
302    source_schema: &SchemaRef,
303) -> Result<(), Error> {
304    plan.apply(|node| {
305        if let LogicalPlan::TableScan(scan) = node {
306            let provider = source_as_provider(&scan.source)?;
307            let provider = provider
308                .downcast_ref::<DfTableProviderAdapter>()
309                .ok_or_else(|| {
310                    datafusion::error::DataFusionError::Plan(
311                        "Streaming flow source scan does not use a table provider".into(),
312                    )
313                })?;
314            let table_info = provider.table().table_info();
315            if table_info.ident.table_id != source_table_id
316                || table_info.meta.schema.as_ref() != source_schema.as_ref()
317            {
318                return Err(datafusion::error::DataFusionError::Plan(format!(
319                    "Streaming flow source scan does not match source table {source_table_id} schema version {}",
320                    source_schema.version()
321                )));
322            }
323        }
324        Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue)
325    })
326    .context(DatafusionSnafu {
327        context: "Failed to validate streaming flow source scan",
328    })?;
329    Ok(())
330}
331
332/// Rejects execution when the source metadata changed after the flow plan was retained.
333/// The outer adapter replans a flow when it observes a new source schema version; this guard
334/// rejects a request if the version changes again before execution.
335fn validate_source_schema_version(
336    retained_version: u32,
337    current_version: u32,
338) -> Result<(), Error> {
339    ensure!(
340        retained_version == current_version,
341        InvalidQuerySnafu {
342            reason: format!(
343                "Source schema version changed from {retained_version} to {current_version}; recreate or recover the flow before writing"
344            )
345        }
346    );
347    Ok(())
348}
349
350fn synthesize_auto_values(columns: &[ColumnSchema], now: Timestamp) -> Result<Vec<Value>, Error> {
351    columns
352        .iter()
353        .map(|column| {
354            let timestamp_type = column.data_type.as_timestamp().context(InvalidQuerySnafu {
355                reason: format!("Auto sink column {} is not a timestamp", column.name),
356            })?;
357            let value = if column.name == crate::adapter::AUTO_CREATED_UPDATE_AT_TS_COL {
358                now.convert_to(timestamp_type.unit())
359                    .context(InvalidQuerySnafu {
360                        reason: "Current timestamp cannot be represented in sink timestamp unit",
361                    })?
362            } else if column.name == crate::adapter::AUTO_CREATED_PLACEHOLDER_TS_COL {
363                Timestamp::new(0, timestamp_type.unit())
364            } else {
365                return InvalidQuerySnafu {
366                    reason: format!("Unsupported auto sink column {}", column.name),
367                }
368                .fail();
369            };
370            Ok(Value::Timestamp(value))
371        })
372        .collect()
373}
374
375/// Executes one mirror write using only the supplied batch and writes its output.
376pub(crate) async fn execute(
377    flow: &StatelessFlow,
378    rows: &[DiffRow],
379    batch_datatypes: &[datatypes::data_type::ConcreteDataType],
380    query_engine: &Arc<dyn QueryEngine>,
381    frontend_client: &Arc<FrontendClient>,
382    current_source_schema_version: u32,
383) -> Result<usize, Error> {
384    validate_source_schema_version(flow.source_schema_version, current_source_schema_version)?;
385    let values = rows.iter().map(|(row, _, _)| row.clone()).collect();
386    let batch = crate::expr::Batch::try_from_rows_with_types(values, batch_datatypes)
387        .map_err(BoxedError::new)
388        .context(ExternalSnafu)?;
389    let batch = RecordBatch::new(flow.source_schema.clone(), batch.batch().to_vec())
390        .map_err(BoxedError::new)
391        .context(ExternalSnafu)?;
392    let provider = input_provider(&batch)?;
393    let source_ref = TableReference::full(
394        flow.source_table_name[0].clone(),
395        flow.source_table_name[1].clone(),
396        flow.source_table_name[2].clone(),
397    );
398    let plan = replace_source(flow.plan.clone(), &source_ref, provider)?;
399    let output = query_engine
400        .execute(plan, flow.query_ctx.clone())
401        .await
402        .map_err(BoxedError::new)
403        .context(ExternalSnafu)?;
404    let batches = match output.data {
405        OutputData::RecordBatches(batches) => batches,
406        OutputData::Stream(stream) => RecordBatches::try_collect(stream)
407            .await
408            .map_err(BoxedError::new)
409            .context(ExternalSnafu)?,
410        OutputData::AffectedRows(_) => {
411            return UnexpectedSnafu {
412                reason: "Streaming flow query returned affected rows",
413            }
414            .fail();
415        }
416    };
417
418    let output_schema = batches
419        .schema()
420        .column_schemas()
421        .iter()
422        .cloned()
423        .map(|mut column| {
424            column.data_type = map_dictionary_to_values_data_type(&column.data_type);
425            column
426        })
427        .collect::<Vec<_>>();
428    crate::adapter::validate_sink_layout_with_suffix(
429        &output_schema,
430        &flow.sink_schema,
431        &flow.auto_columns,
432    )?;
433    let mut output_rows = Vec::new();
434    for batch in batches {
435        let vectors = datatypes::vectors::Helper::try_into_vectors(batch.columns())
436            .map_err(BoxedError::new)
437            .context(ExternalSnafu)?;
438        output_rows.extend(vectors_to_rows(vectors.iter(), batch.num_rows()));
439    }
440    if output_rows.is_empty() {
441        return Ok(0);
442    }
443
444    // Auto columns are deliberately synthesized here rather than in the query plan. This keeps
445    // one current timestamp for the whole request and preserves the sink's timestamp precision.
446    let auto_values = synthesize_auto_values(&flow.auto_columns, Timestamp::current_millis())?
447        .into_iter()
448        .map(to_grpc_value)
449        .collect::<Vec<_>>();
450    for row in &mut output_rows {
451        row.values.extend(auto_values.iter().cloned());
452    }
453
454    let output_row_count = output_rows.len();
455    let proto_schema = column_schemas_to_proto(flow.sink_schema.clone(), &flow.sink_primary_keys)?;
456    let request = Request::RowInserts(RowInsertRequests {
457        inserts: vec![RowInsertRequest {
458            table_name: flow.sink_table_name[2].clone(),
459            rows: Some(Rows {
460                schema: proto_schema,
461                rows: output_rows,
462            }),
463        }],
464    });
465    let mut peer = None;
466    frontend_client
467        .handle_insert_once(
468            request,
469            &flow.sink_table_name[0],
470            &flow.sink_table_name[1],
471            &mut peer,
472        )
473        .await
474        .map_err(BoxedError::new)
475        .context(ExternalSnafu)?;
476    Ok(output_row_count)
477}
478
479#[cfg(test)]
480mod tests {
481    use datafusion::catalog::MemTable;
482    use datafusion::logical_expr::LogicalPlanBuilder;
483    use datatypes::data_type::ConcreteDataType;
484    use datatypes::schema::{ColumnSchema, Schema};
485    use datatypes::vectors::{Int32Vector, TimestampMillisecondVector};
486    use session::context::QueryContext;
487
488    use super::*;
489
490    #[test]
491    fn validation_accepts_distinct_over_one_source() {
492        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
493            "number",
494            ConcreteDataType::int32_datatype(),
495            false,
496        )]));
497        let plan = LogicalPlanBuilder::scan(
498            TableReference::bare("source"),
499            provider_as_source(provider(&schema, 1)),
500            None,
501        )
502        .unwrap()
503        .project(vec![datafusion_expr::col("number")])
504        .unwrap()
505        .distinct()
506        .unwrap()
507        .build()
508        .unwrap();
509        assert!(validate_plan(&plan).is_ok());
510    }
511
512    #[test]
513    fn validation_rejects_plan_without_source_scan() {
514        let plan = LogicalPlan::EmptyRelation(datafusion_expr::logical_plan::EmptyRelation {
515            produce_one_row: false,
516            schema: Arc::new(DFSchema::empty()),
517        });
518        assert!(validate_plan(&plan).is_err());
519    }
520
521    fn provider(schema: &SchemaRef, value: i32) -> Arc<dyn TableProvider> {
522        let batch = RecordBatch::new(
523            schema.clone(),
524            vec![Arc::new(Int32Vector::from_slice([value])) as datatypes::prelude::VectorRef],
525        )
526        .unwrap();
527        let arrow = batch.df_record_batch().clone();
528        Arc::new(MemTable::try_new(arrow.schema(), vec![vec![arrow]]).unwrap())
529    }
530
531    #[test]
532    fn source_schema_version_must_match_retained_plan() {
533        assert!(validate_source_schema_version(7, 7).is_ok());
534
535        let error = validate_source_schema_version(7, 8).unwrap_err();
536        assert!(matches!(error, Error::InvalidQuery { reason, .. } if
537            reason.contains("Source schema version changed from 7 to 8")
538                && reason.contains("recreate or recover")
539        ));
540    }
541
542    #[test]
543    fn auto_values_use_the_sink_timestamp_units() {
544        let update_at = ColumnSchema::new(
545            crate::adapter::AUTO_CREATED_UPDATE_AT_TS_COL,
546            datatypes::data_type::ConcreteDataType::timestamp_second_datatype(),
547            true,
548        );
549        let placeholder = ColumnSchema::new(
550            crate::adapter::AUTO_CREATED_PLACEHOLDER_TS_COL,
551            datatypes::data_type::ConcreteDataType::timestamp_nanosecond_datatype(),
552            true,
553        );
554        let values = synthesize_auto_values(
555            &[update_at],
556            Timestamp::new(1_234, common_time::timestamp::TimeUnit::Millisecond),
557        )
558        .unwrap();
559        assert_eq!(values.len(), 1);
560        assert_eq!(
561            values[0].as_timestamp().unwrap().unit(),
562            common_time::timestamp::TimeUnit::Second
563        );
564        assert!(values[0].as_timestamp().unwrap().value() > 0);
565
566        let values = synthesize_auto_values(
567            &[
568                ColumnSchema::new(
569                    crate::adapter::AUTO_CREATED_UPDATE_AT_TS_COL,
570                    datatypes::data_type::ConcreteDataType::timestamp_microsecond_datatype(),
571                    true,
572                ),
573                placeholder,
574            ],
575            Timestamp::new(1_234, common_time::timestamp::TimeUnit::Millisecond),
576        )
577        .unwrap();
578        assert_eq!(values[1].as_timestamp().unwrap().value(), 0);
579        assert_eq!(
580            values[1].as_timestamp().unwrap().unit(),
581            common_time::timestamp::TimeUnit::Nanosecond
582        );
583    }
584
585    #[test]
586    fn auto_values_reject_arbitrary_columns() {
587        let column = ColumnSchema::new(
588            "other",
589            datatypes::data_type::ConcreteDataType::timestamp_millisecond_datatype(),
590            true,
591        );
592        assert!(synthesize_auto_values(&[column], Timestamp::current_millis()).is_err());
593    }
594
595    #[test]
596    fn rewrite_source_timestamp_appends_collision_free_hidden_output() {
597        let schema = Arc::new(Schema::new(vec![
598            ColumnSchema::new("number", ConcreteDataType::int32_datatype(), false),
599            ColumnSchema::new(
600                "ts",
601                ConcreteDataType::timestamp_millisecond_datatype(),
602                false,
603            )
604            .with_time_index(true),
605            ColumnSchema::new(
606                "__flow_source_timestamp",
607                ConcreteDataType::int32_datatype(),
608                false,
609            ),
610        ]));
611        let batch = RecordBatch::new(
612            schema.clone(),
613            vec![
614                Arc::new(Int32Vector::from_slice([2])) as datatypes::prelude::VectorRef,
615                Arc::new(TimestampMillisecondVector::from_slice([42]))
616                    as datatypes::prelude::VectorRef,
617                Arc::new(Int32Vector::from_slice([7])) as datatypes::prelude::VectorRef,
618            ],
619        )
620        .unwrap();
621        let plan = datafusion_expr::LogicalPlanBuilder::scan(
622            TableReference::bare("source"),
623            provider_as_source(input_provider(&batch).unwrap()),
624            None,
625        )
626        .unwrap()
627        .filter(datafusion_expr::col("number").gt(datafusion_expr::lit(1)))
628        .unwrap()
629        .project(vec![datafusion_expr::col("number")])
630        .unwrap()
631        .build()
632        .unwrap();
633        let rewritten =
634            rewrite_source_timestamp(plan, &TableReference::bare("source"), "ts").unwrap();
635        assert_eq!(rewritten.schema().fields().len(), 2);
636        assert!(
637            rewritten
638                .schema()
639                .field(1)
640                .name()
641                .starts_with("__flow_source_timestamp")
642        );
643    }
644
645    #[tokio::test]
646    async fn finite_projection_filter_does_not_retain_previous_batch() {
647        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
648            "number",
649            ConcreteDataType::int32_datatype(),
650            false,
651        )]));
652        let initial_provider = provider(&schema, 0);
653        let plan = test_source_plan(TableReference::bare("source"), initial_provider);
654        assert!(validate_plan(&plan).is_ok());
655
656        let engine = crate::test_utils::create_test_query_engine();
657        let run = |input: Arc<dyn TableProvider>| {
658            let engine = engine.clone();
659            let plan = plan.clone();
660            async move {
661                let plan = replace_source(plan, &TableReference::bare("source"), input).unwrap();
662                let output = engine.execute(plan, QueryContext::arc()).await.unwrap();
663                match output.data {
664                    OutputData::Stream(stream) => RecordBatches::try_collect(stream)
665                        .await
666                        .unwrap()
667                        .iter()
668                        .map(|batch| batch.num_rows())
669                        .sum(),
670                    OutputData::RecordBatches(batches) => {
671                        batches.iter().map(|batch| batch.num_rows()).sum()
672                    }
673                    OutputData::AffectedRows(_) => 0,
674                }
675            }
676        };
677
678        assert_eq!(run(provider(&schema, 1)).await, 0);
679        assert_eq!(run(provider(&schema, 2)).await, 1);
680    }
681
682    #[tokio::test]
683    async fn finite_distinct_collapses_duplicates_per_request() {
684        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
685            "number",
686            ConcreteDataType::int32_datatype(),
687            false,
688        )]));
689        let input = |values: &[i32]| {
690            let batch = RecordBatch::new(
691                schema.clone(),
692                vec![Arc::new(Int32Vector::from_slice(values)) as datatypes::prelude::VectorRef],
693            )
694            .unwrap();
695            input_provider(&batch).unwrap()
696        };
697        let plan = LogicalPlanBuilder::scan(
698            TableReference::bare("source"),
699            provider_as_source(input(&[1])),
700            None,
701        )
702        .unwrap()
703        .project(vec![datafusion_expr::col("number")])
704        .unwrap()
705        .distinct()
706        .unwrap()
707        .build()
708        .unwrap();
709        assert!(validate_plan(&plan).is_ok());
710
711        let engine = crate::test_utils::create_test_query_engine();
712        let run = |provider: Arc<dyn TableProvider>| {
713            let engine = engine.clone();
714            let plan = plan.clone();
715            async move {
716                let plan = replace_source(plan, &TableReference::bare("source"), provider).unwrap();
717                let output = engine.execute(plan, QueryContext::arc()).await.unwrap();
718                let batches = match output.data {
719                    OutputData::Stream(stream) => RecordBatches::try_collect(stream).await.unwrap(),
720                    OutputData::RecordBatches(batches) => batches,
721                    OutputData::AffectedRows(_) => panic!("unexpected affected rows"),
722                };
723                batches.iter().map(|batch| batch.num_rows()).sum::<usize>()
724            }
725        };
726
727        assert_eq!(run(input(&[1, 1, 2])).await, 2);
728        assert_eq!(run(input(&[1, 1])).await, 1);
729    }
730}
flow/adapter/stateless.rs

flow/adapter/
stateless.rs
