[arrow-datafusion] branch master updated: Use .get() to avoid panic (#3201)

[yj...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

yjshen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new 82da46d77 Use .get() to avoid panic (#3201)
82da46d77 is described below

commit 82da46d7725d6410074de7dd53b37c5f7df6a4c2
Author: Jack Klamer <jf...@gmail.com>
AuthorDate: Wed Aug 24 22:47:49 2022 -0500

    Use .get() to avoid panic (#3201)
    
    * Use .get() to avoid panic
    
    * Refactor system var detection into common
    
    * Clean up other known index issue
    
    * PR comments, and refactorings
---
 datafusion/core/src/execution/context.rs     |  4 ++--
 datafusion/physical-expr/src/planner.rs      |  3 ++-
 datafusion/physical-expr/src/var_provider.rs | 25 +++++++++++++++++++++++++
 datafusion/sql/src/planner.rs                |  2 +-
 4 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/datafusion/core/src/execution/context.rs b/datafusion/core/src/execution/context.rs
index a00babfe4..a0dee62de 100644
--- a/datafusion/core/src/execution/context.rs
+++ b/datafusion/core/src/execution/context.rs
@@ -43,6 +43,7 @@ use crate::{
     },
 };
 pub use datafusion_physical_expr::execution_props::ExecutionProps;
+use datafusion_physical_expr::var_provider::is_system_variables;
 use parking_lot::RwLock;
 use std::sync::Arc;
 use std::{
@@ -1563,8 +1564,7 @@ impl ContextProvider for SessionState {
             return None;
         }
 
-        let first_variable = &variable_names[0];
-        let provider_type = if first_variable.len() > 1 && &first_variable[0..2] == "@@" {
+        let provider_type = if is_system_variables(variable_names) {
             VarType::System
         } else {
             VarType::UserDefined
diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs
index 44fe0595a..f84a5fbb9 100644
--- a/datafusion/physical-expr/src/planner.rs
+++ b/datafusion/physical-expr/src/planner.rs
@@ -16,6 +16,7 @@
 // under the License.
 
 use crate::expressions::try_cast;
+use crate::var_provider::is_system_variables;
 use crate::{
     execution_props::ExecutionProps,
     expressions::{
@@ -51,7 +52,7 @@ pub fn create_physical_expr(
         }
         Expr::Literal(value) => Ok(Arc::new(Literal::new(value.clone()))),
         Expr::ScalarVariable(_, variable_names) => {
-            if &variable_names[0][0..2] == "@@" {
+            if is_system_variables(variable_names) {
                 match execution_props.get_var_provider(VarType::System) {
                     Some(provider) => {
                         let scalar_value = provider.get_value(variable_names.clone())?;
diff --git a/datafusion/physical-expr/src/var_provider.rs b/datafusion/physical-expr/src/var_provider.rs
index db68200d4..0afff91fc 100644
--- a/datafusion/physical-expr/src/var_provider.rs
+++ b/datafusion/physical-expr/src/var_provider.rs
@@ -37,3 +37,28 @@ pub trait VarProvider {
     /// Return the type of the given variable
     fn get_type(&self, var_names: &[String]) -> Option<DataType>;
 }
+
+pub fn is_system_variables(variable_names: &[String]) -> bool {
+    !variable_names.is_empty() && variable_names[0].get(0..2) == Some("@@")
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_is_system_variables() {
+        assert!(!is_system_variables(&["N".into(), "\"\"".into()]));
+        assert!(is_system_variables(&["@@D".into(), "F".into(), "J".into()]));
+        assert!(!is_system_variables(&[
+            "J".into(),
+            "@@F".into(),
+            "K".into()
+        ]));
+        assert!(!is_system_variables(&[]));
+        assert!(is_system_variables(&[
+            "@@longvariablenamethatIdontknowhwyanyonewoulduse".into()
+        ]));
+        assert!(!is_system_variables(&["@F".into()]));
+    }
+}
diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs
index 2f5bc3470..094b980f1 100644
--- a/datafusion/sql/src/planner.rs
+++ b/datafusion/sql/src/planner.rs
@@ -1817,7 +1817,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             SQLExpr::CompoundIdentifier(ids) => {
                 let mut var_names: Vec<_> = ids.into_iter().map(|s| normalize_ident(&s)).collect();
 
-                if &var_names[0][0..1] == "@" {
+                if var_names[0].get(0..1) == Some("@") {
                     let ty = self
                         .schema_provider
                         .get_variable_type(&var_names)