diff --git a/datafusion/core/tests/sql/select.rs b/datafusion/core/tests/sql/select.rs index c5c1ae54e75f8..fae2fb5bc5e76 100644 --- a/datafusion/core/tests/sql/select.rs +++ b/datafusion/core/tests/sql/select.rs @@ -387,6 +387,79 @@ async fn test_query_parameters_with_metadata() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_limit_offset_parameters_named() -> Result<()> { + let ctx = SessionContext::new(); + let df = ctx + .sql( + "SELECT value FROM (VALUES (10), (20), (30)) AS t(value) \ + ORDER BY value LIMIT $rows OFFSET $skip", + ) + .await?; + assert_eq!( + df.logical_plan().get_parameter_types()?, + HashMap::from([ + ("$rows".to_string(), Some(DataType::Int64)), + ("$skip".to_string(), Some(DataType::Int64)), + ]) + ); + let results = df + .with_param_values(vec![ + ("rows", ScalarValue::Int64(Some(1))), + ("skip", ScalarValue::Int64(Some(1))), + ])? + .collect() + .await?; + datafusion::assert_batches_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| 20 |", + "+-------+" + ], + &results + ); + Ok(()) +} + +#[tokio::test] +async fn test_limit_offset_parameters_keep_field_metadata() -> Result<()> { + let ctx = SessionContext::new(); + let metadata = HashMap::from([("some_key".to_string(), "some_value".to_string())]); + let schema = Arc::new(Schema::new(vec![ + Field::new("value", DataType::Int32, false).with_metadata(metadata.clone()), + ])); + ctx.register_batch( + "t", + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 1]))])?, + )?; + + for clause in ["LIMIT", "OFFSET"] { + let sql = format!("SELECT $1 AS value FROM t WHERE value = $1 {clause} $1"); + let df = ctx.sql(&sql).await?; + let results = df + .with_param_values(ParamValues::List(vec![ScalarAndMetadata::new( + ScalarValue::Int32(Some(1)), + Some(metadata.clone().into()), + )]))? + .collect() + .await?; + datafusion::assert_batches_eq!( + [ + "+-------+", + "| value |", + "+-------+", + "| 1 |", + "+-------+" + ], + &results + ); + } + + Ok(()) +} + #[tokio::test] async fn test_version_function() { let expected_version = format!( diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 92c1c03c14b9a..b727b700d1b17 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -1836,14 +1836,21 @@ impl LogicalPlan { .collect()) } - /// Walk the logical plan, find any `Placeholder` tokens, and return a map of their IDs and FieldRefs + /// Walk the logical plan, find any `Placeholder` tokens, and return a map of their IDs and FieldRefs. + /// Bare `LIMIT`/`OFFSET` parameters default to `Int64` if no occurrence provides a type. pub fn get_parameter_fields( &self, ) -> Result>, DataFusionError> { let mut param_types: HashMap> = HashMap::new(); + let mut row_count_parameters: HashSet = HashSet::new(); self.apply_with_subqueries(|plan| { plan.apply_expressions(|expr| { + if matches!(plan, LogicalPlan::Limit(_)) + && let Expr::Placeholder(Placeholder { id, field: None }) = expr + { + row_count_parameters.insert(id.clone()); + } expr.apply(|expr| { if let Expr::Placeholder(Placeholder { id, field }) = expr { let prev = param_types.get(id); @@ -1860,15 +1867,22 @@ impl LogicalPlan { param_types.insert(id.clone(), Some(Arc::clone(field))); } _ => { - param_types.insert(id.clone(), None); + param_types.entry(id.clone()).or_insert(None); } } } Ok(TreeNodeRecursion::Continue) }) }) - }) - .map(|_| param_types) + })?; + + for id in row_count_parameters { + param_types + .entry(id) + .or_default() + .get_or_insert_with(|| Arc::new(Field::new("", DataType::Int64, true))); + } + Ok(param_types) } // ------------ diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index c45bc5678ac87..930048339ab4a 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -100,6 +100,25 @@ impl AnalyzerRule for TypeCoercion { fn analyze(&self, plan: LogicalPlan, config: &ConfigOptions) -> Result { static EMPTY_SCHEMA: LazyLock = LazyLock::new(DFSchema::empty); + // Record bare row-count parameter types before coercion makes them indistinguishable + // from explicit casts. A type inferred elsewhere takes precedence over the default. + let parameter_fields = plan.get_parameter_fields()?; + let plan = plan + .transform_up_with_subqueries(|plan| { + if !matches!(plan, LogicalPlan::Limit(_)) { + return Ok(Transformed::no(plan)); + } + plan.map_expressions(|expr| match expr { + Expr::Placeholder(mut placeholder) if placeholder.field.is_none() => { + placeholder.field = + parameter_fields.get(&placeholder.id).cloned().flatten(); + Ok(Transformed::yes(Expr::Placeholder(placeholder))) + } + expr => Ok(Transformed::no(expr)), + }) + })? + .data; + // recurse let transformed_plan = plan .transform_up_with_subqueries(|plan| analyze_internal(&EMPTY_SCHEMA, plan))? @@ -1631,10 +1650,10 @@ mod test { use arrow::datatypes::{DataType, Field, Schema, SchemaBuilder, TimeUnit}; use insta::assert_snapshot; - use crate::analyzer::Analyzer; use crate::analyzer::type_coercion::{ TypeCoercion, TypeCoercionRewriter, coerce_case_expression, }; + use crate::analyzer::{Analyzer, AnalyzerRule}; use crate::assert_analyzed_plan_with_config_eq_snapshot; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{TransformedResult, TreeNode}; @@ -1646,9 +1665,9 @@ mod test { use datafusion_expr::test::function_stub::avg_udaf; use datafusion_expr::{ AccumulatorFactoryFunction, AggregateUDF, BinaryExpr, Case, ColumnarValue, Expr, - ExprSchemable, Filter, LogicalPlan, Operator, ScalarFunctionArgs, ScalarUDF, - ScalarUDFImpl, Signature, SimpleAggregateUDF, Subquery, Union, Volatility, cast, - col, create_udaf, is_true, lit, + ExprSchemable, Filter, Limit, LogicalPlan, Operator, ScalarFunctionArgs, + ScalarUDF, ScalarUDFImpl, Signature, SimpleAggregateUDF, Subquery, Union, + Volatility, cast, col, create_udaf, is_true, lit, placeholder, }; use datafusion_functions_aggregate::average::AvgAccumulator; @@ -1733,6 +1752,19 @@ mod test { Ok(()) } + #[test] + fn untyped_limit_parameter_remains_inferable_after_coercion() -> Result<()> { + let plan = LogicalPlan::Limit(Limit { + input: empty(), + fetch: Some(Box::new(placeholder("$1"))), + skip: None, + }); + let analyzed = TypeCoercion::new().analyze(plan, &ConfigOptions::default())?; + + assert_eq!(analyzed.get_parameter_types()?["$1"], Some(DataType::Int64)); + Ok(()) + } + #[test] fn simple_case() -> Result<()> { let expr = col("a").lt(lit(2_u32)); diff --git a/datafusion/optimizer/tests/optimizer_integration.rs b/datafusion/optimizer/tests/optimizer_integration.rs index 26b48c5e1f352..eb4088e5e31b4 100644 --- a/datafusion/optimizer/tests/optimizer_integration.rs +++ b/datafusion/optimizer/tests/optimizer_integration.rs @@ -670,6 +670,75 @@ fn optimize_plan(plan: LogicalPlan) -> Result { optimizer.optimize(plan, &config, observe) } +#[test] +fn limit_offset_parameters_keep_type_after_optimization() -> Result<()> { + for sql in [ + "SELECT 20 AS value LIMIT $1", + "SELECT 20 AS value OFFSET $1", + ] { + let optimized = test_sql(sql)?; + assert_eq!( + optimized.get_parameter_types()?, + HashMap::from([("$1".to_string(), Some(DataType::Int64))]), + "{sql}" + ); + } + Ok(()) +} + +#[test] +fn limit_offset_parameters_leave_cast_inputs_unresolved_after_optimization() -> Result<()> +{ + for sql in [ + "SELECT $1 AS value", + "SELECT $1 AS value LIMIT CAST($1 AS INT)", + "SELECT $1 AS value LIMIT CAST($1 AS BIGINT)", + "SELECT $1 AS value FROM (VALUES (1), (2)) AS t(v) OFFSET CAST($1 AS BIGINT)", + "SELECT $1 AS value FROM (SELECT 1 LIMIT CAST($1 AS BIGINT)) AS t", + ] { + let optimized = test_sql(sql)?; + assert_eq!( + optimized.get_parameter_types()?, + HashMap::from([("$1".to_string(), None)]), + "{sql}" + ); + } + Ok(()) +} + +#[test] +fn limit_parameter_binding_with_inferred_type() -> Result<()> { + let statements = Parser::parse_sql( + &GenericDialect {}, + "SELECT $1 + CAST(1 AS INT) AS value LIMIT $1", + )?; + let context_provider = MyContextProvider::default(); + let plan = + SqlToRel::new(&context_provider).sql_statement_to_plan(statements[0].clone())?; + let analyzed = + Analyzer::new().execute_and_check(plan, &ConfigOptions::default(), |_, _| {})?; + let bound = analyzed.with_param_values(vec![ScalarValue::Int32(Some(1))])?; + let optimized = optimize_plan(bound)?; + let limit = match &optimized { + LogicalPlan::Limit(limit) => limit, + LogicalPlan::Projection(projection) => { + let LogicalPlan::Limit(limit) = projection.input.as_ref() else { + panic!("expected LIMIT under projection, got {optimized:?}"); + }; + limit + } + _ => panic!("expected LIMIT plan, got {optimized:?}"), + }; + assert!( + matches!( + limit.get_fetch_type()?, + datafusion_expr::FetchType::Literal(Some(1)) + ), + "unexpected optimized LIMIT: {optimized}" + ); + Ok(()) +} + /// Extension node that does NOT implement `necessary_children_exprs`. /// Used to test that the optimizer still processes subtrees below such nodes. #[derive(Debug, Hash, PartialEq, Eq)] diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 10420b9a61f02..ed6eeab28445c 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -880,16 +880,18 @@ impl SqlToRel<'_, S> { if fields.is_empty() { let map_types = plan.get_parameter_fields()?; - let param_types: Vec<_> = (1..=map_types.len()) - .filter_map(|i| { + let param_types: Option> = (1..=map_types.len()) + .map(|i| { let key = format!("${i}"); map_types.get(&key).and_then(|opt| opt.clone()) }) .collect(); - fields.extend(param_types.iter().cloned()); - planner_context.with_prepare_param_data_types( - param_types.into_iter().map(Some).collect(), - ); + if let Some(param_types) = param_types { + fields.extend(param_types.iter().cloned()); + planner_context.with_prepare_param_data_types( + param_types.into_iter().map(Some).collect(), + ); + } } Ok(LogicalPlan::Statement(PlanStatement::Prepare(Prepare { diff --git a/datafusion/sql/tests/cases/params.rs b/datafusion/sql/tests/cases/params.rs index 3d51974071eb6..4a1c71b79a8c3 100644 --- a/datafusion/sql/tests/cases/params.rs +++ b/datafusion/sql/tests/cases/params.rs @@ -987,32 +987,122 @@ fn test_prepare_statement_to_plan_having() { #[test] fn test_prepare_statement_to_plan_limit() { - let sql = "PREPARE my_plan(BIGINT, BIGINT) AS - SELECT id FROM person \ - OFFSET $1 LIMIT $2"; - let (plan, dt) = generate_prepare_stmt_and_data_types(sql); - assert_snapshot!( - plan, - @r#" + for sql in [ + "PREPARE my_plan(BIGINT, BIGINT) AS SELECT id FROM person OFFSET $1 LIMIT $2", + "PREPARE my_plan AS SELECT id FROM person OFFSET $1 LIMIT $2", + ] { + let test = ParameterTest { + sql, + expected_types: vec![ + ("$1", Some(DataType::Int64)), + ("$2", Some(DataType::Int64)), + ], + param_values: vec![ + ScalarValue::Int64(Some(10)), + ScalarValue::Int64(Some(200)), + ], + }; + insta::allow_duplicates! { + assert_snapshot!(test.run(), @r#" + ** Initial Plan: Prepare: "my_plan" [Int64, Int64] Limit: skip=$1, fetch=$2 Projection: person.id TableScan: person - "# - ); - assert_snapshot!(dt, @"Int64, Int64"); - - // replace params with values - let param_values = vec![ScalarValue::Int64(Some(10)), ScalarValue::Int64(Some(200))]; - let plan_with_params = plan.with_param_values(param_values).unwrap(); - assert_snapshot!( - plan_with_params, - @r" + ** Final Plan: Limit: skip=10, fetch=200 Projection: person.id TableScan: person - " - ); + "#); + } + } +} + +#[test] +fn test_limit_offset_parameters() { + let mut plans = Vec::new(); + for sql in [ + "SELECT 20 AS value LIMIT $1", + "SELECT 20 AS value OFFSET $1", + "SELECT $1 AS value LIMIT $1", + "SELECT $1 AS value FROM (VALUES (0), (0)) AS t(dummy) OFFSET $1", + "SELECT value FROM (SELECT 20 AS value LIMIT $1) AS t", + "SELECT (SELECT value FROM (VALUES (10), (20)) AS t(value) \ + ORDER BY value LIMIT $1) AS value", + "SELECT 20 AS value WHERE EXISTS \ + (SELECT value FROM (VALUES (10), (20)) AS t(value) OFFSET $1)", + ] { + let test = ParameterTest { + sql, + expected_types: vec![("$1", Some(DataType::Int64))], + param_values: vec![ScalarValue::Int64(Some(1))], + }; + plans.push(format!("{sql}\n{}", test.run())); + } + + let test = ParameterTest { + sql: "SELECT $1 + CAST(1 AS INT) AS value \ + FROM (SELECT $1) AS t LIMIT $1", + expected_types: vec![("$1", Some(DataType::Int32))], + param_values: vec![ScalarValue::Int32(Some(1))], + }; + plans.push(test.run()); + assert_snapshot!(plans.join("\n")); +} + +#[test] +fn test_limit_offset_parameters_leave_cast_inputs_unresolved() { + let mut plans = Vec::new(); + for sql in [ + "SELECT $1 AS value", + "SELECT $1 AS value LIMIT CAST($1 AS INT)", + "SELECT $1 AS value LIMIT CAST($1 AS BIGINT)", + "SELECT $1 AS value FROM (VALUES (1), (2)) AS t(v) \ + OFFSET CAST($1 AS BIGINT)", + "SELECT $1 AS value FROM (SELECT 1 LIMIT CAST($1 AS BIGINT)) AS t", + ] { + let test = ParameterTest { + sql, + expected_types: vec![("$1", None)], + param_values: vec![ScalarValue::from("1")], + }; + plans.push(format!("{sql}\n{}", test.run())); + } + assert_snapshot!(plans.join("\n")); +} + +#[test] +fn test_limit_offset_parameters_keep_field_metadata() { + let field = Field::new("value", DataType::Utf8, true) + .with_metadata(Metadata::new().with("ARROW:extension:name", "example.string")); + let mut plans = Vec::new(); + for clause in ["LIMIT", "OFFSET"] { + let sql = format!( + "SELECT $1 AS value FROM string_with_extension WHERE value = $1 {clause} $1" + ); + let test = ParameterTestWithMetadata { + sql: &sql, + expected_types: vec![("$1", Some(field.clone().into()))], + param_values: vec![ScalarAndMetadata::new( + ScalarValue::from("1"), + Some(field.metadata().into()), + )], + }; + plans.push(test.run()); + } + assert_snapshot!(plans.join("\n")); +} + +#[test] +fn test_limit_offset_parameters_reject_type_conflicts() { + let plan = logical_plan( + "SELECT $1 + CAST(1 AS INT) AS a, $1 + CAST(1 AS DOUBLE) AS b LIMIT $1", + ) + .unwrap(); + let error = plan + .get_parameter_fields() + .expect_err("conflicting parameter types must fail"); + assert_contains!(error.to_string(), "Conflicting types for id $1"); } #[test] diff --git a/datafusion/sql/tests/cases/snapshots/sql_integration__cases__params__limit_offset_parameters.snap b/datafusion/sql/tests/cases/snapshots/sql_integration__cases__params__limit_offset_parameters.snap new file mode 100644 index 0000000000000..91c91d481819f --- /dev/null +++ b/datafusion/sql/tests/cases/snapshots/sql_integration__cases__params__limit_offset_parameters.snap @@ -0,0 +1,111 @@ +--- +source: datafusion/sql/tests/cases/params.rs +expression: "plans.join(\"\\n\")" +--- +SELECT 20 AS value LIMIT $1 +** Initial Plan: +Limit: skip=0, fetch=$1 + Projection: Int64(20) AS value + EmptyRelation: rows=1 +** Final Plan: +Limit: skip=0, fetch=1 + Projection: Int64(20) AS value + EmptyRelation: rows=1 +SELECT 20 AS value OFFSET $1 +** Initial Plan: +Limit: skip=$1, fetch=None + Projection: Int64(20) AS value + EmptyRelation: rows=1 +** Final Plan: +Limit: skip=1, fetch=None + Projection: Int64(20) AS value + EmptyRelation: rows=1 +SELECT $1 AS value LIMIT $1 +** Initial Plan: +Limit: skip=0, fetch=$1 + Projection: $1 AS value + EmptyRelation: rows=1 +** Final Plan: +Limit: skip=0, fetch=1 + Projection: Int64(1) AS value + EmptyRelation: rows=1 +SELECT $1 AS value FROM (VALUES (0), (0)) AS t(dummy) OFFSET $1 +** Initial Plan: +Limit: skip=$1, fetch=None + Projection: $1 AS value + SubqueryAlias: t + Projection: column1 AS dummy + Values: (Int64(0)), (Int64(0)) +** Final Plan: +Limit: skip=1, fetch=None + Projection: Int64(1) AS value + SubqueryAlias: t + Projection: column1 AS dummy + Values: (Int64(0)), (Int64(0)) +SELECT value FROM (SELECT 20 AS value LIMIT $1) AS t +** Initial Plan: +Projection: t.value + SubqueryAlias: t + Limit: skip=0, fetch=$1 + Projection: Int64(20) AS value + EmptyRelation: rows=1 +** Final Plan: +Projection: t.value + SubqueryAlias: t + Limit: skip=0, fetch=1 + Projection: Int64(20) AS value + EmptyRelation: rows=1 +SELECT (SELECT value FROM (VALUES (10), (20)) AS t(value) ORDER BY value LIMIT $1) AS value +** Initial Plan: +Projection: () AS value + Subquery: + Limit: skip=0, fetch=$1 + Sort: t.value ASC NULLS LAST + Projection: t.value + SubqueryAlias: t + Projection: column1 AS value + Values: (Int64(10)), (Int64(20)) + EmptyRelation: rows=1 +** Final Plan: +Projection: () AS value + Subquery: + Limit: skip=0, fetch=1 + Sort: t.value ASC NULLS LAST + Projection: t.value + SubqueryAlias: t + Projection: column1 AS value + Values: (Int64(10)), (Int64(20)) + EmptyRelation: rows=1 +SELECT 20 AS value WHERE EXISTS (SELECT value FROM (VALUES (10), (20)) AS t(value) OFFSET $1) +** Initial Plan: +Projection: Int64(20) AS value + Filter: EXISTS () + Subquery: + Limit: skip=$1, fetch=None + Projection: t.value + SubqueryAlias: t + Projection: column1 AS value + Values: (Int64(10)), (Int64(20)) + EmptyRelation: rows=1 +** Final Plan: +Projection: Int64(20) AS value + Filter: EXISTS () + Subquery: + Limit: skip=1, fetch=None + Projection: t.value + SubqueryAlias: t + Projection: column1 AS value + Values: (Int64(10)), (Int64(20)) + EmptyRelation: rows=1 +** Initial Plan: +Limit: skip=0, fetch=$1 + Projection: $1 + CAST(Int64(1) AS Int32) AS value + SubqueryAlias: t + Projection: $1 + EmptyRelation: rows=1 +** Final Plan: +Limit: skip=0, fetch=Int32(1) + Projection: Int32(1) + CAST(Int64(1) AS Int32) AS value + SubqueryAlias: t + Projection: Int32(1) AS $1 + EmptyRelation: rows=1 diff --git a/datafusion/sql/tests/cases/snapshots/sql_integration__cases__params__limit_offset_parameters_keep_field_metadata.snap b/datafusion/sql/tests/cases/snapshots/sql_integration__cases__params__limit_offset_parameters_keep_field_metadata.snap new file mode 100644 index 0000000000000..3be40e4451025 --- /dev/null +++ b/datafusion/sql/tests/cases/snapshots/sql_integration__cases__params__limit_offset_parameters_keep_field_metadata.snap @@ -0,0 +1,24 @@ +--- +source: datafusion/sql/tests/cases/params.rs +expression: "plans.join(\"\\n\")" +--- +** Initial Plan: +Limit: skip=0, fetch=$1 + Projection: $1 AS value + Filter: string_with_extension.value = $1 + TableScan: string_with_extension +** Final Plan: +Limit: skip=0, fetch=Utf8("1") FieldMetadata { inner: {"ARROW:extension:name": "example.string"} } + Projection: Utf8("1") FieldMetadata { inner: {"ARROW:extension:name": "example.string"} } AS value + Filter: string_with_extension.value = Utf8("1") FieldMetadata { inner: {"ARROW:extension:name": "example.string"} } + TableScan: string_with_extension +** Initial Plan: +Limit: skip=$1, fetch=None + Projection: $1 AS value + Filter: string_with_extension.value = $1 + TableScan: string_with_extension +** Final Plan: +Limit: skip=Utf8("1") FieldMetadata { inner: {"ARROW:extension:name": "example.string"} }, fetch=None + Projection: Utf8("1") FieldMetadata { inner: {"ARROW:extension:name": "example.string"} } AS value + Filter: string_with_extension.value = Utf8("1") FieldMetadata { inner: {"ARROW:extension:name": "example.string"} } + TableScan: string_with_extension diff --git a/datafusion/sql/tests/cases/snapshots/sql_integration__cases__params__limit_offset_parameters_leave_cast_inputs_unresolved.snap b/datafusion/sql/tests/cases/snapshots/sql_integration__cases__params__limit_offset_parameters_leave_cast_inputs_unresolved.snap new file mode 100644 index 0000000000000..86ab25e91d0fc --- /dev/null +++ b/datafusion/sql/tests/cases/snapshots/sql_integration__cases__params__limit_offset_parameters_leave_cast_inputs_unresolved.snap @@ -0,0 +1,55 @@ +--- +source: datafusion/sql/tests/cases/params.rs +expression: "plans.join(\"\\n\")" +--- +SELECT $1 AS value +** Initial Plan: +Projection: $1 AS value + EmptyRelation: rows=1 +** Final Plan: +Projection: Utf8("1") AS value + EmptyRelation: rows=1 +SELECT $1 AS value LIMIT CAST($1 AS INT) +** Initial Plan: +Limit: skip=0, fetch=CAST($1 AS Int32) + Projection: $1 AS value + EmptyRelation: rows=1 +** Final Plan: +Limit: skip=0, fetch=CAST(Utf8("1") AS Int32) + Projection: Utf8("1") AS value + EmptyRelation: rows=1 +SELECT $1 AS value LIMIT CAST($1 AS BIGINT) +** Initial Plan: +Limit: skip=0, fetch=CAST($1 AS Int64) + Projection: $1 AS value + EmptyRelation: rows=1 +** Final Plan: +Limit: skip=0, fetch=CAST(Utf8("1") AS Int64) + Projection: Utf8("1") AS value + EmptyRelation: rows=1 +SELECT $1 AS value FROM (VALUES (1), (2)) AS t(v) OFFSET CAST($1 AS BIGINT) +** Initial Plan: +Limit: skip=CAST($1 AS Int64), fetch=None + Projection: $1 AS value + SubqueryAlias: t + Projection: column1 AS v + Values: (Int64(1)), (Int64(2)) +** Final Plan: +Limit: skip=CAST(Utf8("1") AS Int64), fetch=None + Projection: Utf8("1") AS value + SubqueryAlias: t + Projection: column1 AS v + Values: (Int64(1)), (Int64(2)) +SELECT $1 AS value FROM (SELECT 1 LIMIT CAST($1 AS BIGINT)) AS t +** Initial Plan: +Projection: $1 AS value + SubqueryAlias: t + Limit: skip=0, fetch=CAST($1 AS Int64) + Projection: Int64(1) + EmptyRelation: rows=1 +** Final Plan: +Projection: Utf8("1") AS value + SubqueryAlias: t + Limit: skip=0, fetch=CAST(Utf8("1") AS Int64) + Projection: Int64(1) + EmptyRelation: rows=1 diff --git a/datafusion/sqllogictest/test_files/prepare.slt b/datafusion/sqllogictest/test_files/prepare.slt index 108b023ccf100..00abf742b7f95 100644 --- a/datafusion/sqllogictest/test_files/prepare.slt +++ b/datafusion/sqllogictest/test_files/prepare.slt @@ -19,6 +19,182 @@ ## Prepare Statement Tests ########## +statement ok +PREPARE page AS SELECT value FROM (VALUES (10), (20), (30)) AS t(value) ORDER BY value LIMIT $1 OFFSET $2; + +query I +EXECUTE page(1, 1); +---- +20 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT 20 AS value LIMIT $1; + +query I +EXECUTE page(1); +---- +20 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT value FROM (VALUES (10), (20)) AS t(value) ORDER BY value OFFSET $1; + +query I +EXECUTE page(1); +---- +20 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT $1 AS value LIMIT $1; + +query I +EXECUTE page(1); +---- +1 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT $1 AS value FROM (VALUES (0), (0)) AS t(dummy) OFFSET $1; + +query I +EXECUTE page(1); +---- +1 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT value FROM (SELECT 20 AS value LIMIT $1) AS t; + +query I +EXECUTE page(1); +---- +20 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT (SELECT value FROM (VALUES (10), (20)) AS t(value) ORDER BY value LIMIT $1) AS value; + +query I +EXECUTE page(1); +---- +10 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT 20 AS value WHERE EXISTS (SELECT value FROM (VALUES (10), (20)) AS t(value) OFFSET $1); + +query I +EXECUTE page(1); +---- +20 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT $1 AS value LIMIT CAST($1 AS INT); + +query T +EXECUTE page('1'); +---- +1 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT $1 AS value LIMIT CAST($1 AS BIGINT); + +query T +EXECUTE page('1'); +---- +1 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT $1 AS value FROM (VALUES (1), (2)) AS t(v) OFFSET CAST($1 AS BIGINT); + +query T +EXECUTE page('1'); +---- +1 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT $1 AS value FROM (SELECT 1 LIMIT CAST($1 AS BIGINT)) AS t; + +query T +EXECUTE page('1'); +---- +1 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT $1 + CAST(1 AS INT) AS value FROM (SELECT $1) AS t LIMIT $1; + +query I +EXECUTE page(1); +---- +2 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT $1 + CAST(1 AS INT) AS value LIMIT $1; + +query I +EXECUTE page(1); +---- +2 + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT $1 AS value LIMIT $2; + +query T +EXECUTE page('row', 1); +---- +row + +statement ok +DEALLOCATE page; + +statement ok +PREPARE page AS SELECT $1 AS value FROM (VALUES (0), (0)) AS t(dummy) OFFSET $2; + +query T +EXECUTE page('row', 1); +---- +row + +statement ok +DEALLOCATE page; + statement ok create table person (id int, first_name varchar, last_name varchar, age int, state varchar, salary double, birthday timestamp, "😀" int) as values (1, 'jane', 'smith', 20, 'MA', 100000.45, '2000-11-12T00:00:00'::timestamp, 99); @@ -379,6 +555,28 @@ EXECUTE get_N_rand_ints_from_last_run(2); statement ok DEALLOCATE get_N_rand_ints_from_last_run; +statement ok +PREPARE limit_after_typed_filter AS SELECT id FROM test WHERE id = $1 LIMIT $2; + +query I +EXECUTE limit_after_typed_filter(1, 1); +---- +1 + +statement ok +DEALLOCATE limit_after_typed_filter; + +statement ok +PREPARE limit_after_inferred_int32 AS SELECT $1 + CAST(1 AS INT) AS value LIMIT $1; + +query I +EXECUTE limit_after_inferred_int32(1); +---- +2 + +statement ok +DEALLOCATE limit_after_inferred_int32; + statement ok DROP TABLE test;