iceberg/expr/
accessor.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
18use std::sync::Arc;
19
20use serde_derive::{Deserialize, Serialize};
21
22use crate::spec::{Datum, Literal, PrimitiveType, Struct};
23use crate::{Error, ErrorKind, Result};
24
25#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
26pub struct StructAccessor {
27    position: usize,
28    r#type: PrimitiveType,
29    inner: Option<Box<StructAccessor>>,
30}
31
32pub(crate) type StructAccessorRef = Arc<StructAccessor>;
33
34impl StructAccessor {
35    pub(crate) fn new(position: usize, r#type: PrimitiveType) -> Self {
36        StructAccessor {
37            position,
38            r#type,
39            inner: None,
40        }
41    }
42
43    pub(crate) fn wrap(position: usize, inner: Box<StructAccessor>) -> Self {
44        StructAccessor {
45            position,
46            r#type: inner.r#type().clone(),
47            inner: Some(inner),
48        }
49    }
50
51    pub(crate) fn position(&self) -> usize {
52        self.position
53    }
54
55    pub(crate) fn r#type(&self) -> &PrimitiveType {
56        &self.r#type
57    }
58
59    pub(crate) fn get<'a>(&'a self, container: &'a Struct) -> Result<Option<Datum>> {
60        match &self.inner {
61            None => match &container[self.position] {
62                None => Ok(None),
63                Some(Literal::Primitive(literal)) => {
64                    Ok(Some(Datum::new(self.r#type().clone(), literal.clone())))
65                }
66                Some(_) => Err(Error::new(
67                    ErrorKind::Unexpected,
68                    "Expected Literal to be Primitive",
69                )),
70            },
71            Some(inner) => {
72                if let Some(Literal::Struct(wrapped)) = &container[self.position] {
73                    inner.get(wrapped)
74                } else {
75                    Err(Error::new(
76                        ErrorKind::Unexpected,
77                        "Nested accessor should only be wrapping a Struct",
78                    ))
79                }
80            }
81        }
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use crate::expr::accessor::StructAccessor;
88    use crate::spec::{Datum, Literal, PrimitiveType, Struct};
89
90    #[test]
91    fn test_single_level_accessor() {
92        let accessor = StructAccessor::new(1, PrimitiveType::Boolean);
93
94        assert_eq!(accessor.r#type(), &PrimitiveType::Boolean);
95        assert_eq!(accessor.position(), 1);
96
97        let test_struct =
98            Struct::from_iter(vec![Some(Literal::bool(false)), Some(Literal::bool(true))]);
99
100        assert_eq!(accessor.get(&test_struct).unwrap(), Some(Datum::bool(true)));
101    }
102
103    #[test]
104    fn test_single_level_accessor_null() {
105        let accessor = StructAccessor::new(1, PrimitiveType::Boolean);
106
107        assert_eq!(accessor.r#type(), &PrimitiveType::Boolean);
108        assert_eq!(accessor.position(), 1);
109
110        let test_struct = Struct::from_iter(vec![Some(Literal::bool(false)), None]);
111
112        assert_eq!(accessor.get(&test_struct).unwrap(), None);
113    }
114
115    #[test]
116    fn test_nested_accessor() {
117        let nested_accessor = StructAccessor::new(1, PrimitiveType::Boolean);
118        let accessor = StructAccessor::wrap(2, Box::new(nested_accessor));
119
120        assert_eq!(accessor.r#type(), &PrimitiveType::Boolean);
121        //assert_eq!(accessor.position(), 1);
122
123        let nested_test_struct =
124            Struct::from_iter(vec![Some(Literal::bool(false)), Some(Literal::bool(true))]);
125
126        let test_struct = Struct::from_iter(vec![
127            Some(Literal::bool(false)),
128            Some(Literal::bool(false)),
129            Some(Literal::Struct(nested_test_struct)),
130        ]);
131
132        assert_eq!(accessor.get(&test_struct).unwrap(), Some(Datum::bool(true)));
133    }
134
135    #[test]
136    fn test_nested_accessor_null() {
137        let nested_accessor = StructAccessor::new(0, PrimitiveType::Boolean);
138        let accessor = StructAccessor::wrap(2, Box::new(nested_accessor));
139
140        assert_eq!(accessor.r#type(), &PrimitiveType::Boolean);
141        //assert_eq!(accessor.position(), 1);
142
143        let nested_test_struct = Struct::from_iter(vec![None, Some(Literal::bool(true))]);
144
145        let test_struct = Struct::from_iter(vec![
146            Some(Literal::bool(false)),
147            Some(Literal::bool(false)),
148            Some(Literal::Struct(nested_test_struct)),
149        ]);
150
151        assert_eq!(accessor.get(&test_struct).unwrap(), None);
152    }
153}