/Users/andrewlamb/Software/arrow-rs/arrow-array/src/builder/primitive_builder.rs
Line | Count | Source |
1 | | // Licensed to the Apache Software Foundation (ASF) under one |
2 | | // or more contributor license agreements. See the NOTICE file |
3 | | // distributed with this work for additional information |
4 | | // regarding copyright ownership. The ASF licenses this file |
5 | | // to you under the Apache License, Version 2.0 (the |
6 | | // "License"); you may not use this file except in compliance |
7 | | // with the License. You may obtain a copy of the License at |
8 | | // |
9 | | // http://www.apache.org/licenses/LICENSE-2.0 |
10 | | // |
11 | | // Unless required by applicable law or agreed to in writing, |
12 | | // software distributed under the License is distributed on an |
13 | | // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
14 | | // KIND, either express or implied. See the License for the |
15 | | // specific language governing permissions and limitations |
16 | | // under the License. |
17 | | |
18 | | use crate::builder::ArrayBuilder; |
19 | | use crate::types::*; |
20 | | use crate::{Array, ArrayRef, PrimitiveArray}; |
21 | | use arrow_buffer::{Buffer, MutableBuffer, NullBufferBuilder, ScalarBuffer}; |
22 | | use arrow_data::ArrayData; |
23 | | use arrow_schema::{ArrowError, DataType}; |
24 | | use std::any::Any; |
25 | | use std::sync::Arc; |
26 | | |
27 | | /// A signed 8-bit integer array builder. |
28 | | pub type Int8Builder = PrimitiveBuilder<Int8Type>; |
29 | | /// A signed 16-bit integer array builder. |
30 | | pub type Int16Builder = PrimitiveBuilder<Int16Type>; |
31 | | /// A signed 32-bit integer array builder. |
32 | | pub type Int32Builder = PrimitiveBuilder<Int32Type>; |
33 | | /// A signed 64-bit integer array builder. |
34 | | pub type Int64Builder = PrimitiveBuilder<Int64Type>; |
35 | | /// An usigned 8-bit integer array builder. |
36 | | pub type UInt8Builder = PrimitiveBuilder<UInt8Type>; |
37 | | /// An usigned 16-bit integer array builder. |
38 | | pub type UInt16Builder = PrimitiveBuilder<UInt16Type>; |
39 | | /// An usigned 32-bit integer array builder. |
40 | | pub type UInt32Builder = PrimitiveBuilder<UInt32Type>; |
41 | | /// An usigned 64-bit integer array builder. |
42 | | pub type UInt64Builder = PrimitiveBuilder<UInt64Type>; |
43 | | /// A 16-bit floating point array builder. |
44 | | pub type Float16Builder = PrimitiveBuilder<Float16Type>; |
45 | | /// A 32-bit floating point array builder. |
46 | | pub type Float32Builder = PrimitiveBuilder<Float32Type>; |
47 | | /// A 64-bit floating point array builder. |
48 | | pub type Float64Builder = PrimitiveBuilder<Float64Type>; |
49 | | |
50 | | /// A timestamp second array builder. |
51 | | pub type TimestampSecondBuilder = PrimitiveBuilder<TimestampSecondType>; |
52 | | /// A timestamp millisecond array builder. |
53 | | pub type TimestampMillisecondBuilder = PrimitiveBuilder<TimestampMillisecondType>; |
54 | | /// A timestamp microsecond array builder. |
55 | | pub type TimestampMicrosecondBuilder = PrimitiveBuilder<TimestampMicrosecondType>; |
56 | | /// A timestamp nanosecond array builder. |
57 | | pub type TimestampNanosecondBuilder = PrimitiveBuilder<TimestampNanosecondType>; |
58 | | |
59 | | /// A 32-bit date array builder. |
60 | | pub type Date32Builder = PrimitiveBuilder<Date32Type>; |
61 | | /// A 64-bit date array builder. |
62 | | pub type Date64Builder = PrimitiveBuilder<Date64Type>; |
63 | | |
64 | | /// A 32-bit elaspsed time in seconds array builder. |
65 | | pub type Time32SecondBuilder = PrimitiveBuilder<Time32SecondType>; |
66 | | /// A 32-bit elaspsed time in milliseconds array builder. |
67 | | pub type Time32MillisecondBuilder = PrimitiveBuilder<Time32MillisecondType>; |
68 | | /// A 64-bit elaspsed time in microseconds array builder. |
69 | | pub type Time64MicrosecondBuilder = PrimitiveBuilder<Time64MicrosecondType>; |
70 | | /// A 64-bit elaspsed time in nanoseconds array builder. |
71 | | pub type Time64NanosecondBuilder = PrimitiveBuilder<Time64NanosecondType>; |
72 | | |
73 | | /// A “calendar” interval in months array builder. |
74 | | pub type IntervalYearMonthBuilder = PrimitiveBuilder<IntervalYearMonthType>; |
75 | | /// A “calendar” interval in days and milliseconds array builder. |
76 | | pub type IntervalDayTimeBuilder = PrimitiveBuilder<IntervalDayTimeType>; |
77 | | /// A “calendar” interval in months, days, and nanoseconds array builder. |
78 | | pub type IntervalMonthDayNanoBuilder = PrimitiveBuilder<IntervalMonthDayNanoType>; |
79 | | |
80 | | /// An elapsed time in seconds array builder. |
81 | | pub type DurationSecondBuilder = PrimitiveBuilder<DurationSecondType>; |
82 | | /// An elapsed time in milliseconds array builder. |
83 | | pub type DurationMillisecondBuilder = PrimitiveBuilder<DurationMillisecondType>; |
84 | | /// An elapsed time in microseconds array builder. |
85 | | pub type DurationMicrosecondBuilder = PrimitiveBuilder<DurationMicrosecondType>; |
86 | | /// An elapsed time in nanoseconds array builder. |
87 | | pub type DurationNanosecondBuilder = PrimitiveBuilder<DurationNanosecondType>; |
88 | | |
89 | | /// A decimal 32 array builder |
90 | | pub type Decimal32Builder = PrimitiveBuilder<Decimal32Type>; |
91 | | /// A decimal 64 array builder |
92 | | pub type Decimal64Builder = PrimitiveBuilder<Decimal64Type>; |
93 | | /// A decimal 128 array builder |
94 | | pub type Decimal128Builder = PrimitiveBuilder<Decimal128Type>; |
95 | | /// A decimal 256 array builder |
96 | | pub type Decimal256Builder = PrimitiveBuilder<Decimal256Type>; |
97 | | |
98 | | /// Builder for [`PrimitiveArray`] |
99 | | #[derive(Debug)] |
100 | | pub struct PrimitiveBuilder<T: ArrowPrimitiveType> { |
101 | | values_builder: Vec<T::Native>, |
102 | | null_buffer_builder: NullBufferBuilder, |
103 | | data_type: DataType, |
104 | | } |
105 | | |
106 | | impl<T: ArrowPrimitiveType> ArrayBuilder for PrimitiveBuilder<T> { |
107 | | /// Returns the builder as a non-mutable `Any` reference. |
108 | 0 | fn as_any(&self) -> &dyn Any { |
109 | 0 | self |
110 | 0 | } |
111 | | |
112 | | /// Returns the builder as a mutable `Any` reference. |
113 | 9 | fn as_any_mut(&mut self) -> &mut dyn Any { |
114 | 9 | self |
115 | 9 | } |
116 | | |
117 | | /// Returns the boxed builder as a box of `Any`. |
118 | 0 | fn into_box_any(self: Box<Self>) -> Box<dyn Any> { |
119 | 0 | self |
120 | 0 | } |
121 | | |
122 | | /// Returns the number of array slots in the builder |
123 | 140 | fn len(&self) -> usize { |
124 | 140 | self.values_builder.len() |
125 | 140 | } |
126 | | |
127 | | /// Builds the array and reset this builder. |
128 | 11 | fn finish(&mut self) -> ArrayRef { |
129 | 11 | Arc::new(self.finish()) |
130 | 11 | } |
131 | | |
132 | | /// Builds the array without resetting the builder. |
133 | 0 | fn finish_cloned(&self) -> ArrayRef { |
134 | 0 | Arc::new(self.finish_cloned()) |
135 | 0 | } |
136 | | } |
137 | | |
138 | | impl<T: ArrowPrimitiveType> Default for PrimitiveBuilder<T> { |
139 | | fn default() -> Self { |
140 | | Self::new() |
141 | | } |
142 | | } |
143 | | |
144 | | impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> { |
145 | | /// Creates a new primitive array builder |
146 | 26 | pub fn new() -> Self { |
147 | 26 | Self::with_capacity(1024) |
148 | 26 | } |
149 | | |
150 | | /// Creates a new primitive array builder with capacity no of items |
151 | 83 | pub fn with_capacity(capacity: usize) -> Self { |
152 | 83 | Self { |
153 | 83 | values_builder: Vec::with_capacity(capacity), |
154 | 83 | null_buffer_builder: NullBufferBuilder::new(capacity), |
155 | 83 | data_type: T::DATA_TYPE, |
156 | 83 | } |
157 | 83 | } |
158 | | |
159 | | /// Creates a new primitive array builder from buffers |
160 | | pub fn new_from_buffer( |
161 | | values_buffer: MutableBuffer, |
162 | | null_buffer: Option<MutableBuffer>, |
163 | | ) -> Self { |
164 | | let values_builder: Vec<T::Native> = ScalarBuffer::<T::Native>::from(values_buffer).into(); |
165 | | |
166 | | let null_buffer_builder = null_buffer |
167 | | .map(|buffer| NullBufferBuilder::new_from_buffer(buffer, values_builder.len())) |
168 | | .unwrap_or_else(|| NullBufferBuilder::new_with_len(values_builder.len())); |
169 | | |
170 | | Self { |
171 | | values_builder, |
172 | | null_buffer_builder, |
173 | | data_type: T::DATA_TYPE, |
174 | | } |
175 | | } |
176 | | |
177 | | /// By default [`PrimitiveBuilder`] uses [`ArrowPrimitiveType::DATA_TYPE`] as the |
178 | | /// data type of the generated array. |
179 | | /// |
180 | | /// This method allows overriding the data type, to allow specifying timezones |
181 | | /// for [`DataType::Timestamp`] or precision and scale for [`DataType::Decimal32`], |
182 | | /// [`DataType::Decimal64`], [`DataType::Decimal128`] and [`DataType::Decimal256`] |
183 | | /// |
184 | | /// # Panics |
185 | | /// |
186 | | /// This method panics if `data_type` is not [PrimitiveArray::is_compatible] |
187 | 57 | pub fn with_data_type(self, data_type: DataType) -> Self { |
188 | 57 | assert!( |
189 | 57 | PrimitiveArray::<T>::is_compatible(&data_type), |
190 | 0 | "incompatible data type for builder, expected {} got {}", |
191 | 0 | T::DATA_TYPE, |
192 | | data_type |
193 | | ); |
194 | 57 | Self { data_type, ..self } |
195 | 57 | } |
196 | | |
197 | | /// Returns the capacity of this builder measured in slots of type `T` |
198 | | pub fn capacity(&self) -> usize { |
199 | | self.values_builder.capacity() |
200 | | } |
201 | | |
202 | | /// Appends a value of type `T` into the builder |
203 | | #[inline] |
204 | 226 | pub fn append_value(&mut self, v: T::Native) { |
205 | 226 | self.null_buffer_builder.append_non_null(); |
206 | 226 | self.values_builder.push(v); |
207 | 226 | } |
208 | | |
209 | | /// Appends a value of type `T` into the builder `n` times |
210 | | #[inline] |
211 | | pub fn append_value_n(&mut self, v: T::Native, n: usize) { |
212 | | self.null_buffer_builder.append_n_non_nulls(n); |
213 | | self.values_builder.extend(std::iter::repeat_n(v, n)); |
214 | | } |
215 | | |
216 | | /// Appends a null slot into the builder |
217 | | #[inline] |
218 | 4 | pub fn append_null(&mut self) { |
219 | 4 | self.null_buffer_builder.append_null(); |
220 | 4 | self.values_builder.push(T::Native::default()); |
221 | 4 | } |
222 | | |
223 | | /// Appends `n` no. of null's into the builder |
224 | | #[inline] |
225 | 0 | pub fn append_nulls(&mut self, n: usize) { |
226 | 0 | self.null_buffer_builder.append_n_nulls(n); |
227 | 0 | self.values_builder |
228 | 0 | .extend(std::iter::repeat_n(T::Native::default(), n)); |
229 | 0 | } |
230 | | |
231 | | /// Appends an `Option<T>` into the builder |
232 | | #[inline] |
233 | 0 | pub fn append_option(&mut self, v: Option<T::Native>) { |
234 | 0 | match v { |
235 | 0 | None => self.append_null(), |
236 | 0 | Some(v) => self.append_value(v), |
237 | | }; |
238 | 0 | } |
239 | | |
240 | | /// Appends a slice of type `T` into the builder |
241 | | #[inline] |
242 | | pub fn append_slice(&mut self, v: &[T::Native]) { |
243 | | self.null_buffer_builder.append_n_non_nulls(v.len()); |
244 | | self.values_builder.extend_from_slice(v); |
245 | | } |
246 | | |
247 | | /// Appends values from a slice of type `T` and a validity boolean slice |
248 | | /// |
249 | | /// # Panics |
250 | | /// |
251 | | /// Panics if `values` and `is_valid` have different lengths |
252 | | #[inline] |
253 | | pub fn append_values(&mut self, values: &[T::Native], is_valid: &[bool]) { |
254 | | assert_eq!( |
255 | | values.len(), |
256 | | is_valid.len(), |
257 | | "Value and validity lengths must be equal" |
258 | | ); |
259 | | self.null_buffer_builder.append_slice(is_valid); |
260 | | self.values_builder.extend_from_slice(values); |
261 | | } |
262 | | |
263 | | /// Appends array values and null to this builder as is |
264 | | /// (this means that underlying null values are copied as is). |
265 | | /// |
266 | | /// # Panics |
267 | | /// |
268 | | /// Panics if `array` and `self` data types are different |
269 | | #[inline] |
270 | 188 | pub fn append_array(&mut self, array: &PrimitiveArray<T>) { |
271 | 188 | assert_eq!( |
272 | 188 | &self.data_type, |
273 | 188 | array.data_type(), |
274 | 0 | "array data type mismatch" |
275 | | ); |
276 | | |
277 | 188 | self.values_builder.extend_from_slice(array.values()); |
278 | 188 | if let Some(null_buffer8 ) = array.nulls() { |
279 | 8 | self.null_buffer_builder.append_buffer(null_buffer); |
280 | 180 | } else { |
281 | 180 | self.null_buffer_builder.append_n_non_nulls(array.len()); |
282 | 180 | } |
283 | 188 | } |
284 | | |
285 | | /// Appends values from a trusted length iterator. |
286 | | /// |
287 | | /// # Safety |
288 | | /// This requires the iterator be a trusted length. This could instead require |
289 | | /// the iterator implement `TrustedLen` once that is stabilized. |
290 | | #[inline] |
291 | 0 | pub unsafe fn append_trusted_len_iter(&mut self, iter: impl IntoIterator<Item = T::Native>) { |
292 | 0 | let iter = iter.into_iter(); |
293 | 0 | let len = iter |
294 | 0 | .size_hint() |
295 | 0 | .1 |
296 | 0 | .expect("append_trusted_len_iter requires an upper bound"); |
297 | | |
298 | 0 | self.null_buffer_builder.append_n_non_nulls(len); |
299 | 0 | self.values_builder.extend(iter); |
300 | 0 | } |
301 | | |
302 | | /// Builds the [`PrimitiveArray`] and reset this builder. |
303 | 119 | pub fn finish(&mut self) -> PrimitiveArray<T> { |
304 | 119 | let len = self.len(); |
305 | 119 | let nulls = self.null_buffer_builder.finish(); |
306 | 119 | let builder = ArrayData::builder(self.data_type.clone()) |
307 | 119 | .len(len) |
308 | 119 | .add_buffer(std::mem::take(&mut self.values_builder).into()) |
309 | 119 | .nulls(nulls); |
310 | | |
311 | 119 | let array_data = unsafe { builder.build_unchecked() }; |
312 | 119 | PrimitiveArray::<T>::from(array_data) |
313 | 119 | } |
314 | | |
315 | | /// Builds the [`PrimitiveArray`] without resetting the builder. |
316 | 0 | pub fn finish_cloned(&self) -> PrimitiveArray<T> { |
317 | 0 | let len = self.len(); |
318 | 0 | let nulls = self.null_buffer_builder.finish_cloned(); |
319 | 0 | let values_buffer = Buffer::from_slice_ref(self.values_builder.as_slice()); |
320 | 0 | let builder = ArrayData::builder(self.data_type.clone()) |
321 | 0 | .len(len) |
322 | 0 | .add_buffer(values_buffer) |
323 | 0 | .nulls(nulls); |
324 | | |
325 | 0 | let array_data = unsafe { builder.build_unchecked() }; |
326 | 0 | PrimitiveArray::<T>::from(array_data) |
327 | 0 | } |
328 | | |
329 | | /// Returns the current values buffer as a slice |
330 | | pub fn values_slice(&self) -> &[T::Native] { |
331 | | self.values_builder.as_slice() |
332 | | } |
333 | | |
334 | | /// Returns the current values buffer as a mutable slice |
335 | | pub fn values_slice_mut(&mut self) -> &mut [T::Native] { |
336 | | self.values_builder.as_mut_slice() |
337 | | } |
338 | | |
339 | | /// Returns the current null buffer as a slice |
340 | | pub fn validity_slice(&self) -> Option<&[u8]> { |
341 | | self.null_buffer_builder.as_slice() |
342 | | } |
343 | | |
344 | | /// Returns the current null buffer as a mutable slice |
345 | | pub fn validity_slice_mut(&mut self) -> Option<&mut [u8]> { |
346 | | self.null_buffer_builder.as_slice_mut() |
347 | | } |
348 | | |
349 | | /// Returns the current values buffer and null buffer as a slice |
350 | | pub fn slices_mut(&mut self) -> (&mut [T::Native], Option<&mut [u8]>) { |
351 | | ( |
352 | | self.values_builder.as_mut_slice(), |
353 | | self.null_buffer_builder.as_slice_mut(), |
354 | | ) |
355 | | } |
356 | | } |
357 | | |
358 | | impl<P: DecimalType> PrimitiveBuilder<P> { |
359 | | /// Sets the precision and scale |
360 | 12 | pub fn with_precision_and_scale(self, precision: u8, scale: i8) -> Result<Self, ArrowError> { |
361 | 12 | validate_decimal_precision_and_scale::<P>(precision, scale)?0 ; |
362 | 12 | Ok(Self { |
363 | 12 | data_type: P::TYPE_CONSTRUCTOR(precision, scale), |
364 | 12 | ..self |
365 | 12 | }) |
366 | 12 | } |
367 | | } |
368 | | |
369 | | impl<P: ArrowTimestampType> PrimitiveBuilder<P> { |
370 | | /// Sets the timezone |
371 | | pub fn with_timezone(self, timezone: impl Into<Arc<str>>) -> Self { |
372 | | self.with_timezone_opt(Some(timezone.into())) |
373 | | } |
374 | | |
375 | | /// Sets an optional timezone |
376 | | pub fn with_timezone_opt<S: Into<Arc<str>>>(self, timezone: Option<S>) -> Self { |
377 | | Self { |
378 | | data_type: DataType::Timestamp(P::UNIT, timezone.map(Into::into)), |
379 | | ..self |
380 | | } |
381 | | } |
382 | | } |
383 | | |
384 | | impl<P: ArrowPrimitiveType> Extend<Option<P::Native>> for PrimitiveBuilder<P> { |
385 | | #[inline] |
386 | | fn extend<T: IntoIterator<Item = Option<P::Native>>>(&mut self, iter: T) { |
387 | | for v in iter { |
388 | | self.append_option(v) |
389 | | } |
390 | | } |
391 | | } |
392 | | |
393 | | #[cfg(test)] |
394 | | mod tests { |
395 | | use super::*; |
396 | | use arrow_buffer::{NullBuffer, ScalarBuffer}; |
397 | | use arrow_schema::TimeUnit; |
398 | | |
399 | | use crate::array::Array; |
400 | | use crate::array::BooleanArray; |
401 | | use crate::array::Date32Array; |
402 | | use crate::array::Int32Array; |
403 | | use crate::array::TimestampSecondArray; |
404 | | |
405 | | #[test] |
406 | | fn test_primitive_array_builder_i32() { |
407 | | let mut builder = Int32Array::builder(5); |
408 | | for i in 0..5 { |
409 | | builder.append_value(i); |
410 | | } |
411 | | let arr = builder.finish(); |
412 | | assert_eq!(5, arr.len()); |
413 | | assert_eq!(0, arr.offset()); |
414 | | assert_eq!(0, arr.null_count()); |
415 | | for i in 0..5 { |
416 | | assert!(!arr.is_null(i)); |
417 | | assert!(arr.is_valid(i)); |
418 | | assert_eq!(i as i32, arr.value(i)); |
419 | | } |
420 | | } |
421 | | |
422 | | #[test] |
423 | | fn test_primitive_array_builder_i32_append_iter() { |
424 | | let mut builder = Int32Array::builder(5); |
425 | | unsafe { builder.append_trusted_len_iter(0..5) }; |
426 | | let arr = builder.finish(); |
427 | | assert_eq!(5, arr.len()); |
428 | | assert_eq!(0, arr.offset()); |
429 | | assert_eq!(0, arr.null_count()); |
430 | | for i in 0..5 { |
431 | | assert!(!arr.is_null(i)); |
432 | | assert!(arr.is_valid(i)); |
433 | | assert_eq!(i as i32, arr.value(i)); |
434 | | } |
435 | | } |
436 | | |
437 | | #[test] |
438 | | fn test_primitive_array_builder_i32_append_nulls() { |
439 | | let mut builder = Int32Array::builder(5); |
440 | | builder.append_nulls(5); |
441 | | let arr = builder.finish(); |
442 | | assert_eq!(5, arr.len()); |
443 | | assert_eq!(0, arr.offset()); |
444 | | assert_eq!(5, arr.null_count()); |
445 | | for i in 0..5 { |
446 | | assert!(arr.is_null(i)); |
447 | | assert!(!arr.is_valid(i)); |
448 | | } |
449 | | } |
450 | | |
451 | | #[test] |
452 | | fn test_primitive_array_builder_date32() { |
453 | | let mut builder = Date32Array::builder(5); |
454 | | for i in 0..5 { |
455 | | builder.append_value(i); |
456 | | } |
457 | | let arr = builder.finish(); |
458 | | assert_eq!(5, arr.len()); |
459 | | assert_eq!(0, arr.offset()); |
460 | | assert_eq!(0, arr.null_count()); |
461 | | for i in 0..5 { |
462 | | assert!(!arr.is_null(i)); |
463 | | assert!(arr.is_valid(i)); |
464 | | assert_eq!(i as i32, arr.value(i)); |
465 | | } |
466 | | } |
467 | | |
468 | | #[test] |
469 | | fn test_primitive_array_builder_timestamp_second() { |
470 | | let mut builder = TimestampSecondArray::builder(5); |
471 | | for i in 0..5 { |
472 | | builder.append_value(i); |
473 | | } |
474 | | let arr = builder.finish(); |
475 | | assert_eq!(5, arr.len()); |
476 | | assert_eq!(0, arr.offset()); |
477 | | assert_eq!(0, arr.null_count()); |
478 | | for i in 0..5 { |
479 | | assert!(!arr.is_null(i)); |
480 | | assert!(arr.is_valid(i)); |
481 | | assert_eq!(i as i64, arr.value(i)); |
482 | | } |
483 | | } |
484 | | |
485 | | #[test] |
486 | | fn test_primitive_array_builder_bool() { |
487 | | // 00000010 01001000 |
488 | | let buf = Buffer::from([72_u8, 2_u8]); |
489 | | let mut builder = BooleanArray::builder(10); |
490 | | for i in 0..10 { |
491 | | if i == 3 || i == 6 || i == 9 { |
492 | | builder.append_value(true); |
493 | | } else { |
494 | | builder.append_value(false); |
495 | | } |
496 | | } |
497 | | |
498 | | let arr = builder.finish(); |
499 | | assert_eq!(&buf, arr.values().inner()); |
500 | | assert_eq!(10, arr.len()); |
501 | | assert_eq!(0, arr.offset()); |
502 | | assert_eq!(0, arr.null_count()); |
503 | | for i in 0..10 { |
504 | | assert!(!arr.is_null(i)); |
505 | | assert!(arr.is_valid(i)); |
506 | | assert_eq!(i == 3 || i == 6 || i == 9, arr.value(i), "failed at {i}") |
507 | | } |
508 | | } |
509 | | |
510 | | #[test] |
511 | | fn test_primitive_array_builder_append_option() { |
512 | | let arr1 = Int32Array::from(vec![Some(0), None, Some(2), None, Some(4)]); |
513 | | |
514 | | let mut builder = Int32Array::builder(5); |
515 | | builder.append_option(Some(0)); |
516 | | builder.append_option(None); |
517 | | builder.append_option(Some(2)); |
518 | | builder.append_option(None); |
519 | | builder.append_option(Some(4)); |
520 | | let arr2 = builder.finish(); |
521 | | |
522 | | assert_eq!(arr1.len(), arr2.len()); |
523 | | assert_eq!(arr1.offset(), arr2.offset()); |
524 | | assert_eq!(arr1.null_count(), arr2.null_count()); |
525 | | for i in 0..5 { |
526 | | assert_eq!(arr1.is_null(i), arr2.is_null(i)); |
527 | | assert_eq!(arr1.is_valid(i), arr2.is_valid(i)); |
528 | | if arr1.is_valid(i) { |
529 | | assert_eq!(arr1.value(i), arr2.value(i)); |
530 | | } |
531 | | } |
532 | | } |
533 | | |
534 | | #[test] |
535 | | fn test_primitive_array_builder_append_null() { |
536 | | let arr1 = Int32Array::from(vec![Some(0), Some(2), None, None, Some(4)]); |
537 | | |
538 | | let mut builder = Int32Array::builder(5); |
539 | | builder.append_value(0); |
540 | | builder.append_value(2); |
541 | | builder.append_null(); |
542 | | builder.append_null(); |
543 | | builder.append_value(4); |
544 | | let arr2 = builder.finish(); |
545 | | |
546 | | assert_eq!(arr1.len(), arr2.len()); |
547 | | assert_eq!(arr1.offset(), arr2.offset()); |
548 | | assert_eq!(arr1.null_count(), arr2.null_count()); |
549 | | for i in 0..5 { |
550 | | assert_eq!(arr1.is_null(i), arr2.is_null(i)); |
551 | | assert_eq!(arr1.is_valid(i), arr2.is_valid(i)); |
552 | | if arr1.is_valid(i) { |
553 | | assert_eq!(arr1.value(i), arr2.value(i)); |
554 | | } |
555 | | } |
556 | | } |
557 | | |
558 | | #[test] |
559 | | fn test_primitive_array_builder_append_slice() { |
560 | | let arr1 = Int32Array::from(vec![Some(0), Some(2), None, None, Some(4)]); |
561 | | |
562 | | let mut builder = Int32Array::builder(5); |
563 | | builder.append_slice(&[0, 2]); |
564 | | builder.append_null(); |
565 | | builder.append_null(); |
566 | | builder.append_value(4); |
567 | | let arr2 = builder.finish(); |
568 | | |
569 | | assert_eq!(arr1.len(), arr2.len()); |
570 | | assert_eq!(arr1.offset(), arr2.offset()); |
571 | | assert_eq!(arr1.null_count(), arr2.null_count()); |
572 | | for i in 0..5 { |
573 | | assert_eq!(arr1.is_null(i), arr2.is_null(i)); |
574 | | assert_eq!(arr1.is_valid(i), arr2.is_valid(i)); |
575 | | if arr1.is_valid(i) { |
576 | | assert_eq!(arr1.value(i), arr2.value(i)); |
577 | | } |
578 | | } |
579 | | } |
580 | | |
581 | | #[test] |
582 | | fn test_primitive_array_builder_finish() { |
583 | | let mut builder = Int32Builder::new(); |
584 | | builder.append_slice(&[2, 4, 6, 8]); |
585 | | let mut arr = builder.finish(); |
586 | | assert_eq!(4, arr.len()); |
587 | | assert_eq!(0, builder.len()); |
588 | | |
589 | | builder.append_slice(&[1, 3, 5, 7, 9]); |
590 | | arr = builder.finish(); |
591 | | assert_eq!(5, arr.len()); |
592 | | assert_eq!(0, builder.len()); |
593 | | } |
594 | | |
595 | | #[test] |
596 | | fn test_primitive_array_builder_finish_cloned() { |
597 | | let mut builder = Int32Builder::new(); |
598 | | builder.append_value(23); |
599 | | builder.append_value(45); |
600 | | let result = builder.finish_cloned(); |
601 | | assert_eq!(result, Int32Array::from(vec![23, 45])); |
602 | | builder.append_value(56); |
603 | | assert_eq!(builder.finish_cloned(), Int32Array::from(vec![23, 45, 56])); |
604 | | |
605 | | builder.append_slice(&[2, 4, 6, 8]); |
606 | | let mut arr = builder.finish(); |
607 | | assert_eq!(7, arr.len()); |
608 | | assert_eq!(arr, Int32Array::from(vec![23, 45, 56, 2, 4, 6, 8])); |
609 | | assert_eq!(0, builder.len()); |
610 | | |
611 | | builder.append_slice(&[1, 3, 5, 7, 9]); |
612 | | arr = builder.finish(); |
613 | | assert_eq!(5, arr.len()); |
614 | | assert_eq!(0, builder.len()); |
615 | | } |
616 | | |
617 | | #[test] |
618 | | fn test_primitive_array_builder_with_data_type() { |
619 | | let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(1, 2)); |
620 | | builder.append_value(1); |
621 | | let array = builder.finish(); |
622 | | assert_eq!(array.precision(), 1); |
623 | | assert_eq!(array.scale(), 2); |
624 | | |
625 | | let data_type = DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())); |
626 | | let mut builder = TimestampNanosecondBuilder::new().with_data_type(data_type.clone()); |
627 | | builder.append_value(1); |
628 | | let array = builder.finish(); |
629 | | assert_eq!(array.data_type(), &data_type); |
630 | | } |
631 | | |
632 | | #[test] |
633 | | #[should_panic(expected = "incompatible data type for builder, expected Int32 got Int64")] |
634 | | fn test_invalid_with_data_type() { |
635 | | Int32Builder::new().with_data_type(DataType::Int64); |
636 | | } |
637 | | |
638 | | #[test] |
639 | | fn test_extend() { |
640 | | let mut builder = PrimitiveBuilder::<Int16Type>::new(); |
641 | | builder.extend([1, 2, 3, 5, 2, 4, 4].into_iter().map(Some)); |
642 | | builder.extend([2, 4, 6, 2].into_iter().map(Some)); |
643 | | let array = builder.finish(); |
644 | | assert_eq!(array.values(), &[1, 2, 3, 5, 2, 4, 4, 2, 4, 6, 2]); |
645 | | } |
646 | | |
647 | | #[test] |
648 | | fn test_primitive_array_append_array() { |
649 | | let input = vec![ |
650 | | Some(1), |
651 | | None, |
652 | | Some(3), |
653 | | None, |
654 | | Some(5), |
655 | | None, |
656 | | None, |
657 | | None, |
658 | | Some(7), |
659 | | Some(9), |
660 | | Some(8), |
661 | | Some(6), |
662 | | Some(4), |
663 | | ]; |
664 | | let arr1 = Int32Array::from(input[..5].to_vec()); |
665 | | let arr2 = Int32Array::from(input[5..8].to_vec()); |
666 | | let arr3 = Int32Array::from(input[8..].to_vec()); |
667 | | |
668 | | let mut builder = Int32Array::builder(5); |
669 | | builder.append_array(&arr1); |
670 | | builder.append_array(&arr2); |
671 | | builder.append_array(&arr3); |
672 | | let actual = builder.finish(); |
673 | | let expected = Int32Array::from(input); |
674 | | |
675 | | assert_eq!(actual, expected); |
676 | | } |
677 | | |
678 | | #[test] |
679 | | fn test_append_array_add_underlying_null_values() { |
680 | | let array = Int32Array::new( |
681 | | ScalarBuffer::from(vec![2, 3, 4, 5]), |
682 | | Some(NullBuffer::from(&[true, true, false, false])), |
683 | | ); |
684 | | |
685 | | let mut builder = Int32Array::builder(5); |
686 | | builder.append_array(&array); |
687 | | let actual = builder.finish(); |
688 | | |
689 | | assert_eq!(actual, array); |
690 | | assert_eq!(actual.values(), array.values()) |
691 | | } |
692 | | |
693 | | #[test] |
694 | | #[should_panic(expected = "array data type mismatch")] |
695 | | fn test_invalid_with_data_type_in_append_array() { |
696 | | let array = { |
697 | | let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(1, 2)); |
698 | | builder.append_value(1); |
699 | | builder.finish() |
700 | | }; |
701 | | |
702 | | let mut builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(2, 3)); |
703 | | builder.append_array(&array) |
704 | | } |
705 | | } |