Make expr member of PyExpr public (apache#375)

jdye64 · web-flow · commit bdcbb9f12e59 · 2023-05-09T06:39:50.000-06:00
diff --git a/src/common/data_type.rs b/src/common/data_type.rs
@@ -21,6 +21,17 @@ use pyo3::prelude::*;
 
 use crate::errors::py_datafusion_err;
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[pyclass(name = "RexType", module = "datafusion.common")]
+pub enum RexType {
+    Alias,
+    Literal,
+    Call,
+    Reference,
+    ScalarSubquery,
+    Other,
+}
+
 /// These bindings are tying together several disparate systems.
 /// You have SQL types for the SQL strings and RDBMS systems itself.
 /// Rust types for the DataFusion code
diff --git a/src/expr.rs b/src/expr.rs
@@ -22,6 +22,7 @@ use datafusion::arrow::datatypes::DataType;
 use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion_expr::{col, lit, Cast, Expr, GetIndexedField};
 
+use crate::common::data_type::RexType;
 use crate::errors::py_runtime_err;
 use crate::expr::aggregate_expr::PyAggregateFunction;
 use crate::expr::binary_expr::PyBinaryExpr;
@@ -83,7 +84,7 @@ pub mod union;
 #[pyclass(name = "Expr", module = "datafusion.expr", subclass)]
 #[derive(Debug, Clone)]
 pub struct PyExpr {
-    pub(crate) expr: Expr,
+    pub expr: Expr,
 }
 
 impl From<PyExpr> for Expr {
@@ -228,6 +229,51 @@ impl PyExpr {
         let expr = Expr::Cast(Cast::new(Box::new(self.expr.clone()), to.0));
         expr.into()
     }
+
+    /// A Rex (Row Expression) specifies a single row of data. That specification
+    /// could include user defined functions or types. RexType identifies the row
+    /// as one of the possible valid `RexTypes`.
+    pub fn rex_type(&self) -> PyResult<RexType> {
+        Ok(match self.expr {
+            Expr::Alias(..) => RexType::Alias,
+            Expr::Column(..) | Expr::QualifiedWildcard { .. } | Expr::GetIndexedField { .. } => {
+                RexType::Reference
+            }
+            Expr::ScalarVariable(..) | Expr::Literal(..) => RexType::Literal,
+            Expr::BinaryExpr { .. }
+            | Expr::Not(..)
+            | Expr::IsNotNull(..)
+            | Expr::Negative(..)
+            | Expr::IsNull(..)
+            | Expr::Like { .. }
+            | Expr::ILike { .. }
+            | Expr::SimilarTo { .. }
+            | Expr::Between { .. }
+            | Expr::Case { .. }
+            | Expr::Cast { .. }
+            | Expr::TryCast { .. }
+            | Expr::Sort { .. }
+            | Expr::ScalarFunction { .. }
+            | Expr::AggregateFunction { .. }
+            | Expr::WindowFunction { .. }
+            | Expr::AggregateUDF { .. }
+            | Expr::InList { .. }
+            | Expr::Wildcard
+            | Expr::ScalarUDF { .. }
+            | Expr::Exists { .. }
+            | Expr::InSubquery { .. }
+            | Expr::GroupingSet(..)
+            | Expr::IsTrue(..)
+            | Expr::IsFalse(..)
+            | Expr::IsUnknown(_)
+            | Expr::IsNotTrue(..)
+            | Expr::IsNotFalse(..)
+            | Expr::Placeholder { .. }
+            | Expr::OuterReferenceColumn(_, _)
+            | Expr::IsNotUnknown(_) => RexType::Call,
+            Expr::ScalarSubquery(..) => RexType::ScalarSubquery,
+        })
+    }
 }
 
 /// Initializes the `expr` module to match the pattern of `datafusion-expr` https://docs.rs/datafusion-expr/latest/datafusion_expr/
