vortex_expr/
identity.rs

1use std::any::Any;
2use std::fmt::Display;
3use std::sync::{Arc, LazyLock};
4
5use vortex_array::{Array, ArrayRef};
6use vortex_dtype::DType;
7use vortex_error::VortexResult;
8
9use crate::{ExprRef, VortexExpr};
10
11static IDENTITY: LazyLock<ExprRef> = LazyLock::new(|| Arc::new(Identity));
12
13#[derive(Debug, PartialEq, Eq, Hash)]
14pub struct Identity;
15
16impl Identity {
17    pub fn new_expr() -> ExprRef {
18        IDENTITY.clone()
19    }
20}
21
22impl Display for Identity {
23    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24        write!(f, "$")
25    }
26}
27
28impl VortexExpr for Identity {
29    fn as_any(&self) -> &dyn Any {
30        self
31    }
32
33    fn unchecked_evaluate(&self, batch: &dyn Array) -> VortexResult<ArrayRef> {
34        Ok(batch.to_array())
35    }
36
37    fn children(&self) -> Vec<&ExprRef> {
38        vec![]
39    }
40
41    fn replacing_children(self: Arc<Self>, children: Vec<ExprRef>) -> ExprRef {
42        assert_eq!(children.len(), 0);
43        self
44    }
45
46    fn return_dtype(&self, scope_dtype: &DType) -> VortexResult<DType> {
47        Ok(scope_dtype.clone())
48    }
49}
50
51// Return a global pointer to the identity token.
52pub fn ident() -> ExprRef {
53    Identity::new_expr()
54}
55
56#[cfg(test)]
57mod tests {
58    use std::sync::Arc;
59
60    use vortex_dtype::{DType, Nullability, PType};
61
62    use crate::{ident, test_harness};
63
64    #[test]
65    fn dtype() {
66        let dtype = test_harness::struct_dtype();
67        assert_eq!(ident().return_dtype(&dtype).unwrap(), dtype);
68        assert_eq!(ident().return_dtype(&dtype).unwrap(), dtype);
69    }
70
71    #[test]
72    fn list_dtype() {
73        let in_dtype = DType::List(
74            Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)),
75            Nullability::Nullable,
76        );
77        assert_eq!(ident().return_dtype(&in_dtype).unwrap(), in_dtype);
78    }
79}