Skip to content

Commit 535f10e

Browse files
alambde-bgunter
authored andcommitted
Fix internal error "Physical input schema should be the same as the one converted from logical input schema." (apache#18412)
## Which issue does this PR close? - Closes apache#18337 ## Rationale for this change It is a bug we are seeing in our production related to a schema mismatch ## What changes are included in this PR? 1. New slt test for the issue 2. Properly compute the output field from window functions ## Are these changes tested? Yes ## Are there any user-facing changes?
1 parent 783b794 commit 535f10e

2 files changed

Lines changed: 51 additions & 22 deletions

File tree

datafusion/expr/src/expr_schema.rs

Lines changed: 13 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ use crate::type_coercion::functions::fields_with_udf;
2525
use crate::udf::ReturnFieldArgs;
2626
use crate::{LogicalPlan, Projection, Subquery, WindowFunctionDefinition, utils};
2727
use arrow::compute::can_cast_types;
28-
use arrow::datatypes::{DataType, Field};
28+
use arrow::datatypes::{DataType, Field, FieldRef};
2929
use datafusion_common::datatype::FieldExt;
3030
use datafusion_common::metadata::FieldMetadata;
3131
use datafusion_common::{
@@ -156,9 +156,10 @@ impl ExprSchemable for Expr {
156156
let return_type = self.to_field(schema)?.1.data_type().clone();
157157
Ok(return_type)
158158
}
159-
Expr::WindowFunction(window_function) => self
160-
.data_type_and_nullable_with_window_function(schema, window_function)
161-
.map(|(return_type, _)| return_type),
159+
Expr::WindowFunction(window_function) => Ok(self
160+
.window_function_field(schema, window_function)?
161+
.data_type()
162+
.clone()),
162163
Expr::AggregateFunction(AggregateFunction {
163164
func,
164165
params: AggregateFunctionParams { args, .. },
@@ -357,12 +358,9 @@ impl ExprSchemable for Expr {
357358
Expr::AggregateFunction(AggregateFunction { func, .. }) => {
358359
Ok(func.is_nullable())
359360
}
360-
Expr::WindowFunction(window_function) => self
361-
.data_type_and_nullable_with_window_function(
362-
input_schema,
363-
window_function,
364-
)
365-
.map(|(_, nullable)| nullable),
361+
Expr::WindowFunction(window_function) => Ok(self
362+
.window_function_field(input_schema, window_function)?
363+
.is_nullable()),
366364
Expr::ScalarVariable(field, _) => Ok(field.is_nullable()),
367365
Expr::TryCast { .. } | Expr::Unnest(_) | Expr::Placeholder(_) => Ok(true),
368366
Expr::IsNull(_)
@@ -458,7 +456,7 @@ impl ExprSchemable for Expr {
458456
/// with the default implementation returning empty field metadata
459457
/// - **Aggregate functions**: Generate metadata via function's [`return_field`] method,
460458
/// with the default implementation returning empty field metadata
461-
/// - **Window functions**: field metadata is empty
459+
/// - **Window functions**: field metadata follows the function's return field
462460
///
463461
/// ## Table Reference Scoping
464462
/// - Establishes proper qualified field references when columns belong to specific tables
@@ -534,11 +532,7 @@ impl ExprSchemable for Expr {
534532
)))
535533
}
536534
Expr::WindowFunction(window_function) => {
537-
let (dt, nullable) = self.data_type_and_nullable_with_window_function(
538-
schema,
539-
window_function,
540-
)?;
541-
Ok(Arc::new(Field::new(&schema_name, dt, nullable)))
535+
self.window_function_field(schema, window_function)
542536
}
543537
Expr::AggregateFunction(aggregate_function) => {
544538
let AggregateFunction {
@@ -698,11 +692,11 @@ impl Expr {
698692
///
699693
/// Otherwise, returns an error if there's a type mismatch between
700694
/// the window function's signature and the provided arguments.
701-
fn data_type_and_nullable_with_window_function(
695+
fn window_function_field(
702696
&self,
703697
schema: &dyn ExprSchema,
704698
window_function: &WindowFunction,
705-
) -> Result<(DataType, bool)> {
699+
) -> Result<FieldRef> {
706700
let WindowFunction {
707701
fun,
708702
params: WindowFunctionParams { args, .. },
@@ -738,9 +732,7 @@ impl Expr {
738732
.into_iter()
739733
.collect::<Vec<_>>();
740734

741-
let return_field = udaf.return_field(&new_fields)?;
742-
743-
Ok((return_field.data_type().clone(), return_field.is_nullable()))
735+
udaf.return_field(&new_fields)
744736
}
745737
WindowFunctionDefinition::WindowUDF(udwf) => {
746738
let data_types = fields
@@ -769,7 +761,6 @@ impl Expr {
769761
let field_args = WindowUDFFieldArgs::new(&new_fields, &function_name);
770762

771763
udwf.field(field_args)
772-
.map(|field| (field.data_type().clone(), field.is_nullable()))
773764
}
774765
}
775766
}

datafusion/sqllogictest/test_files/metadata.slt

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,22 @@
2424
## in the test harness as there is no way to define schema
2525
## with metadata in SQL.
2626

27+
query ITTPT
28+
select * from table_with_metadata;
29+
----
30+
1 NULL NULL 2020-09-08T13:42:29.190855123 no_foo
31+
NULL bar l_bar 2020-09-08T13:42:29.190855123 no_bar
32+
3 baz l_baz 2020-09-08T13:42:29.190855123 no_baz
33+
34+
query TTT
35+
describe table_with_metadata;
36+
----
37+
id Int32 YES
38+
name Utf8 YES
39+
l_name Utf8 YES
40+
ts Timestamp(ns) NO
41+
nonnull_name Utf8 NO
42+
2743
query IT
2844
select id, name from table_with_metadata;
2945
----
@@ -235,6 +251,28 @@ order by 1 asc nulls last;
235251
3 1
236252
NULL 1
237253

254+
# Reproducer for https://github.com/apache/datafusion/issues/18337
255+
# this query should not get an internal error
256+
query TI
257+
SELECT
258+
'foo' AS name,
259+
COUNT(
260+
CASE
261+
WHEN prev_value = 'no_bar' AND value = 'no_baz' THEN 1
262+
ELSE NULL
263+
END
264+
) AS count_rises
265+
FROM
266+
(
267+
SELECT
268+
nonnull_name as value,
269+
LAG(nonnull_name) OVER (ORDER BY ts) AS prev_value
270+
FROM
271+
table_with_metadata
272+
);
273+
----
274+
foo 1
275+
238276
# Regression test: first_value should preserve metadata
239277
query IT
240278
select first_value(id order by id asc nulls last), arrow_metadata(first_value(id order by id asc nulls last), 'metadata_key')

0 commit comments

Comments
 (0)