[DF] WIP: Upgrade to DataFusion rev 8ea59a

dask_planner/Cargo.toml

@@ -9,12 +9,12 @@ edition = "2021"
 rust-version = "1.62"
 
 [dependencies]
-arrow = { version = "22.0.0", features = ["prettyprint"] }
+arrow = { version = "23.0.0", features = ["prettyprint"] }
 async-trait = "0.1.41"
-datafusion-common = "12.0.0"
-datafusion-expr = "12.0.0"
-datafusion-optimizer = "12.0.0"
-datafusion-sql = "12.0.0"
+datafusion-common = { git = "https://github.com/apache/arrow-datafusion/", rev = "1261741af2a5e142fa0c7916e759859cc18ea59a" }
+datafusion-expr = { git = "https://github.com/apache/arrow-datafusion/", rev = "1261741af2a5e142fa0c7916e759859cc18ea59a" }
+datafusion-optimizer = { git = "https://github.com/apache/arrow-datafusion/", rev = "1261741af2a5e142fa0c7916e759859cc18ea59a" }
+datafusion-sql = { git = "https://github.com/apache/arrow-datafusion/", rev = "1261741af2a5e142fa0c7916e759859cc18ea59a" }
 env_logger = "0.9"
 log = "^0.4"
 mimalloc = { version = "*", default-features = false }

dask_planner/src/expression.rs

@@ -529,152 +529,211 @@ impl PyExpr {
     /// TODO: I can't express how much I dislike explicity listing all of these methods out
     /// but PyO3 makes it necessary since its annotations cannot be used in trait impl blocks
     #[pyo3(name = "getFloat32Value")]
-    pub fn float_32_value(&mut self) -> PyResult<f32> {
+    pub fn float_32_value(&mut self) -> PyResult<Option<f32>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::Float32(iv) => Ok(iv.unwrap()),
+                ScalarValue::Float32(iv) => Ok(*iv),
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getFloat64Value")]
-    pub fn float_64_value(&mut self) -> PyResult<f64> {
+    pub fn float_64_value(&mut self) -> PyResult<Option<f64>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::Float64(iv) => Ok(iv.unwrap()),
+                ScalarValue::Float64(iv) => Ok(*iv),
+                _ => Err(py_type_err("getValue<T>() - Unexpected value")),
+            },
+            _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
+        }
+    }
+
+    #[pyo3(name = "getDecimal128Value")]
+    pub fn decimal_128_value(&mut self) -> PyResult<(Option<i128>, u8, u8)> {
+        match &self.expr {
+            Expr::Literal(scalar_value) => match scalar_value {
+                ScalarValue::Decimal128(value, precision, scale) => {
+                    Ok((*value, *precision, *scale))
+                }
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getInt8Value")]
-    pub fn int_8_value(&mut self) -> PyResult<i8> {
+    pub fn int_8_value(&mut self) -> PyResult<Option<i8>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::Int8(iv) => Ok(iv.unwrap()),
+                ScalarValue::Int8(iv) => Ok(*iv),
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getInt16Value")]
-    pub fn int_16_value(&mut self) -> PyResult<i16> {
+    pub fn int_16_value(&mut self) -> PyResult<Option<i16>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::Int16(iv) => Ok(iv.unwrap()),
+                ScalarValue::Int16(iv) => Ok(*iv),
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getInt32Value")]
-    pub fn int_32_value(&mut self) -> PyResult<i32> {
+    pub fn int_32_value(&mut self) -> PyResult<Option<i32>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::Int32(iv) => Ok(iv.unwrap()),
+                ScalarValue::Int32(iv) => Ok(*iv),
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getInt64Value")]
-    pub fn int_64_value(&mut self) -> PyResult<i64> {
+    pub fn int_64_value(&mut self) -> PyResult<Option<i64>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::Int64(iv) => Ok(iv.unwrap()),
+                ScalarValue::Int64(iv) => Ok(*iv),
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getUInt8Value")]
-    pub fn uint_8_value(&mut self) -> PyResult<u8> {
+    pub fn uint_8_value(&mut self) -> PyResult<Option<u8>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::UInt8(iv) => Ok(iv.unwrap()),
+                ScalarValue::UInt8(iv) => Ok(*iv),
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getUInt16Value")]
-    pub fn uint_16_value(&mut self) -> PyResult<u16> {
+    pub fn uint_16_value(&mut self) -> PyResult<Option<u16>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::UInt16(iv) => Ok(iv.unwrap()),
+                ScalarValue::UInt16(iv) => Ok(*iv),
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getUInt32Value")]
-    pub fn uint_32_value(&mut self) -> PyResult<u32> {
+    pub fn uint_32_value(&mut self) -> PyResult<Option<u32>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::UInt32(iv) => Ok(iv.unwrap()),
+                ScalarValue::UInt32(iv) => Ok(*iv),
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getUInt64Value")]
-    pub fn uint_64_value(&mut self) -> PyResult<u64> {
+    pub fn uint_64_value(&mut self) -> PyResult<Option<u64>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::UInt64(iv) => Ok(iv.unwrap()),
+                ScalarValue::UInt64(iv) => Ok(*iv),
+                _ => Err(py_type_err("getValue<T>() - Unexpected value")),
+            },
+            _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
+        }
+    }
+
+    #[pyo3(name = "getDate32Value")]
+    pub fn date_32_value(&mut self) -> PyResult<Option<i32>> {
+        match &self.expr {
+            Expr::Literal(scalar_value) => match scalar_value {
+                ScalarValue::Date32(iv) => Ok(*iv),
+                _ => Err(py_type_err("getValue<T>() - Unexpected value")),
+            },
+            _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
+        }
+    }
+
+    #[pyo3(name = "getDate64Value")]
+    pub fn date_64_value(&mut self) -> PyResult<Option<i64>> {
+        match &self.expr {
+            Expr::Literal(scalar_value) => match scalar_value {
+                ScalarValue::Date64(iv) => Ok(*iv),
+                _ => Err(py_type_err("getValue<T>() - Unexpected value")),
+            },
+            _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
+        }
+    }
+
+    #[pyo3(name = "getTime64Value")]
+    pub fn time_64_value(&mut self) -> PyResult<Option<i64>> {
+        match &self.expr {
+            Expr::Literal(scalar_value) => match scalar_value {
+                ScalarValue::Time64(iv) => Ok(*iv),
+                _ => Err(py_type_err("getValue<T>() - Unexpected value")),
+            },
+            _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
+        }
+    }
+
+    #[pyo3(name = "getTimestampValue")]
+    pub fn timestamp_value(&mut self) -> PyResult<(Option<i64>, Option<String>)> {
+        match &self.expr {
+            Expr::Literal(scalar_value) => match scalar_value {
+                ScalarValue::TimestampNanosecond(iv, tz)
+                | ScalarValue::TimestampMicrosecond(iv, tz)
+                | ScalarValue::TimestampMillisecond(iv, tz)
+                | ScalarValue::TimestampSecond(iv, tz) => Ok((*iv, tz.clone())),
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getBoolValue")]
-    pub fn bool_value(&mut self) -> PyResult<bool> {
+    pub fn bool_value(&mut self) -> PyResult<Option<bool>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::Boolean(Some(iv)) => Ok(*iv),
+                ScalarValue::Boolean(iv) => Ok(*iv),
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getStringValue")]
-    pub fn string_value(&mut self) -> PyResult<String> {
+    pub fn string_value(&mut self) -> PyResult<Option<String>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::Utf8(iv) => Ok(iv.clone().unwrap()),
+                ScalarValue::Utf8(iv) => Ok(iv.clone()),
                 _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
             _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }
 
     #[pyo3(name = "getIntervalDayTimeValue")]
-    pub fn interval_day_time_value(&mut self) -> (i32, i32) {
+    pub fn interval_day_time_value(&mut self) -> PyResult<Option<(i32, i32)>> {
         match &self.expr {
             Expr::Literal(scalar_value) => match scalar_value {
-                ScalarValue::IntervalDayTime(iv) => {
-                    let interval = iv.unwrap() as u64;
+                ScalarValue::IntervalDayTime(Some(iv)) => {
+                    let interval = *iv as u64;
                     let days = (interval >> 32) as i32;
                     let ms = interval as i32;
-                    (days, ms)
-                }
-                _ => {
-                    panic!("getValue<T>() - Unexpected value")
+                    Ok(Some((days, ms)))
                 }
+                ScalarValue::IntervalDayTime(None) => Ok(None),
+                _ => Err(py_type_err("getValue<T>() - Unexpected value")),
             },
-            _ => panic!("getValue<T>() - Non literal value encountered"),
+            _ => Err(py_type_err("getValue<T>() - Non literal value encountered")),
         }
     }