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
16#[cfg(feature = "proto")]
17pub(crate) mod proto {
18    use vortex_error::VortexResult;
19    use vortex_proto::expr::kind;
20    use vortex_proto::expr::kind::Kind;
21
22    use crate::identity::IDENTITY;
23    use crate::{ExprDeserialize, ExprRef, ExprSerializable, Id, Identity};
24
25    pub(crate) struct IdentitySerde;
26
27    impl Id for IdentitySerde {
28        fn id(&self) -> &'static str {
29            "identity"
30        }
31    }
32
33    impl ExprDeserialize for IdentitySerde {
34        fn deserialize(&self, _expr: &Kind, _children: Vec<ExprRef>) -> VortexResult<ExprRef> {
35            Ok(IDENTITY.clone())
36        }
37    }
38
39    impl ExprSerializable for Identity {
40        fn id(&self) -> &'static str {
41            IdentitySerde.id()
42        }
43
44        fn serialize_kind(&self) -> VortexResult<Kind> {
45            Ok(Kind::Identity(kind::Identity {}))
46        }
47    }
48}
49
50impl Identity {
51    pub fn new_expr() -> ExprRef {
52        IDENTITY.clone()
53    }
54}
55
56impl Display for Identity {
57    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        write!(f, "$")
59    }
60}
61
62impl VortexExpr for Identity {
63    fn as_any(&self) -> &dyn Any {
64        self
65    }
66
67    fn unchecked_evaluate(&self, batch: &dyn Array) -> VortexResult<ArrayRef> {
68        Ok(batch.to_array())
69    }
70
71    fn children(&self) -> Vec<&ExprRef> {
72        vec![]
73    }
74
75    fn replacing_children(self: Arc<Self>, children: Vec<ExprRef>) -> ExprRef {
76        assert_eq!(children.len(), 0);
77        self
78    }
79
80    fn return_dtype(&self, scope_dtype: &DType) -> VortexResult<DType> {
81        Ok(scope_dtype.clone())
82    }
83}
84
85// Return a global pointer to the identity token.
86pub fn ident() -> ExprRef {
87    Identity::new_expr()
88}
89
90#[cfg(test)]
91mod tests {
92    use std::sync::Arc;
93
94    use vortex_dtype::{DType, Nullability, PType};
95
96    use crate::{ident, test_harness};
97
98    #[test]
99    fn dtype() {
100        let dtype = test_harness::struct_dtype();
101        assert_eq!(ident().return_dtype(&dtype).unwrap(), dtype);
102        assert_eq!(ident().return_dtype(&dtype).unwrap(), dtype);
103    }
104
105    #[test]
106    fn list_dtype() {
107        let in_dtype = DType::List(
108            Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)),
109            Nullability::Nullable,
110        );
111        assert_eq!(ident().return_dtype(&in_dtype).unwrap(), in_dtype);
112    }
113}