datafusion_expr/var_provider.rs
1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements. See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership. The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License. You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied. See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Variable provider
19
20use arrow::datatypes::DataType;
21use datafusion_common::{Result, ScalarValue};
22
23/// Variable type, system/user defined
24#[derive(Debug, Clone, PartialEq, Eq, Hash)]
25pub enum VarType {
26 /// System variable, like @@version
27 System,
28 /// User defined variable, like @name
29 UserDefined,
30}
31
32/// A var provider for `@variable` and `@@variable` runtime values.
33pub trait VarProvider: std::fmt::Debug {
34 /// Get variable value
35 fn get_value(&self, var_names: Vec<String>) -> Result<ScalarValue>;
36
37 /// Return the type of the given variable
38 fn get_type(&self, var_names: &[String]) -> Option<DataType>;
39}
40
41/// Returns true if the specified string is a "system" variable such as
42/// `@@version`
43///
44/// See [`SessionContext::register_variable`] for more details
45///
46/// [`SessionContext::register_variable`]: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html#method.register_variable
47pub fn is_system_variables(variable_names: &[String]) -> bool {
48 !variable_names.is_empty() && variable_names[0].get(0..2) == Some("@@")
49}
50
51#[cfg(test)]
52mod test {
53 use super::*;
54
55 #[test]
56 fn test_is_system_variables() {
57 assert!(!is_system_variables(&["N".into(), "\"\"".into()]));
58 assert!(is_system_variables(&["@@D".into(), "F".into(), "J".into()]));
59 assert!(!is_system_variables(&[
60 "J".into(),
61 "@@F".into(),
62 "K".into()
63 ]));
64 assert!(!is_system_variables(&[]));
65 assert!(is_system_variables(&[
66 "@@longvariablenamethatIdontknowhwyanyonewoulduse".into()
67 ]));
68 assert!(!is_system_variables(&["@F".into()]));
69 }
70}