From 35f59271e6ee90cf1522ee2ca173e3bce9af59fd Mon Sep 17 00:00:00 2001
From: Jeremy Dyer <jdye64@gmail.com>
Date: Wed, 25 Oct 2023 14:30:59 -0400
Subject: [PATCH] Resolve syntax issues encountered due to bumping to ADP 32
 from 28

---
 src/expression.rs                              |  4 ++++
 src/sql/logical.rs                             |  1 +
 src/sql/optimizer/decorrelate_where_in.rs      | 13 +++++++++++--
 src/sql/optimizer/dynamic_partition_pruning.rs |  3 +--
 4 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/src/expression.rs b/src/expression.rs
index fccfa9d87..8f4a60935 100644
--- a/src/expression.rs
+++ b/src/expression.rs
@@ -546,12 +546,16 @@ impl PyExpr {
                         "Bitwise operators unsupported in get_type".to_string(),
                     ));
                 }
+                Operator::AtArrow | Operator::ArrowAt => {
+                    todo!()
+                }
             },
             Expr::Literal(scalar_value) => match scalar_value {
                 ScalarValue::Boolean(_value) => "Boolean",
                 ScalarValue::Float32(_value) => "Float32",
                 ScalarValue::Float64(_value) => "Float64",
                 ScalarValue::Decimal128(_value, ..) => "Decimal128",
+                ScalarValue::Decimal256(_, _, _) => "Decimal256",
                 ScalarValue::Dictionary(..) => "Dictionary",
                 ScalarValue::Int8(_value) => "Int8",
                 ScalarValue::Int16(_value) => "Int16",
diff --git a/src/sql/logical.rs b/src/sql/logical.rs
index e8f5f9f6f..9ca3278a0 100644
--- a/src/sql/logical.rs
+++ b/src/sql/logical.rs
@@ -372,6 +372,7 @@ impl PyLogicalPlan {
                 }
             }
             LogicalPlan::Unnest(_unnest) => "Unnest",
+            LogicalPlan::Copy(_) => "Copy",
         })
     }
 
diff --git a/src/sql/optimizer/decorrelate_where_in.rs b/src/sql/optimizer/decorrelate_where_in.rs
index 014f22092..cb4e3736c 100644
--- a/src/sql/optimizer/decorrelate_where_in.rs
+++ b/src/sql/optimizer/decorrelate_where_in.rs
@@ -18,7 +18,7 @@
 use std::sync::Arc;
 
 use datafusion_python::{
-    datafusion_common::{alias::AliasGenerator, context, Column, Result},
+    datafusion_common::{alias::AliasGenerator, context, Column, DataFusionError, Result},
     datafusion_expr::{
         expr::InSubquery,
         expr_rewriter::unnormalize_col,
@@ -155,7 +155,7 @@ fn optimize_where_in(
     left: &LogicalPlan,
     alias: &AliasGenerator,
 ) -> Result<LogicalPlan> {
-    let projection = Projection::try_from_plan(&query_info.query.subquery)
+    let projection = try_from_plan(&query_info.query.subquery)
         .map_err(|e| context!("a projection is required", e))?;
     let subquery_input = projection.input.clone();
     // TODO add the validate logic to Analyzer
@@ -241,6 +241,15 @@ fn remove_duplicated_filter(filters: Vec<Expr>, in_predicate: Expr) -> Vec<Expr>
         .collect::<Vec<_>>()
 }
 
+fn try_from_plan(plan: &LogicalPlan) -> Result<&Projection> {
+    match plan {
+        LogicalPlan::Projection(it) => Ok(it),
+        _ => Err(DataFusionError::Internal(
+            "Could not coerce into Projection!".to_string(),
+        )),
+    }
+}
+
 struct SubqueryInfo {
     query: Subquery,
     where_in_expr: Expr,
diff --git a/src/sql/optimizer/dynamic_partition_pruning.rs b/src/sql/optimizer/dynamic_partition_pruning.rs
index d7e1a8be5..d94359f3a 100644
--- a/src/sql/optimizer/dynamic_partition_pruning.rs
+++ b/src/sql/optimizer/dynamic_partition_pruning.rs
@@ -24,7 +24,6 @@ use datafusion_python::{
     datafusion_expr::{
         expr::InList,
         logical_plan::LogicalPlan,
-        utils::from_plan,
         Expr,
         JoinType,
         Operator,
@@ -1081,7 +1080,7 @@ fn optimize_children(
         new_inputs.push(new_input.unwrap_or_else(|| input.clone()))
     }
     if plan_is_changed {
-        Ok(Some(from_plan(plan, &new_exprs, &new_inputs)?))
+        Ok(Some(plan.with_new_exprs(new_exprs, &new_inputs)?))
     } else {
         Ok(None)
     }