Skip to content

Commit 7f115d2

Browse files
robert3005robert3005
committed
fix
1 parent 3237b54 commit 7f115d2

File tree

10 files changed

+61
-35
lines changed

10 files changed

+61
-35
lines changed

‎vortex-array/src/arrays/constant/compute/take.rs

+10-9
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
use vortex_error::VortexResult;
2-
use vortex_mask::AllOr;
2+
use vortex_mask::{AllOr, Mask};
33
use vortex_scalar::Scalar;
44

55
use crate::arrays::{ConstantArray, ConstantEncoding};
6-
use crate::builders::{ArrayBuilderExt, builder_with_capacity};
6+
use crate::builders::builder_with_capacity;
77
use crate::compute::TakeFn;
88
use crate::{Array, ArrayRef};
99

@@ -19,15 +19,16 @@ impl TakeFn<&ConstantArray> for ConstantEncoding {
1919
)
2020
.into_array()),
2121
AllOr::Some(v) => {
22+
let arr = ConstantArray::new(array.scalar().clone(), indices.len()).into_array();
23+
24+
if array.scalar().is_null() {
25+
return Ok(arr);
26+
}
27+
2228
let mut result_builder =
2329
builder_with_capacity(&array.dtype().as_nullable(), indices.len());
24-
for valid in v.iter() {
25-
if valid {
26-
result_builder.append_scalar_value(array.scalar().value().clone())?;
27-
} else {
28-
result_builder.append_null();
29-
}
30-
}
30+
result_builder.extend_from_array(&arr)?;
31+
result_builder.set_validity(Mask::from_buffer(v.clone()));
3132
Ok(result_builder.finish())
3233
}
3334
}

‎vortex-array/src/builders/bool.rs

+6
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use std::any::Any;
33
use arrow_buffer::BooleanBufferBuilder;
44
use vortex_dtype::{DType, Nullability};
55
use vortex_error::{VortexResult, vortex_bail};
6+
use vortex_mask::Mask;
67

78
use crate::arrays::BoolArray;
89
use crate::builders::ArrayBuilder;
@@ -85,6 +86,11 @@ impl ArrayBuilder for BoolBuilder {
8586
Ok(())
8687
}
8788

89+
fn set_validity(&mut self, validity: Mask) {
90+
self.nulls = LazyNullBufferBuilder::new(validity.len());
91+
self.nulls.append_validity_mask(validity);
92+
}
93+
8894
fn finish(&mut self) -> ArrayRef {
8995
assert_eq!(
9096
self.nulls.len(),

‎vortex-array/src/builders/extension.rs

+5
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use std::sync::Arc;
33

44
use vortex_dtype::{DType, ExtDType};
55
use vortex_error::{VortexResult, vortex_bail};
6+
use vortex_mask::Mask;
67
use vortex_scalar::ExtScalar;
78

89
use crate::arrays::ExtensionArray;
@@ -82,6 +83,10 @@ impl ArrayBuilder for ExtensionBuilder {
8283
array.storage().append_to_builder(self.storage.as_mut())
8384
}
8485

86+
fn set_validity(&mut self, validity: Mask) {
87+
self.storage.set_validity(validity);
88+
}
89+
8590
fn finish(&mut self) -> ArrayRef {
8691
let storage = self.storage.finish();
8792
ExtensionArray::new(self.ext_dtype(), storage).into_array()

‎vortex-array/src/builders/lazy_validity_builder.rs

-2
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,6 @@ impl LazyNullBufferBuilder {
5252
.append_n(n, false);
5353
}
5454

55-
#[allow(dead_code)]
5655
#[inline]
5756
pub fn append_null(&mut self) {
5857
self.materialize_if_needed();
@@ -62,7 +61,6 @@ impl LazyNullBufferBuilder {
6261
.append(false);
6362
}
6463

65-
#[allow(dead_code)]
6664
#[inline]
6765
pub fn append(&mut self, not_null: bool) {
6866
if not_null {

‎vortex-array/src/builders/list.rs

+6
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ use std::sync::Arc;
44
use vortex_dtype::Nullability::NonNullable;
55
use vortex_dtype::{DType, NativePType, Nullability};
66
use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
7+
use vortex_mask::Mask;
78
use vortex_scalar::{BinaryNumericOperator, ListScalar};
89

910
use crate::arrays::{ConstantArray, ListArray, OffsetPType};
@@ -150,6 +151,11 @@ impl<O: OffsetPType> ArrayBuilder for ListBuilder<O> {
150151
Ok(())
151152
}
152153

154+
fn set_validity(&mut self, validity: Mask) {
155+
self.nulls = LazyNullBufferBuilder::new(validity.len());
156+
self.nulls.append_validity_mask(validity);
157+
}
158+
153159
fn finish(&mut self) -> ArrayRef {
154160
assert_eq!(
155161
self.index_builder.len(),

‎vortex-array/src/builders/mod.rs

+4
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ pub use primitive::*;
1717
pub use varbinview::*;
1818
use vortex_dtype::{DType, match_each_native_ptype};
1919
use vortex_error::{VortexResult, vortex_bail, vortex_err};
20+
use vortex_mask::Mask;
2021
use vortex_scalar::{
2122
BinaryScalar, BoolScalar, ExtScalar, ListScalar, PrimitiveScalar, Scalar, ScalarValue,
2223
StructScalar, Utf8Scalar,
@@ -57,6 +58,9 @@ pub trait ArrayBuilder: Send {
5758
/// Extends the array with the provided array, canonicalizing if necessary.
5859
fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()>;
5960

61+
/// Override builders validity with the one provided
62+
fn set_validity(&mut self, validity: Mask);
63+
6064
/// Constructs an Array from the builder components.
6165
///
6266
/// # Panics

‎vortex-array/src/builders/null.rs

+5
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use std::any::Any;
22

33
use vortex_dtype::DType;
44
use vortex_error::VortexResult;
5+
use vortex_mask::Mask;
56

67
use crate::arrays::NullArray;
78
use crate::builders::ArrayBuilder;
@@ -54,6 +55,10 @@ impl ArrayBuilder for NullBuilder {
5455
Ok(())
5556
}
5657

58+
fn set_validity(&mut self, validity: Mask) {
59+
self.length = validity.len();
60+
}
61+
5762
fn finish(&mut self) -> ArrayRef {
5863
NullArray::new(self.length).into_array()
5964
}

‎vortex-array/src/builders/primitive.rs

+5
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,11 @@ impl<T: NativePType> ArrayBuilder for PrimitiveBuilder<T> {
178178
Ok(())
179179
}
180180

181+
fn set_validity(&mut self, validity: Mask) {
182+
self.nulls = LazyNullBufferBuilder::new(validity.len());
183+
self.nulls.append_validity_mask(validity);
184+
}
185+
181186
fn finish(&mut self) -> ArrayRef {
182187
self.finish_into_primitive().into_array()
183188
}

‎vortex-array/src/builders/struct_.rs

+15-24
Original file line numberDiff line numberDiff line change
@@ -4,18 +4,18 @@ use std::sync::Arc;
44
use itertools::Itertools;
55
use vortex_dtype::{DType, Nullability, StructDType};
66
use vortex_error::{VortexExpect, VortexResult, vortex_bail};
7+
use vortex_mask::Mask;
78
use vortex_scalar::StructScalar;
89

910
use crate::arrays::StructArray;
10-
use crate::builders::{ArrayBuilder, ArrayBuilderExt, BoolBuilder, builder_with_capacity};
11-
use crate::validity::Validity;
11+
use crate::builders::lazy_validity_builder::LazyNullBufferBuilder;
12+
use crate::builders::{ArrayBuilder, ArrayBuilderExt, builder_with_capacity};
1213
use crate::variants::StructArrayTrait;
1314
use crate::{Array, ArrayRef, Canonical};
1415

1516
pub struct StructBuilder {
1617
builders: Vec<Box<dyn ArrayBuilder>>,
17-
// TODO(ngates): this should be a NullBufferBuilder? Or mask builder?
18-
validity: BoolBuilder,
18+
validity: LazyNullBufferBuilder,
1919
struct_dtype: Arc<StructDType>,
2020
nullability: Nullability,
2121
dtype: DType,
@@ -34,7 +34,7 @@ impl StructBuilder {
3434

3535
Self {
3636
builders,
37-
validity: BoolBuilder::with_capacity(Nullability::NonNullable, capacity),
37+
validity: LazyNullBufferBuilder::new(capacity),
3838
struct_dtype: struct_dtype.clone(),
3939
nullability,
4040
dtype: DType::Struct(struct_dtype, nullability),
@@ -54,7 +54,7 @@ impl StructBuilder {
5454
for (builder, field) in self.builders.iter_mut().zip_eq(fields) {
5555
builder.append_scalar(&field)?;
5656
}
57-
self.validity.append_value(true);
57+
self.validity.append_non_null();
5858
} else {
5959
self.append_null()
6060
}
@@ -84,7 +84,7 @@ impl ArrayBuilder for StructBuilder {
8484
self.builders
8585
.iter_mut()
8686
.for_each(|builder| builder.append_zeros(n));
87-
self.validity.append_values(true, n);
87+
self.validity.append_n_non_nulls(n);
8888
}
8989

9090
fn append_nulls(&mut self, n: usize) {
@@ -93,7 +93,7 @@ impl ArrayBuilder for StructBuilder {
9393
// We push zero values into our children when appending a null in case the children are
9494
// themselves non-nullable.
9595
.for_each(|builder| builder.append_zeros(n));
96-
self.validity.append_value(false);
96+
self.validity.append_null();
9797
}
9898

9999
fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()> {
@@ -120,21 +120,15 @@ impl ArrayBuilder for StructBuilder {
120120
a.append_to_builder(builder.as_mut())?;
121121
}
122122

123-
match array.validity() {
124-
Validity::NonNullable | Validity::AllValid => {
125-
self.validity.append_values(true, array.len());
126-
}
127-
Validity::AllInvalid => {
128-
self.validity.append_values(false, array.len());
129-
}
130-
Validity::Array(validity) => {
131-
validity.append_to_builder(&mut self.validity)?;
132-
}
133-
}
134-
123+
self.validity.append_validity_mask(array.validity_mask()?);
135124
Ok(())
136125
}
137126

127+
fn set_validity(&mut self, validity: Mask) {
128+
self.validity = LazyNullBufferBuilder::new(validity.len());
129+
self.validity.append_validity_mask(validity);
130+
}
131+
138132
fn finish(&mut self) -> ArrayRef {
139133
let len = self.len();
140134
let fields = self
@@ -156,10 +150,7 @@ impl ArrayBuilder for StructBuilder {
156150
}
157151
}
158152

159-
let validity = match self.nullability {
160-
Nullability::NonNullable => Validity::NonNullable,
161-
Nullability::Nullable => Validity::Array(self.validity.finish()),
162-
};
153+
let validity = self.validity.finish_with_nullability(self.nullability);
163154

164155
StructArray::try_new(self.struct_dtype.names().clone(), fields, len, validity)
165156
.vortex_expect("Fields must all have same length.")

‎vortex-array/src/builders/varbinview.rs

+5
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,11 @@ impl ArrayBuilder for VarBinViewBuilder {
177177
Ok(())
178178
}
179179

180+
fn set_validity(&mut self, validity: Mask) {
181+
self.null_buffer_builder = LazyNullBufferBuilder::new(validity.len());
182+
self.null_buffer_builder.append_validity_mask(validity);
183+
}
184+
180185
fn finish(&mut self) -> ArrayRef {
181186
self.flush_in_progress();
182187
let buffers = std::mem::take(&mut self.completed);

0 commit comments

Comments
 (0)