reifydb_core/value/column/
nulls.rs1use 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 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 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 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 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 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}