Infer placeholder datatype after LIMIT clause as DataType::Int64 (#81)

UPSTREAM NOTE: Upstream PR has been created but not merged yet. Should be available in DF49
apache#15980

Author: kczimm

datafusion/core/src/datasource/listing/table.rs

@@ -35,7 +35,6 @@ use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuil
 use datafusion_datasource::metadata::MetadataColumn;
 use datafusion_expr::dml::InsertOp;
 use datafusion_expr::{Expr, SortExpr, TableProviderFilterPushDown, TableType};
-use datafusion_physical_expr::create_physical_expr;
 use datafusion_physical_plan::empty::EmptyExec;
 use datafusion_physical_plan::{ExecutionPlan, Statistics};
 

datafusion/core/tests/sql/path_partition.rs

@@ -40,16 +40,15 @@ use datafusion_catalog::memory::DataSourceExec;
 use datafusion_catalog::TableProvider;
 use datafusion_common::stats::Precision;
 use datafusion_common::test_util::batches_to_sort_string;
-use datafusion_common::{Column, ScalarValue};
+use datafusion_common::ScalarValue;
 use datafusion_datasource::file_scan_config::FileScanConfig;
 use datafusion_datasource::metadata::MetadataColumn;
-use datafusion_datasource_parquet::source::ParquetSource;
 use datafusion_execution::config::SessionConfig;
 
 use async_trait::async_trait;
 use bytes::Bytes;
 use chrono::{TimeZone, Utc};
-use datafusion_expr::{col, lit, BinaryExpr, Expr, Operator};
+use datafusion_expr::{col, lit, Expr};
 use futures::stream::{self, BoxStream};
 use insta::assert_snapshot;
 use object_store::{

datafusion/expr/src/expr.rs

@@ -1985,7 +1985,7 @@ impl Expr {
                             &Expr::Column(Column {
                                 relation: None,
                                 name: first_field.name().clone(),
-                                spans: Spans::new(),
+                                spans: Spans::default(),
                             }),
                             schema,
                         )?;
@@ -3525,7 +3525,7 @@ mod test {
         let subquery_filter = Expr::BinaryExpr(BinaryExpr {
             left: Box::new(col("B")),
             op: Operator::Gt,
-            right: Box::new(Expr::Literal(ScalarValue::Int32(Some(3)))),
+            right: Box::new(Expr::Literal(ScalarValue::Int32(Some(3)), None)),
         });
 
         let subquery_scan = LogicalPlan::TableScan(TableScan {

datafusion/expr/src/logical_plan/plan.rs

@@ -1497,6 +1497,14 @@ impl LogicalPlan {
         let mut param_types: HashMap<String, Option<DataType>> = HashMap::new();
 
         self.apply_with_subqueries(|plan| {
+            if let LogicalPlan::Limit(Limit { fetch: Some(e), .. }) = plan {
+                if let Expr::Placeholder(Placeholder { id, data_type }) = &**e {
+                    param_types.insert(
+                        id.clone(),
+                        Some(data_type.clone().unwrap_or(DataType::Int64)),
+                    );
+                }
+            }
             plan.apply_expressions(|expr| {
                 expr.apply(|expr| {
                     if let Expr::Placeholder(Placeholder { id, data_type }) = expr {
@@ -1510,6 +1518,9 @@ impl LogicalPlan {
                             (_, Some(dt)) => {
                                 param_types.insert(id.clone(), Some(dt.clone()));
                             }
+                            (Some(Some(_)), None) => {
+                                // we have already inferred the datatype
+                            }
                             _ => {
                                 param_types.insert(id.clone(), None);
                             }
@@ -5613,6 +5624,42 @@ mod tests {
             "USING join should have all fields"
         );
         assert_eq!(using_join.join_constraint, JoinConstraint::Using);
+        Ok(())
+    }
+
+    #[test]
+    fn test_resolved_placeholder_limit() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("A", DataType::Int32, true)]));
+        let source = Arc::new(LogicalTableSource::new(schema.clone()));
+
+        let placeholder_value = "$1";
+
+        // SELECT * FROM my_table LIMIT $1
+        let plan = LogicalPlan::Limit(Limit {
+            skip: None,
+            fetch: Some(Box::new(Expr::Placeholder(Placeholder {
+                id: placeholder_value.to_string(),
+                data_type: None,
+            }))),
+            input: Arc::new(LogicalPlan::TableScan(TableScan {
+                table_name: TableReference::from("my_table"),
+                source,
+                projected_schema: Arc::new(DFSchema::try_from(schema.clone())?),
+                projection: None,
+                filters: vec![],
+                fetch: None,
+            })),
+        });
+
+        let params = plan.get_parameter_types().expect("to infer type");
+        assert_eq!(params.len(), 1);
+
+        let parameter_type = params
+            .clone()
+            .get(placeholder_value)
+            .expect("to get type")
+            .clone();
+        assert_eq!(parameter_type, Some(DataType::Int64));
 
         Ok(())
     }

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -284,7 +284,7 @@ fn roundtrip_crossjoin() -> Result<()> {
         plan_roundtrip,
         @r"
     Projection: j1.j1_id, j2.j2_string
-      Cross Join:
+      Cross Join: 
         TableScan: j1
         TableScan: j2
     "
@@ -1977,7 +1977,7 @@ fn test_complex_order_by_with_grouping() -> Result<()> {
     }, {
         assert_snapshot!(
             sql,
-            @r#"SELECT j1.j1_id, j1.j1_string, lochierarchy FROM (SELECT j1.j1_id, j1.j1_string, (grouping(j1.j1_id) + grouping(j1.j1_string)) AS lochierarchy, grouping(j1.j1_string), grouping(j1.j1_id) FROM j1 GROUP BY ROLLUP (j1.j1_id, j1.j1_string) ORDER BY (grouping(j1.j1_id) + grouping(j1.j1_string)) DESC NULLS FIRST, CASE WHEN ((grouping(j1.j1_id) + grouping(j1.j1_string)) = 0) THEN j1.j1_id END ASC NULLS LAST) LIMIT 100"#
+            @"SELECT j1_id, j1_string, lochierarchy FROM (SELECT j1.j1_id, j1.j1_string, (grouping(j1.j1_id) + grouping(j1.j1_string)) AS lochierarchy, grouping(j1.j1_string), grouping(j1.j1_id) FROM j1 GROUP BY ROLLUP (j1.j1_id, j1.j1_string) ORDER BY (grouping(j1.j1_id) + grouping(j1.j1_string)) DESC NULLS FIRST, CASE WHEN ((grouping(j1.j1_id) + grouping(j1.j1_string)) = 0) THEN j1.j1_id END ASC NULLS LAST) LIMIT 100"
         );
     });
 
@@ -2682,7 +2682,7 @@ fn test_struct_expr() {
     );
     assert_snapshot!(
         statement,
-        @r#"SELECT test."metadata".product FROM (SELECT {product: {"name": 'Product Name'}} AS "metadata") AS test WHERE (test."metadata".product."name" = 'Product Name')"#
+        @r#"SELECT product FROM (SELECT {product: {"name": 'Product Name'}} AS "metadata") AS test WHERE (test."metadata".product."name" = 'Product Name')"#
     );
 
     let statement = generate_round_trip_statement(
@@ -2691,7 +2691,7 @@ fn test_struct_expr() {
     );
     assert_snapshot!(
         statement,
-        @r#"SELECT test."metadata".product FROM (SELECT {product: {"name": 'Product Name'}} AS "metadata") AS test WHERE (test."metadata".product."name" = 'Product Name')"#
+        @r#"SELECT product FROM (SELECT {product: {"name": 'Product Name'}} AS "metadata") AS test WHERE (test."metadata".product."name" = 'Product Name')"#
     );
 }