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
85pub 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}