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

datafusion_functions_nested/
expr_ext.rs
