Skip to content

Commit a62969d

Browse files
limadog9limadog9
committed
fix: Unnest must not keep stale functional dependencies
Remap dependency indices after struct expansion and invalidate or downgrade dependencies affected by list UNNEST. Preserve inherited dependency metadata and cover DISTINCT, GROUP BY, ordered windows, and shifted columns.
1 parent dd1d709 commit a62969d

2 files changed

Lines changed: 203 additions & 3 deletions

File tree

‎datafusion/expr/src/logical_plan/plan.rs‎

Lines changed: 116 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5190,8 +5190,47 @@ impl Unnest {
51905190

51915191
let metadata = input_schema.metadata().clone();
51925192
let df_schema = DFSchema::new_with_metadata(fields, metadata)?;
5193-
// We can use the existing functional dependencies:
5194-
let deps = input_schema.functional_dependencies().clone();
5193+
let list_indices: HashSet<_> =
5194+
list_columns.iter().map(|(index, _)| *index).collect();
5195+
// Struct unnest may replace one input column with several output columns,
5196+
// so remap dependencies using the input index of each output column.
5197+
let remap_indices = |indices: &[usize], exclude_lists: bool| {
5198+
dependency_indices
5199+
.iter()
5200+
.enumerate()
5201+
.filter_map(|(output_index, input_index)| {
5202+
(indices.contains(input_index)
5203+
&& (!exclude_lists || !list_indices.contains(input_index)))
5204+
.then_some(output_index)
5205+
})
5206+
.collect()
5207+
};
5208+
let deps = FunctionalDependencies::new(
5209+
input_schema
5210+
.functional_dependencies()
5211+
.iter()
5212+
// A list input no longer determines the rows produced by unnest.
5213+
// A struct input is also unsafe as a source: a null struct and
5214+
// a non-null struct with all-null fields have the same outputs.
5215+
.filter(|dep| {
5216+
dep.source_indices.iter().all(|index| {
5217+
!list_indices.contains(index) && !struct_columns.contains(index)
5218+
})
5219+
})
5220+
.map(|dep| FunctionalDependence {
5221+
source_indices: remap_indices(&dep.source_indices, false),
5222+
// The unnested values are not determined by the input key.
5223+
target_indices: remap_indices(&dep.target_indices, true),
5224+
// One input row can produce multiple output rows.
5225+
mode: if list_indices.is_empty() {
5226+
dep.mode
5227+
} else {
5228+
Dependency::Multi
5229+
},
5230+
..dep.clone()
5231+
})
5232+
.collect(),
5233+
);
51955234
let schema = Arc::new(df_schema.with_functional_dependencies(deps)?);
51965235

51975236
Ok(Unnest {
@@ -5297,7 +5336,7 @@ mod tests {
52975336
use datafusion_common::tree_node::{
52985337
TransformedResult, TreeNodeRewriter, TreeNodeVisitor,
52995338
};
5300-
use datafusion_common::{Constraint, not_impl_err};
5339+
use datafusion_common::{Constraint, RecursionUnnestOption, not_impl_err};
53015340
use insta::{assert_debug_snapshot, assert_snapshot};
53025341
use std::hash::DefaultHasher;
53035342

@@ -5377,6 +5416,80 @@ mod tests {
53775416
])
53785417
}
53795418

5419+
#[test]
5420+
fn unnest_remaps_functional_dependencies() -> Result<()> {
5421+
let schema = Arc::new(Schema::new(vec![
5422+
Field::new_struct(
5423+
"s",
5424+
vec![
5425+
Field::new("a", DataType::Int32, true),
5426+
Field::new("b", DataType::Int32, true),
5427+
],
5428+
true,
5429+
),
5430+
Field::new("k", DataType::Int32, false),
5431+
Field::new_list(
5432+
"xs",
5433+
Field::new_list(
5434+
"item",
5435+
Field::new_list_field(DataType::Int32, true),
5436+
true,
5437+
),
5438+
true,
5439+
),
5440+
Field::new("payload", DataType::Utf8, true),
5441+
]));
5442+
let constraints = Constraints::new_unverified(vec![
5443+
Constraint::PrimaryKey(vec![1]),
5444+
Constraint::Unique(vec![0]),
5445+
Constraint::Unique(vec![2]),
5446+
]);
5447+
let source: Arc<dyn TableSource> =
5448+
Arc::new(LogicalTableSource::new(schema).with_constraints(constraints));
5449+
5450+
// Expanding a struct preserves row count and remaps keys and targets,
5451+
// but a dependency sourced from the struct itself cannot be kept.
5452+
let struct_only = LogicalPlanBuilder::scan("t", Arc::clone(&source), None)?
5453+
.unnest_column("s")?
5454+
.build()?;
5455+
assert_eq!(
5456+
struct_only.schema().functional_dependencies(),
5457+
&FunctionalDependencies::new(vec![
5458+
FunctionalDependence::new(vec![2], vec![0, 1, 2, 3, 4], false)
5459+
.with_mode(Dependency::Single),
5460+
FunctionalDependence::new(vec![3], vec![0, 1, 2, 3, 4], true)
5461+
.with_mode(Dependency::Single),
5462+
])
5463+
);
5464+
5465+
// The two recursion outputs both replace the same list input. Neither
5466+
// is determined by k, and a list-sourced dependency is no longer valid.
5467+
let options = UnnestOptions::default()
5468+
.with_recursions(RecursionUnnestOption {
5469+
input_column: "xs".into(),
5470+
output_column: "xs_depth_1".into(),
5471+
depth: 1,
5472+
})
5473+
.with_recursions(RecursionUnnestOption {
5474+
input_column: "xs".into(),
5475+
output_column: "xs_depth_2".into(),
5476+
depth: 2,
5477+
});
5478+
let plan = LogicalPlanBuilder::scan("t", source, None)?
5479+
.unnest_columns_with_options(vec!["s".into(), "xs".into()], options)?
5480+
.build()?;
5481+
assert_eq!(plan.schema().fields().len(), 6);
5482+
assert_eq!(
5483+
plan.schema().functional_dependencies(),
5484+
&FunctionalDependencies::new(vec![
5485+
FunctionalDependence::new(vec![2], vec![0, 1, 2, 5], false,)
5486+
.with_mode(Dependency::Multi)
5487+
])
5488+
);
5489+
5490+
Ok(())
5491+
}
5492+
53805493
#[test]
53815494
fn projection_with_leading_computed_column_preserves_pk() -> Result<()> {
53825495
let constraints =

‎datafusion/sqllogictest/test_files/functional_dependencies.slt‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,93 @@ drop table t_null;
296296
statement ok
297297
drop table t_probe;
298298

299+
##########
300+
## List unnest invalidates dependencies on the input row
301+
##########
302+
303+
statement ok
304+
CREATE TABLE t_unnest_fd (k VARCHAR, v INT) AS VALUES ('a', 1), ('a', 2), ('b', 3);
305+
306+
# DISTINCT must retain the unnested value in its grouping key.
307+
query I
308+
WITH g AS (SELECT k, array_agg(v) AS vs FROM t_unnest_fd GROUP BY k),
309+
u AS (SELECT k, unnest(vs) AS v FROM g)
310+
SELECT count(*) AS n FROM (SELECT DISTINCT k, v FROM u);
311+
----
312+
3
313+
314+
# The equivalent GROUP BY must also retain the unnested value.
315+
query I
316+
WITH g AS (SELECT k, array_agg(v) AS vs FROM t_unnest_fd GROUP BY k),
317+
u AS (SELECT k, unnest(vs) AS v FROM g)
318+
SELECT count(*) AS n FROM (SELECT k, v FROM u GROUP BY k, v);
319+
----
320+
3
321+
322+
# There is no projection between Unnest and DISTINCT to discard a stale target.
323+
query I
324+
WITH g AS (SELECT k, array_agg(v) AS vs FROM t_unnest_fd GROUP BY k)
325+
SELECT count(*) AS n FROM (SELECT DISTINCT * FROM (SELECT k, unnest(vs) FROM g));
326+
----
327+
3
328+
329+
# An unnested list of structs also creates multiple values per input key.
330+
query I
331+
WITH g AS (
332+
SELECT k, array_agg(named_struct('id', v)) AS rs
333+
FROM t_unnest_fd GROUP BY k
334+
), u AS (SELECT k, unnest(rs) AS r FROM g)
335+
SELECT count(*) AS n FROM (SELECT DISTINCT k, r['id'] FROM u);
336+
----
337+
3
338+
339+
# Struct expansion shifts the following list output onto an old dependency
340+
# target index. DISTINCT must still include the list output.
341+
query I
342+
WITH g AS (
343+
SELECT k, named_struct('x', min(v), 'y', max(v)) AS s,
344+
array_agg(v) AS vs, sum(v) AS n
345+
FROM t_unnest_fd GROUP BY k
346+
)
347+
SELECT count(*) FROM (
348+
SELECT DISTINCT * FROM (SELECT k, unnest(s), unnest(vs), n FROM g)
349+
);
350+
----
351+
3
352+
353+
# ORDER BY k has ties after unnest. Its default RANGE frame includes both rows with k = 'a'.
354+
query I
355+
WITH g AS (SELECT k, array_agg(v) AS vs FROM t_unnest_fd GROUP BY k),
356+
u AS (SELECT k, unnest(vs) AS v FROM g)
357+
SELECT count(DISTINCT s) AS n FROM (SELECT sum(v) OVER (ORDER BY k) AS s FROM u);
358+
----
359+
2
360+
361+
# Expanding a struct before k shifts the inherited dependency's column indices.
362+
# The two sy values tie, so their default RANGE frames have the same sum.
363+
query I
364+
WITH g AS (
365+
SELECT named_struct('x', min(v), 'y', length(k)) AS s,
366+
k, sum(v) AS n
367+
FROM t_unnest_fd GROUP BY k
368+
), u(sx, sy, k, n) AS (SELECT unnest(s), k, n FROM g)
369+
SELECT count(DISTINCT running_sum) FROM (
370+
SELECT sum(n) OVER (ORDER BY sy) AS running_sum FROM u
371+
);
372+
----
373+
1
374+
375+
# A row-preserving, single-element unnest remains correct.
376+
query I
377+
SELECT count(*) AS n
378+
FROM (SELECT DISTINCT k, v
379+
FROM (SELECT k, unnest(make_array(v)) AS v FROM t_unnest_fd));
380+
----
381+
3
382+
383+
statement ok
384+
drop table t_unnest_fd;
385+
299386
##########
300387
## Cleanup
301388
##########

0 commit comments

Comments
 (0)