vortex_array/scalar_fn/fns/
is_not_null.rs1use std::fmt::Display;
5use std::fmt::Formatter;
6
7use vortex_error::VortexExpect as _;
8use vortex_error::VortexResult;
9use vortex_session::VortexSession;
10use vortex_session::registry::CachedId;
11
12use crate::ArrayRef;
13use crate::ExecutionCtx;
14use crate::IntoArray;
15use crate::arrays::ConstantArray;
16use crate::arrays::ScalarFnArray;
17use crate::dtype::DType;
18use crate::dtype::Nullability;
19use crate::expr::display::ExprDisplay;
20use crate::scalar_fn::Arity;
21use crate::scalar_fn::ChildName;
22use crate::scalar_fn::EmptyOptions;
23use crate::scalar_fn::ExecutionArgs;
24use crate::scalar_fn::ScalarFnId;
25use crate::scalar_fn::ScalarFnVTable;
26use crate::scalar_fn::ScalarFnVTableExt;
27use crate::validity::Validity;
28
29#[derive(Clone)]
31pub struct IsNotNull;
32
33impl IsNotNull {
34 #[expect(clippy::new_ret_no_self, reason = "constructs the lazy result array")]
36 pub fn new(input: ArrayRef) -> ScalarFnArray {
37 ScalarFnArray::try_new(IsNotNull.bind(EmptyOptions), vec![input])
38 .vortex_expect("IsNotNull has one child and an infallible return dtype")
39 }
40}
41
42impl ScalarFnVTable for IsNotNull {
43 type Options = EmptyOptions;
44
45 fn id(&self) -> ScalarFnId {
46 static ID: CachedId = CachedId::new("vortex.is_not_null");
47 *ID
48 }
49
50 fn serialize(&self, _instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
51 Ok(Some(vec![]))
52 }
53
54 fn deserialize(
55 &self,
56 _metadata: &[u8],
57 _session: &VortexSession,
58 ) -> VortexResult<Self::Options> {
59 Ok(EmptyOptions)
60 }
61
62 fn arity(&self, _options: &Self::Options) -> Arity {
63 Arity::Exact(1)
64 }
65
66 fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName {
67 match child_idx {
68 0 => ChildName::from("input"),
69 _ => unreachable!("Invalid child index {} for IsNotNull expression", child_idx),
70 }
71 }
72
73 fn fmt_sql(
74 &self,
75 _options: &Self::Options,
76 expr: &dyn ExprDisplay,
77 f: &mut Formatter<'_>,
78 ) -> std::fmt::Result {
79 write!(f, "is_not_null(")?;
80 Display::fmt(expr.display_child(0), f)?;
81 write!(f, ")")
82 }
83
84 fn return_dtype(&self, _options: &Self::Options, _arg_dtypes: &[DType]) -> VortexResult<DType> {
85 Ok(DType::Bool(Nullability::NonNullable))
86 }
87
88 fn execute(
89 &self,
90 _data: &Self::Options,
91 args: &dyn ExecutionArgs,
92 _ctx: &mut ExecutionCtx,
93 ) -> VortexResult<ArrayRef> {
94 let child = args.get(0)?;
95 match child.validity()? {
96 Validity::NonNullable | Validity::AllValid => {
97 Ok(ConstantArray::new(true, args.row_count()).into_array())
98 }
99 Validity::AllInvalid => Ok(ConstantArray::new(false, args.row_count()).into_array()),
100 Validity::Array(a) => Ok(a),
101 }
102 }
103
104 fn is_strict(&self, _instance: &Self::Options) -> bool {
105 false
107 }
108
109 fn is_infallible(&self, _instance: &Self::Options) -> bool {
110 true
111 }
112}
113
114#[cfg(test)]
115mod tests {
116 use std::sync::LazyLock;
117
118 use vortex_buffer::buffer;
119 use vortex_error::VortexExpect as _;
120 use vortex_error::VortexResult;
121 use vortex_session::VortexSession;
122
123 use crate::IntoArray;
124 use crate::VortexSessionExecute;
125 use crate::array_session;
126 use crate::arrays::PrimitiveArray;
127 use crate::arrays::StructArray;
128 use crate::dtype::DType;
129 use crate::dtype::Nullability;
130 use crate::expr::col;
131 use crate::expr::eq;
132 use crate::expr::get_item;
133 use crate::expr::is_not_null;
134 use crate::expr::or;
135 use crate::expr::root;
136 use crate::expr::test_harness;
137 use crate::scalar::Scalar;
138 use crate::scalar_fn::EmptyOptions;
139 use crate::scalar_fn::ScalarFnVTableExt;
140 use crate::scalar_fn::internal::row_count::RowCount;
141 use crate::stats::StatsSession;
142 use crate::stats::all_null;
143 use crate::stats::null_count;
144
145 static STATS_SESSION: LazyLock<VortexSession> =
146 LazyLock::new(|| VortexSession::empty().with::<StatsSession>());
147
148 #[test]
149 fn dtype() {
150 let dtype = test_harness::struct_dtype();
151 assert_eq!(
152 is_not_null(root()).return_dtype(&dtype).unwrap(),
153 DType::Bool(Nullability::NonNullable)
154 );
155 }
156
157 #[test]
158 fn replace_children() {
159 let expr = is_not_null(root());
160 expr.with_children([root()])
161 .vortex_expect("operation should succeed in test");
162 }
163
164 #[test]
165 fn evaluate_mask() {
166 let test_array =
167 PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
168 .into_array();
169 let expected = [true, false, true, false, true];
170
171 let result = test_array.clone().apply(&is_not_null(root())).unwrap();
172
173 assert_eq!(result.len(), test_array.len());
174 assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
175
176 for (i, expected_value) in expected.iter().enumerate() {
177 assert_eq!(
178 result
179 .execute_scalar(i, &mut array_session().create_execution_ctx())
180 .unwrap(),
181 Scalar::bool(*expected_value, Nullability::NonNullable)
182 );
183 }
184 }
185
186 #[test]
187 fn evaluate_all_true() {
188 let test_array = buffer![1, 2, 3, 4, 5].into_array();
189
190 let result = test_array.clone().apply(&is_not_null(root())).unwrap();
191
192 assert_eq!(result.len(), test_array.len());
193 for i in 0..result.len() {
194 assert_eq!(
195 result
196 .execute_scalar(i, &mut array_session().create_execution_ctx())
197 .unwrap(),
198 Scalar::bool(true, Nullability::NonNullable)
199 );
200 }
201 }
202
203 #[test]
204 fn evaluate_all_false() {
205 let test_array =
206 PrimitiveArray::from_option_iter(vec![None::<i32>, None, None, None, None])
207 .into_array();
208
209 let result = test_array.clone().apply(&is_not_null(root())).unwrap();
210
211 assert_eq!(result.len(), test_array.len());
212 for i in 0..result.len() {
213 assert_eq!(
214 result
215 .execute_scalar(i, &mut array_session().create_execution_ctx())
216 .unwrap(),
217 Scalar::bool(false, Nullability::NonNullable)
218 );
219 }
220 }
221
222 #[test]
223 fn evaluate_struct() {
224 let test_array = StructArray::from_fields(&[(
225 "a",
226 PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
227 .into_array(),
228 )])
229 .unwrap()
230 .into_array();
231 let expected = [true, false, true, false, true];
232
233 let result = test_array
234 .clone()
235 .apply(&is_not_null(get_item("a", root())))
236 .unwrap();
237
238 assert_eq!(result.len(), test_array.len());
239 assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
240
241 for (i, expected_value) in expected.iter().enumerate() {
242 assert_eq!(
243 result
244 .execute_scalar(i, &mut array_session().create_execution_ctx())
245 .unwrap(),
246 Scalar::bool(*expected_value, Nullability::NonNullable)
247 );
248 }
249 }
250
251 #[test]
252 fn test_display() {
253 let expr = is_not_null(get_item("name", root()));
254 assert_eq!(expr.to_string(), "is_not_null($.name)");
255
256 let expr2 = is_not_null(root());
257 assert_eq!(expr2.to_string(), "is_not_null($)");
258 }
259
260 #[test]
261 fn test_is_not_null_is_not_strict() {
262 assert!(
263 !is_not_null(col("a"))
264 .as_scalar()
265 .is_some_and(|f| f.signature().is_strict())
266 );
267 }
268
269 #[test]
270 fn test_is_not_null_falsification() -> VortexResult<()> {
271 let expr = is_not_null(col("a"));
272 let dtype = test_harness::struct_dtype();
273
274 assert_eq!(
275 expr.bind(&dtype)?.falsify(&STATS_SESSION)?,
276 Some(
277 or(
278 eq(null_count(col("a")), RowCount.new_expr(EmptyOptions, []),),
279 all_null(col("a")),
280 )
281 .bind(&dtype)?
282 )
283 );
284 Ok(())
285 }
286}