Skip to main content

reifydb_core/value/column/
nulls.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use std::sync::Arc;
5
6use arrow_array::{Array, ArrayRef, make_array, new_null_array};
7use arrow_buffer::NullBuffer;
8use arrow_schema::{DataType, FieldRef};
9use reifydb_value::{
10	Result,
11	util::bitmap,
12	value::{
13		column_view::{ColumnView, ViewData},
14		container::{fixed_array, wide_int_array::wide_array},
15		value_type::ValueType,
16	},
17};
18
19use crate::value::batch::frame_error;
20
21pub fn split_nulls(column: (FieldRef, ArrayRef)) -> Result<((FieldRef, ArrayRef), Option<NullBuffer>)> {
22	let (field, array) = column;
23	let keeps_own = {
24		let view = ColumnView::try_from((&array, field.as_ref()))?;
25		matches!(view.data, ViewData::Any { .. } | ViewData::Digest { .. } | ViewData::None { .. })
26	};
27	let nulls = match array.logical_nulls() {
28		Some(nulls) => Some(nulls),
29		None if field.is_nullable() => Some(NullBuffer::new_valid(array.len())),
30		None => None,
31	};
32	if keeps_own || nulls.is_none() {
33		return Ok(((field, array), nulls));
34	}
35	let data = array.to_data().into_builder().nulls(None).build().map_err(frame_error)?;
36	let field = Arc::new(field.as_ref().clone().with_nullable(false));
37	Ok(((field, make_array(data)), nulls))
38}
39
40pub fn with_nulls(column: (FieldRef, ArrayRef), nulls: NullBuffer) -> Result<(FieldRef, ArrayRef)> {
41	let (field, array) = column;
42	let len = array.len();
43	assert_eq!(nulls.len(), len, "validity of {} rows does not match a column of {len} rows", nulls.len());
44	let field = Arc::new(field.as_ref().clone().with_nullable(true));
45	if array.data_type().is_null() {
46		return Ok((field, array));
47	}
48	let nulls = match array.logical_nulls() {
49		Some(existing) => bitmap::and_nulls(&existing, &nulls),
50		None => nulls,
51	};
52	let data = array.to_data().into_builder().nulls(Some(nulls)).build().map_err(frame_error)?;
53	Ok((field, make_array(data)))
54}
55
56pub fn none_filler(ty: &ValueType, data_type: &DataType) -> ArrayRef {
57	match ty.inner_type() {
58		ValueType::Int16 => {
59			Arc::new(fixed_array::attach_nulls(wide_array([0i128]), Some(NullBuffer::new_null(1))))
60		}
61		_ => new_null_array(data_type, 1),
62	}
63}
64
65#[cfg(test)]
66mod tests {
67	use arrow_buffer::{BooleanBuffer, NullBuffer};
68	use reifydb_value::value::{Value, column_view::ColumnView, value_type::ValueType};
69
70	use super::{split_nulls, with_nulls};
71	use crate::value::column::factory::{any_optional, int4, int4_with_bitvec, none};
72
73	#[test]
74	fn split_then_with_nulls_round_trips_an_optional_column() {
75		// Losing the stripped validity would turn the none row into its hidden default value.
76		let column = int4_with_bitvec("a", [1, 2], BooleanBuffer::from(vec![true, false]));
77
78		let (bare, nulls) = split_nulls(column).unwrap();
79		assert!(!bare.0.is_nullable());
80		assert_eq!(ColumnView::try_from(&bare).unwrap().get_type(), ValueType::Int4);
81
82		let restored = with_nulls(bare, nulls.unwrap()).unwrap();
83		let view = ColumnView::try_from(&restored).unwrap();
84		assert_eq!(view.get_type(), ValueType::Option(Box::new(ValueType::Int4)));
85		assert_eq!(view.get_value(0), Value::Int4(1));
86		assert_eq!(view.get_value(1), Value::none_of(ValueType::Int4));
87	}
88
89	#[test]
90	fn split_keeps_the_nulls_inside_an_any_column() {
91		// An any column stores none rows inside its own encoding, so stripping them would corrupt the rows.
92		let column = any_optional("a", [Some(Value::Int4(1)), None]);
93
94		let (kept, nulls) = split_nulls(column).unwrap();
95
96		assert!(nulls.is_some());
97		assert_eq!(ColumnView::try_from(&kept).unwrap().get_value(1), Value::none_of(ValueType::Any));
98	}
99
100	#[test]
101	fn split_reports_all_valid_nulls_for_a_nullable_column_without_a_buffer() {
102		// A nullable field with no none row must stay optional when the nulls are put back.
103		let column = with_nulls(int4("a", [1, 2]), NullBuffer::new_valid(2)).unwrap();
104
105		let (bare, nulls) = split_nulls(column).unwrap();
106
107		assert!(!bare.0.is_nullable());
108		assert_eq!(nulls.map(|nulls| nulls.null_count()), Some(0));
109	}
110
111	#[test]
112	fn with_nulls_keeps_existing_none_rows() {
113		// Replacing instead of combining would resurrect a row that was already none.
114		let column = int4_with_bitvec("a", [1, 2, 3], BooleanBuffer::from(vec![false, true, true]));
115
116		let merged = with_nulls(column, NullBuffer::new(BooleanBuffer::from(vec![true, true, false]))).unwrap();
117
118		let view = ColumnView::try_from(&merged).unwrap();
119		assert_eq!(view.get_value(0), Value::none_of(ValueType::Int4));
120		assert_eq!(view.get_value(1), Value::Int4(2));
121		assert_eq!(view.get_value(2), Value::none_of(ValueType::Int4));
122	}
123
124	#[test]
125	fn split_keeps_a_none_column_whole() {
126		// An untyped none column has no values to strip, so it must come back unchanged with all rows none.
127		let (kept, nulls) = split_nulls(none("a", 2)).unwrap();
128
129		assert!(ColumnView::try_from(&kept).unwrap().is_none());
130		assert_eq!(nulls.map(|nulls| nulls.null_count()), Some(2));
131	}
132}