datafusion_functions_nested/
expr_ext.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//! Extension methods for Expr.
19
20use datafusion_expr::Expr;
21
22use crate::extract::{array_element, array_slice};
23
24/// Return access to the element field. Example `expr["name"]`
25///
26/// ## Example Access element 2 from column "c1"
27///
28/// For example if column "c1" holds documents like this
29///
30/// ```json
31/// [10, 20, 30, 40]
32/// ```
33///
34/// You can access the value "30" with
35///
36/// ```
37/// # use datafusion_expr::{lit, col, Expr};
38/// # use datafusion_functions_nested::expr_ext::IndexAccessor;
39/// let expr = col("c1").index(lit(3));
40/// assert_eq!(expr.schema_name().to_string(), "c1[Int32(3)]");
41/// ```
42pub trait IndexAccessor {
43    fn index(self, key: Expr) -> Expr;
44}
45
46impl IndexAccessor for Expr {
47    fn index(self, key: Expr) -> Expr {
48        array_element(self, key)
49    }
50}
51
52/// Return elements between `1` based `start` and `stop`, for
53/// example `expr[1:3]`
54///
55/// ## Example: Access element 2, 3, 4 from column "c1"
56///
57/// For example if column "c1" holds documents like this
58///
59/// ```json
60/// [10, 20, 30, 40]
61/// ```
62///
63/// You can access the value `[20, 30, 40]` with
64///
65/// ```
66/// # use datafusion_expr::{lit, col};
67/// # use datafusion_functions_nested::expr_ext::SliceAccessor;
68/// let expr = col("c1").range(lit(2), lit(4));
69/// assert_eq!(expr.schema_name().to_string(), "c1[Int32(2):Int32(4)]");
70/// ```
71pub trait SliceAccessor {
72    fn range(self, start: Expr, stop: Expr) -> Expr;
73}
74
75impl SliceAccessor for Expr {
76    fn range(self, start: Expr, stop: Expr) -> Expr {
77        array_slice(self, start, stop, None)
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    use datafusion_expr::{col, lit};
86
87    #[test]
88    fn test_index() {
89        let expr1 = col("a").index(lit(1));
90        let expr2 = array_element(col("a"), lit(1));
91        assert_eq!(expr1, expr2);
92    }
93
94    #[test]
95    fn test_range() {
96        let expr1 = col("a").range(lit(1), lit(2));
97        let expr2 = array_slice(col("a"), lit(1), lit(2), None);
98        assert_eq!(expr1, expr2);
99    }
100}
datafusion_functions_nested/expr_ext.rs

datafusion_functions_nested/
expr_ext.rs
