kosiew commented on code in PR #25364:
URL: https://github.com/apache/datafusion/pull/25364#discussion_r4193314337
##########
datafusion/core/tests/sql/select.rs:
##########
@@ -387,6 +387,393 @@ 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_optimized_limit_offset_parameters_keep_type() -> Result<()> {
+ let ctx = SessionContext::new();
+
+ for sql in [
+ "SELECT 20 AS value LIMIT $1",
+ "SELECT 20 AS value OFFSET $1",
+ ] {
+ let df = ctx.sql(sql).await?;
+ let optimized = ctx.state().optimize(df.logical_plan())?;
+ assert_eq!(
+ optimized.get_parameter_types()?,
+ HashMap::from([("$1".to_string(), Some(DataType::Int64))]),
+ "{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)",
+ "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 df = ctx.sql(sql).await?;
+ println!("explicit cast SQL: {sql}");
+ assert_eq!(
+ df.logical_plan().get_parameter_types()?,
+ HashMap::from([("$1".to_string(), None)])
+ );
+ let optimized = ctx.state().optimize(df.logical_plan())?;
+ assert_eq!(
+ optimized.get_parameter_types()?,
+ HashMap::from([("$1".to_string(), None)])
+ );
+ println!(
+ "optimized parameters={:?}",
+ optimized.get_parameter_types()?
+ );
+ let results = df
+ .with_param_values(vec![ScalarValue::from("1")])?
+ .collect()
+ .await?;
+ println!("{}", batches_to_sort_string(&results));
+ 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_parameter_binding_with_inferred_type() -> Result<()> {
Review Comment:
There is still quite a bit of parameter-test setup duplicated here. Could
you move the planning and binding assertions into the existing parameter
harnesses where they fit, while keeping the analyzer/optimizer and
execution-specific cases in coverage that actually exercises those phases?
--
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]
