Skip to content

SubqueryAlias, Values, and/or EmptyRelation have incorrect schemas after replacing Placeholder values #18102

Description

@paleolimbot

Describe the bug

When trying to bind parameters and execute queries like "SELECT a, b FROM (VALUES ($1, $2)) AS t(a, b)" and "SELECT $1, $2", I get mismatched schemas and/or an error when I try to execute the query.

I think this is because the SubqueryAlias, Values, and/or EmptyRelation both cache their schema and this is not updated when updating parameter values. (There may be other logical plan items that cache their schemas too)

To Reproduce

For SubqueryAlias:

let df = ctx.sql("SELECT a, b FROM (VALUES ($1, $2)) AS t(a, b)").await?;
let df_with_params_replaced = df.with_param_values(vec![
    ScalarValue::UInt32(Some(1)),
    ScalarValue::Utf8(Some("foofy".to_string())),
])?;
dbg!(df_with_params_replaced.clone().schema().as_arrow());
dbg!(df_with_params_replaced.collect().await?[0].schema());
#> Error: ArrowError(InvalidArgumentError("column types must match schema types, expected Null but found UInt32 at column index 0"), Some(""))

For EmptyRelation (mismatched schemas but execution works):

let df = ctx.sql("SELECT $1, $2").await?;
let df_with_params_replaced = df.with_param_values(vec![
    ScalarValue::UInt32(Some(1)),
    ScalarValue::Utf8(Some("foofy".to_string())),
])?;
dbg!(df_with_params_replaced.clone().schema().as_arrow());
dbg!(df_with_params_replaced.collect().await?[0].schema());
Details
[datafusion/core/tests/sql/select.rs:333:5] df_with_params_replaced.clone().schema().as_arrow() = Schema {
    fields: [
        Field {
            name: "$1",
            data_type: Null,
            nullable: true,
            dict_id: 0,
            dict_is_ordered: false,
            metadata: {},
        },
        Field {
            name: "$2",
            data_type: Null,
            nullable: true,
            dict_id: 0,
            dict_is_ordered: false,
            metadata: {},
        },
    ],
    metadata: {},
}
[datafusion/core/tests/sql/select.rs:334:5] df_with_params_replaced.collect().await?[0].schema() = Schema {
    fields: [
        Field {
            name: "$1",
            data_type: UInt32,
            nullable: false,
            dict_id: 0,
            dict_is_ordered: false,
            metadata: {},
        },
        Field {
            name: "$2",
            data_type: Utf8,
            nullable: false,
            dict_id: 0,
            dict_is_ordered: false,
            metadata: {},
        },
    ],
    metadata: {},
}

Expected behavior

I expected the schema and generated batches to match and the queries to execute without error.

Additional context

Encountered whist writing tests for #17986

SubqueryAlias:

/// Aliased subquery
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
// mark non_exhaustive to encourage use of try_new/new()
#[non_exhaustive]
pub struct SubqueryAlias {
/// The incoming logical plan
pub input: Arc<LogicalPlan>,
/// The alias for the input relation
pub alias: TableReference,
/// The schema with qualified field names
pub schema: DFSchemaRef,
}

Values:

/// Values expression. See
/// [Postgres VALUES](https://www.postgresql.org/docs/current/queries-values.html)
/// documentation for more details.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Values {
/// The table schema
pub schema: DFSchemaRef,
/// Values
pub values: Vec<Vec<Expr>>,
}

EmptyRelation:

/// Relationship produces 0 or 1 placeholder rows with specified output schema
/// In most cases the output schema for `EmptyRelation` would be empty,
/// however, it can be non-empty typically for optimizer rules
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct EmptyRelation {
/// Whether to produce a placeholder row
pub produce_one_row: bool,
/// The schema description of the output
pub schema: DFSchemaRef,
}

Replacing Parameters:

pub fn replace_params_with_values(
self,
param_values: &ParamValues,
) -> Result<LogicalPlan> {
self.transform_up_with_subqueries(|plan| {
let schema = Arc::clone(plan.schema());
let name_preserver = NamePreserver::new(&plan);
plan.map_expressions(|e| {
let (e, has_placeholder) = e.infer_placeholder_types(&schema)?;
if !has_placeholder {
// Performance optimization:
// avoid NamePreserver copy and second pass over expression
// if no placeholders.
Ok(Transformed::no(e))
} else {
let original_name = name_preserver.save(&e);
let transformed_expr = e.transform_up(|e| {
if let Expr::Placeholder(Placeholder { id, .. }) = e {
let value = param_values.get_placeholders_with_values(&id)?;
Ok(Transformed::yes(Expr::Literal(value, None)))
} else {
Ok(Transformed::no(e))
}
})?;
// Preserve name to avoid breaking column references to this expression
Ok(transformed_expr.update_data(|expr| original_name.restore(expr)))
}
})
})
.map(|res| res.data)
}

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

bugSomething isn't working

Type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions