jayzhan211 commented on code in PR #25364:
URL: https://github.com/apache/datafusion/pull/25364#discussion_r4170961821
##########
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<HashMap<String, Option<FieldRef>>, DataFusionError> {
let mut param_types: HashMap<String, Option<FieldRef>> =
HashMap::new();
+ let mut row_count_parameters: HashSet<String> = HashSet::new();
self.apply_with_subqueries(|plan| {
plan.apply_expressions(|expr| {
+ if matches!(plan, LogicalPlan::Limit(_))
Review Comment:
The default only applies to a bare placeholder, but
`TypeCoercion::coerce_limit` rewrites it to `CAST($1 AS Int64)`. So once the
plan has been analyzed or optimized, `get_parameter_types()` is back to `None`:
```rs
let df = ctx.sql("SELECT 20 AS value LIMIT $1").await?;
let optimized = ctx.state().optimize(df.logical_plan())?;
// actual: {"$1": None}, plan shows `Limit: skip=0, fetch=CAST($1 AS Int64)`
assert_eq!(
optimized.get_parameter_types()?,
HashMap::from([("$1".to_string(), Some(DataType::Int64))]),
);
```
Leaving untyped placeholders unwrapped in `coerce_limit_expr`
(`datafusion/optimizer/src/analyzer/type_coercion.rs`) fixes it. Locally,
type_coercion tests, `prepare.slt`, `limit.slt` and this PR's tests all pass
with it. Fine to do in a follow-up; if so, please note "unanalyzed plan" in the
doc comment for now.
```diff
) -> Result<Expr> {
+ // Keep untyped placeholders bare so `get_parameter_fields` can
default them
+ if matches!(&expr, Expr::Placeholder(p) if p.field.is_none()) {
+ return Ok(expr);
+ }
let dt = expr.get_type(schema)?;
```
##########
datafusion/core/tests/sql/select.rs:
##########
@@ -387,6 +387,297 @@ async fn test_query_parameters_with_metadata() ->
Result<()> {
Ok(())
}
+#[tokio::test]
+async fn test_limit_offset_parameters() -> Result<()> {
+ let ctx = SessionContext::new();
+
+ for (sql, values, expected_value) in [
+ ("SELECT 20 AS value LIMIT $1", vec![1_i64], 20_i64),
+ (
+ "SELECT value FROM (VALUES (10), (20)) AS t(value) ORDER BY value
OFFSET $1",
+ vec![1],
+ 20,
+ ),
+ (
+ "SELECT value FROM (VALUES (10), (20), (30)) AS t(value) \
+ ORDER BY value LIMIT $1 OFFSET $2",
+ vec![1, 1],
+ 20,
+ ),
+ ("SELECT $1 AS value LIMIT $1", vec![1], 1),
+ (
+ "SELECT $1 AS value FROM (VALUES (0), (0)) AS t(dummy) OFFSET $1",
+ vec![1],
+ 1,
+ ),
+ (
+ "SELECT value FROM (SELECT 20 AS value LIMIT $1) AS t",
+ vec![1],
+ 20,
+ ),
+ (
+ "SELECT (SELECT value FROM (VALUES (10), (20)) AS t(value) \
+ ORDER BY value LIMIT $1) AS value",
+ vec![1],
+ 10,
+ ),
+ (
+ "SELECT 20 AS value WHERE EXISTS \
+ (SELECT value FROM (VALUES (10), (20)) AS t(value) OFFSET $1)",
+ vec![1],
+ 20,
+ ),
+ ] {
+ let df = ctx.sql(sql).await?;
+ let parameter_types = df.logical_plan().get_parameter_types()?;
+ let expected_types: HashMap<String, Option<DataType>> =
(1..=values.len())
+ .map(|index| (format!("${index}"), Some(DataType::Int64)))
+ .collect();
+ let results = df
+ .with_param_values(
+ values
+ .into_iter()
+ .map(ScalarValue::from)
+ .collect::<Vec<_>>(),
+ )?
+ .collect()
+ .await?;
+
+ println!(
+ "{sql}\nparameters={parameter_types:?}\n{}",
+ batches_to_sort_string(&results)
+ );
+ assert_eq!(results.iter().map(RecordBatch::num_rows).sum::<usize>(),
1);
+ assert_eq!(
+ ScalarValue::try_from_array(results[0].column(0), 0)?,
+ ScalarValue::Int64(Some(expected_value))
+ );
+ assert_eq!(parameter_types, expected_types, "{sql}");
+ }
+
+ 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 fields = df.logical_plan().get_parameter_fields()?;
+ let field = fields["$1"].as_ref().unwrap();
+ assert_eq!(field.data_type(), &DataType::Int32);
+ assert_eq!(field.metadata(), &metadata);
+
+ 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_limit_offset_parameters_leave_cast_inputs_unresolved() ->
Result<()> {
+ let ctx = SessionContext::new();
+
+ for sql in [
+ "SELECT $1 AS value",
+ "SELECT $1 AS value LIMIT CAST($1 AS INT)",
+ ] {
+ let df = ctx.sql(sql).await?;
+ assert_eq!(
+ df.logical_plan().get_parameter_types()?,
+ HashMap::from([("$1".to_string(), None)])
+ );
+ let results = df
+ .with_param_values(vec![ScalarValue::from("1")])?
+ .collect()
+ .await?;
+ datafusion::assert_batches_eq!(
+ [
+ "+-------+",
+ "| value |",
+ "+-------+",
+ "| 1 |",
+ "+-------+"
+ ],
+ &results
+ );
+ }
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_limit_offset_parameters_keep_inferred_type() -> Result<()> {
+ let ctx = SessionContext::new();
+ let sql = "SELECT $1 + CAST(1 AS INT) AS value \
+ FROM (SELECT $1) AS t LIMIT $1";
+ let df = ctx.sql(sql).await?;
+ let parameter_types = df.logical_plan().get_parameter_types()?;
+ let results = df
+ .with_param_values(vec![ScalarValue::Int32(Some(1))])?
+ .collect()
+ .await?;
+
+ assert_eq!(
+ parameter_types,
+ HashMap::from([("$1".to_string(), Some(DataType::Int32))])
+ );
+ datafusion::assert_batches_eq!(
+ [
+ "+-------+",
+ "| value |",
+ "+-------+",
+ "| 2 |",
+ "+-------+"
+ ],
+ &results
+ );
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_limit_offset_parameters_reject_type_conflicts() -> Result<()> {
+ let ctx = SessionContext::new();
+ let df = ctx
+ .sql(
+ "SELECT $1 + CAST(1 AS INT) AS a, \
+ $1 + CAST(1 AS DOUBLE) AS b LIMIT $1",
+ )
+ .await?;
+ let error = df
+ .logical_plan()
+ .get_parameter_fields()
+ .expect_err("conflicting parameter types must fail");
+ assert_contains!(error.to_string(), "Conflicting types for id $1");
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_limit_offset_parameters_prepare_execute() -> Result<()> {
Review Comment:
Both PREPARE tests here also EXECUTE fine on main; only the
`get_parameter_types` assertion fails there. The case main actually rejects is
a schema-typed parameter plus a LIMIT parameter: `Prepare specifies 1 data
types but query has 2 parameters`. Please add it to `prepare.slt`, e.g. before
`DROP TABLE test;`:
```sql
# Typed filter parameter combined with an untyped LIMIT parameter
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;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
