From cc5b894651693368fd83d9e3c2c4451e2b65c4ed Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 23 Sep 2026 14:12:01 +0800 Subject: [PATCH] fix(flow): use DISTINCT output columns as sink key for streaming flows When a streaming flow's source table has no primary key, sink key inference produced an empty key. For SELECT DISTINCT v, the DISTINCT establishes v as the output key, but the inference only looked at source primary keys. Combined with the zero-valued placeholder timestamp, all output rows shared the same storage identity and were deduplicated into a single row. Use all non-time-index output columns as the sink key when the plan is a DISTINCT and no source primary key was inherited. Signed-off-by: discord9 --- src/flow/src/adapter.rs | 20 +++++++++++--- src/flow/src/adapter/tests.rs | 51 ++++++++++++++++++++++++++++++++--- 2 files changed, 65 insertions(+), 6 deletions(-) diff --git a/src/flow/src/adapter.rs b/src/flow/src/adapter.rs index c98510bd7d5..7e17032e2e0 100644 --- a/src/flow/src/adapter.rs +++ b/src/flow/src/adapter.rs @@ -141,8 +141,9 @@ fn relation_desc_from_output( columns: &[ColumnSchema], lineage: &[Option], source_primary_key_indices: &[usize], + plan: &LogicalPlan, ) -> RelationDesc { - let keys = source_primary_key_indices + let mut keys = source_primary_key_indices .iter() .filter_map(|source_index| { lineage @@ -150,6 +151,15 @@ fn relation_desc_from_output( .position(|index| index == &Some(*source_index)) }) .collect_vec(); + // A `DISTINCT` establishes the output key by itself: every non-time-index output + // column is part of the distinct key. When no source primary key is inherited, + // fall back to the DISTINCT output key so the sink relation keeps row identity. + if keys.is_empty() && matches!(plan, LogicalPlan::Distinct(Distinct::All(_))) { + keys = columns + .iter() + .positions(|column| !column.is_time_index()) + .collect(); + } let time_index = columns.iter().position(ColumnSchema::is_time_index); RelationDesc { typ: crate::repr::RelationType { @@ -849,8 +859,12 @@ impl StreamingEngine { let source_primary_key_indices = source_meta.primary_key_indices; stateless::validate_source_scan(&flow_plan, source_table_id, &source_schema)?; let (inferred_schema, lineage) = output_column_schemas(&flow_plan, &source_schema)?; - let inferred_relation = - relation_desc_from_output(&inferred_schema, &lineage, &source_primary_key_indices); + let inferred_relation = relation_desc_from_output( + &inferred_schema, + &lineage, + &source_primary_key_indices, + &flow_plan, + ); let sink_exists = self.fetch_table_pk_schema(sink_table_name).await?.is_some(); if !sink_exists { validate_auto_column_names(&inferred_schema)?; diff --git a/src/flow/src/adapter/tests.rs b/src/flow/src/adapter/tests.rs index 6bc02f60786..dba7a950bb6 100644 --- a/src/flow/src/adapter/tests.rs +++ b/src/flow/src/adapter/tests.rs @@ -20,7 +20,7 @@ use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, TimeUnit as use datafusion::catalog::MemTable; use datafusion::datasource::provider_as_source; use datafusion_common::TableReference; -use datafusion_expr::LogicalPlanBuilder; +use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema, Schema, SchemaBuilder}; use store_api::storage::{ConcreteDataType, TableId}; use table::metadata::{TableInfo, TableInfoBuilder, TableMetaBuilder}; @@ -194,10 +194,55 @@ fn stateless_distinct_preserves_direct_column_lineage() { let (output, lineage) = super::output_column_schemas(&plan, &source).unwrap(); assert_eq!(output[0].name, "dis"); assert_eq!(lineage, vec![Some(0)]); - let relation = super::relation_desc_from_output(&output, &lineage, &[0]); + let relation = super::relation_desc_from_output(&output, &lineage, &[0], &plan); assert_eq!(relation.typ.keys[0].column_indices, vec![0]); } +#[test] +fn stateless_distinct_without_source_pk_uses_distinct_output_as_key() { + let source = Arc::new(Schema::new(vec![ColumnSchema::new( + "v", + ConcreteDataType::int32_datatype(), + true, + )])); + let provider = MemTable::try_new( + Arc::new(datafusion::arrow::datatypes::Schema::new(vec![Field::new( + "v", + ArrowDataType::Int32, + true, + )])), + vec![vec![]], + ) + .unwrap(); + let plan = LogicalPlanBuilder::scan( + TableReference::bare("source"), + provider_as_source(Arc::new(provider)), + None, + ) + .unwrap() + .project(vec![datafusion_expr::col("v")]) + .unwrap() + .distinct() + .unwrap() + .build() + .unwrap(); + + let (output, lineage) = super::output_column_schemas(&plan, &source).unwrap(); + let relation = super::relation_desc_from_output(&output, &lineage, &[], &plan); + assert_eq!(relation.typ.keys.len(), 1); + assert_eq!(relation.typ.keys[0].column_indices, vec![0]); + + // Non-DISTINCT plans keep the empty-key behavior. + let project_plan = match &plan { + LogicalPlan::Distinct(datafusion_expr::logical_plan::Distinct::All(input)) => { + input.as_ref().clone() + } + _ => panic!("expected distinct plan"), + }; + let relation = super::relation_desc_from_output(&output, &lineage, &[], &project_plan); + assert!(relation.typ.keys.is_empty()); +} + #[test] fn stateless_normalizes_dictionary_output_type() { let field = Field::new_dictionary("host", ArrowDataType::UInt32, ArrowDataType::Utf8, true); @@ -220,7 +265,7 @@ fn stateless_normalizes_dictionary_output_type() { assert_eq!(output[0].data_type, ConcreteDataType::string_datatype()); assert_eq!(lineage, vec![Some(0)]); - let relation = super::relation_desc_from_output(&output, &lineage, &[0]); + let relation = super::relation_desc_from_output(&output, &lineage, &[0], &plan); assert_eq!(relation.typ.keys[0].column_indices, vec![0]); }