1use std::fmt;
8
9use crate::schema::{TimeUnit, TimeZone};
10
11#[derive(Clone, PartialEq, Eq)]
13pub struct PrimitiveArray<T> {
14 values: Vec<T>,
15 validity: Option<Vec<bool>>,
16}
17
18impl<T> PrimitiveArray<T> {
19 pub fn new(values: Vec<T>, validity: Option<Vec<bool>>) -> Self {
21 Self { values, validity }
22 }
23
24 pub fn len(&self) -> usize {
26 self.values.len()
27 }
28
29 pub fn is_empty(&self) -> bool {
31 self.values.is_empty()
32 }
33
34 pub fn values(&self) -> &[T] {
36 &self.values
37 }
38
39 pub fn validity(&self) -> Option<&[bool]> {
41 self.validity.as_deref()
42 }
43
44 pub fn is_null(&self, index: usize) -> bool {
50 is_null_at(self.validity.as_deref(), index)
51 }
52}
53
54fn is_null_at(validity: Option<&[bool]>, index: usize) -> bool {
57 validity.is_some_and(|validity| validity.get(index) == Some(&false))
58}
59
60impl<T: fmt::Debug> fmt::Debug for PrimitiveArray<T> {
61 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
62 formatter
63 .debug_struct("PrimitiveArray")
64 .field("values", &self.values)
65 .field("validity", &self.validity)
66 .finish()
67 }
68}
69
70pub type BooleanArray = PrimitiveArray<bool>;
72
73#[derive(Debug, Clone, PartialEq, Eq)]
75pub struct Utf8Array {
76 values: Vec<String>,
77 validity: Option<Vec<bool>>,
78}
79
80impl Utf8Array {
81 pub fn new(values: Vec<String>, validity: Option<Vec<bool>>) -> Self {
83 Self { values, validity }
84 }
85
86 pub fn len(&self) -> usize {
87 self.values.len()
88 }
89
90 pub fn is_empty(&self) -> bool {
91 self.values.is_empty()
92 }
93
94 pub fn values(&self) -> &[String] {
95 &self.values
96 }
97
98 pub fn validity(&self) -> Option<&[bool]> {
99 self.validity.as_deref()
100 }
101
102 pub fn is_null(&self, index: usize) -> bool {
104 is_null_at(self.validity.as_deref(), index)
105 }
106}
107
108#[derive(Debug, Clone, PartialEq, Eq)]
111pub struct BinaryArray {
112 values: Vec<Vec<u8>>,
113 validity: Option<Vec<bool>>,
114}
115
116impl BinaryArray {
117 pub fn new(values: Vec<Vec<u8>>, validity: Option<Vec<bool>>) -> Self {
120 Self { values, validity }
121 }
122
123 pub fn len(&self) -> usize {
124 self.values.len()
125 }
126
127 pub fn is_empty(&self) -> bool {
128 self.values.is_empty()
129 }
130
131 pub fn values(&self) -> &[Vec<u8>] {
132 &self.values
133 }
134
135 pub fn validity(&self) -> Option<&[bool]> {
136 self.validity.as_deref()
137 }
138
139 pub fn is_null(&self, index: usize) -> bool {
141 is_null_at(self.validity.as_deref(), index)
142 }
143}
144
145#[derive(Debug, Clone, PartialEq)]
148#[non_exhaustive]
149pub enum Array {
150 Bool(BooleanArray),
151 Int8(PrimitiveArray<i8>),
152 Int16(PrimitiveArray<i16>),
153 Int32(PrimitiveArray<i32>),
154 Int64(PrimitiveArray<i64>),
155 UInt8(PrimitiveArray<u8>),
156 UInt16(PrimitiveArray<u16>),
157 UInt32(PrimitiveArray<u32>),
158 UInt64(PrimitiveArray<u64>),
159 Float32(PrimitiveArray<f32>),
160 Float64(PrimitiveArray<f64>),
161 Decimal(DecimalArray),
162 Timestamp(TimestampArray),
163 Utf8(Utf8Array),
164 Categorical(Utf8Array),
165 Binary(BinaryArray),
166 FixedBinary(BinaryArray),
167 Date32(PrimitiveArray<i32>),
168}
169
170#[derive(Debug, Clone, PartialEq, Eq)]
172pub struct DecimalArray {
173 values: PrimitiveArray<i64>,
174 precision: u16,
175 scale: i16,
176}
177
178impl DecimalArray {
179 pub fn new(values: Vec<i64>, validity: Option<Vec<bool>>, precision: u16, scale: i16) -> Self {
181 Self {
182 values: PrimitiveArray::new(values, validity),
183 precision,
184 scale,
185 }
186 }
187
188 pub fn values(&self) -> &[i64] {
189 self.values.values()
190 }
191
192 pub fn validity(&self) -> Option<&[bool]> {
193 self.values.validity()
194 }
195
196 pub fn precision(&self) -> u16 {
197 self.precision
198 }
199
200 pub fn scale(&self) -> i16 {
201 self.scale
202 }
203
204 pub fn len(&self) -> usize {
205 self.values.len()
206 }
207
208 pub fn is_empty(&self) -> bool {
209 self.values.is_empty()
210 }
211
212 pub fn is_null(&self, index: usize) -> bool {
213 self.values.is_null(index)
214 }
215}
216
217#[derive(Debug, Clone, PartialEq, Eq)]
219pub struct TimestampArray {
220 values: PrimitiveArray<i64>,
221 unit: TimeUnit,
222 timezone: TimeZone,
223}
224
225impl TimestampArray {
226 pub fn new(
228 values: Vec<i64>,
229 validity: Option<Vec<bool>>,
230 unit: TimeUnit,
231 timezone: TimeZone,
232 ) -> Self {
233 Self {
234 values: PrimitiveArray::new(values, validity),
235 unit,
236 timezone,
237 }
238 }
239
240 pub fn values(&self) -> &[i64] {
241 self.values.values()
242 }
243
244 pub fn validity(&self) -> Option<&[bool]> {
245 self.values.validity()
246 }
247
248 pub fn unit(&self) -> TimeUnit {
249 self.unit
250 }
251
252 pub fn timezone(&self) -> &TimeZone {
253 &self.timezone
254 }
255
256 pub fn len(&self) -> usize {
257 self.values.len()
258 }
259
260 pub fn is_empty(&self) -> bool {
261 self.values.is_empty()
262 }
263
264 pub fn is_null(&self, index: usize) -> bool {
265 self.values.is_null(index)
266 }
267}
268
269impl Array {
270 pub fn len(&self) -> usize {
272 match self {
273 Self::Bool(array) => array.len(),
274 Self::Int8(array) => array.len(),
275 Self::Int16(array) => array.len(),
276 Self::Int32(array) => array.len(),
277 Self::Int64(array) => array.len(),
278 Self::UInt8(array) => array.len(),
279 Self::UInt16(array) => array.len(),
280 Self::UInt32(array) => array.len(),
281 Self::UInt64(array) => array.len(),
282 Self::Float32(array) => array.len(),
283 Self::Float64(array) => array.len(),
284 Self::Decimal(array) => array.len(),
285 Self::Timestamp(array) => array.len(),
286 Self::Utf8(array) | Self::Categorical(array) => array.len(),
287 Self::Binary(array) | Self::FixedBinary(array) => array.len(),
288 Self::Date32(array) => array.len(),
289 }
290 }
291
292 pub fn is_empty(&self) -> bool {
294 self.len() == 0
295 }
296
297 pub fn is_null(&self, index: usize) -> bool {
300 match self {
301 Self::Bool(array) => array.is_null(index),
302 Self::Int8(array) => array.is_null(index),
303 Self::Int16(array) => array.is_null(index),
304 Self::Int32(array) => array.is_null(index),
305 Self::Int64(array) => array.is_null(index),
306 Self::UInt8(array) => array.is_null(index),
307 Self::UInt16(array) => array.is_null(index),
308 Self::UInt32(array) => array.is_null(index),
309 Self::UInt64(array) => array.is_null(index),
310 Self::Float32(array) => array.is_null(index),
311 Self::Float64(array) => array.is_null(index),
312 Self::Decimal(array) => array.is_null(index),
313 Self::Timestamp(array) => array.is_null(index),
314 Self::Utf8(array) | Self::Categorical(array) => array.is_null(index),
315 Self::Binary(array) | Self::FixedBinary(array) => array.is_null(index),
316 Self::Date32(array) => array.is_null(index),
317 }
318 }
319
320 pub fn value_at(&self, index: usize) -> Option<ScalarValue<'_>> {
322 if index >= self.len() {
323 return None;
324 }
325 if self.is_null(index) {
326 return None;
327 }
328 Some(match self {
329 Self::Bool(array) => ScalarValue::Bool(array.values()[index]),
330 Self::Int8(array) => ScalarValue::Int8(array.values()[index]),
331 Self::Int16(array) => ScalarValue::Int16(array.values()[index]),
332 Self::Int32(array) => ScalarValue::Int32(array.values()[index]),
333 Self::Int64(array) => ScalarValue::Int64(array.values()[index]),
334 Self::UInt8(array) => ScalarValue::UInt8(array.values()[index]),
335 Self::UInt16(array) => ScalarValue::UInt16(array.values()[index]),
336 Self::UInt32(array) => ScalarValue::UInt32(array.values()[index]),
337 Self::UInt64(array) => ScalarValue::UInt64(array.values()[index]),
338 Self::Float32(array) => ScalarValue::Float32(array.values()[index]),
339 Self::Float64(array) => ScalarValue::Float64(array.values()[index]),
340 Self::Decimal(array) => ScalarValue::Decimal {
341 unscaled: array.values()[index],
342 precision: array.precision(),
343 scale: array.scale(),
344 },
345 Self::Timestamp(array) => ScalarValue::Timestamp {
346 value: array.values()[index],
347 unit: array.unit(),
348 timezone: array.timezone(),
349 },
350 Self::Utf8(array) => ScalarValue::Utf8(&array.values()[index]),
351 Self::Categorical(array) => ScalarValue::Categorical(&array.values()[index]),
352 Self::Binary(array) => ScalarValue::Binary(&array.values()[index]),
353 Self::FixedBinary(array) => ScalarValue::FixedBinary(&array.values()[index]),
354 Self::Date32(array) => ScalarValue::Date32(array.values()[index]),
355 })
356 }
357
358 pub(crate) fn slice(&self, start: usize, end: usize) -> Self {
362 assert!(start <= end && end <= self.len());
363 match self {
364 Self::Bool(array) => Self::Bool(PrimitiveArray::new(
365 array.values()[start..end].to_vec(),
366 slice_validity(array.validity(), start, end),
367 )),
368 Self::Int8(array) => Self::Int8(PrimitiveArray::new(
369 array.values()[start..end].to_vec(),
370 slice_validity(array.validity(), start, end),
371 )),
372 Self::Int16(array) => Self::Int16(PrimitiveArray::new(
373 array.values()[start..end].to_vec(),
374 slice_validity(array.validity(), start, end),
375 )),
376 Self::Int32(array) => Self::Int32(PrimitiveArray::new(
377 array.values()[start..end].to_vec(),
378 slice_validity(array.validity(), start, end),
379 )),
380 Self::Int64(array) => Self::Int64(PrimitiveArray::new(
381 array.values()[start..end].to_vec(),
382 slice_validity(array.validity(), start, end),
383 )),
384 Self::UInt8(array) => Self::UInt8(PrimitiveArray::new(
385 array.values()[start..end].to_vec(),
386 slice_validity(array.validity(), start, end),
387 )),
388 Self::UInt16(array) => Self::UInt16(PrimitiveArray::new(
389 array.values()[start..end].to_vec(),
390 slice_validity(array.validity(), start, end),
391 )),
392 Self::UInt32(array) => Self::UInt32(PrimitiveArray::new(
393 array.values()[start..end].to_vec(),
394 slice_validity(array.validity(), start, end),
395 )),
396 Self::UInt64(array) => Self::UInt64(PrimitiveArray::new(
397 array.values()[start..end].to_vec(),
398 slice_validity(array.validity(), start, end),
399 )),
400 Self::Float32(array) => Self::Float32(PrimitiveArray::new(
401 array.values()[start..end].to_vec(),
402 slice_validity(array.validity(), start, end),
403 )),
404 Self::Float64(array) => Self::Float64(PrimitiveArray::new(
405 array.values()[start..end].to_vec(),
406 slice_validity(array.validity(), start, end),
407 )),
408 Self::Decimal(array) => Self::Decimal(DecimalArray::new(
409 array.values()[start..end].to_vec(),
410 slice_validity(array.validity(), start, end),
411 array.precision(),
412 array.scale(),
413 )),
414 Self::Timestamp(array) => Self::Timestamp(TimestampArray::new(
415 array.values()[start..end].to_vec(),
416 slice_validity(array.validity(), start, end),
417 array.unit(),
418 array.timezone().clone(),
419 )),
420 Self::Utf8(array) => Self::Utf8(Utf8Array::new(
421 array.values()[start..end].to_vec(),
422 slice_validity(array.validity(), start, end),
423 )),
424 Self::Categorical(array) => Self::Categorical(Utf8Array::new(
425 array.values()[start..end].to_vec(),
426 slice_validity(array.validity(), start, end),
427 )),
428 Self::Binary(array) => Self::Binary(BinaryArray::new(
429 array.values()[start..end].to_vec(),
430 slice_validity(array.validity(), start, end),
431 )),
432 Self::FixedBinary(array) => Self::FixedBinary(BinaryArray::new(
433 array.values()[start..end].to_vec(),
434 slice_validity(array.validity(), start, end),
435 )),
436 Self::Date32(array) => Self::Date32(PrimitiveArray::new(
437 array.values()[start..end].to_vec(),
438 slice_validity(array.validity(), start, end),
439 )),
440 }
441 }
442
443 pub(crate) fn take(&self, indices: &[usize]) -> Self {
447 assert!(indices.iter().all(|&index| index < self.len()));
448 match self {
449 Self::Bool(array) => Self::Bool(PrimitiveArray::new(
450 take_values(array.values(), indices),
451 take_validity(array.validity(), indices),
452 )),
453 Self::Int8(array) => Self::Int8(PrimitiveArray::new(
454 take_values(array.values(), indices),
455 take_validity(array.validity(), indices),
456 )),
457 Self::Int16(array) => Self::Int16(PrimitiveArray::new(
458 take_values(array.values(), indices),
459 take_validity(array.validity(), indices),
460 )),
461 Self::Int32(array) => Self::Int32(PrimitiveArray::new(
462 take_values(array.values(), indices),
463 take_validity(array.validity(), indices),
464 )),
465 Self::Int64(array) => Self::Int64(PrimitiveArray::new(
466 take_values(array.values(), indices),
467 take_validity(array.validity(), indices),
468 )),
469 Self::UInt8(array) => Self::UInt8(PrimitiveArray::new(
470 take_values(array.values(), indices),
471 take_validity(array.validity(), indices),
472 )),
473 Self::UInt16(array) => Self::UInt16(PrimitiveArray::new(
474 take_values(array.values(), indices),
475 take_validity(array.validity(), indices),
476 )),
477 Self::UInt32(array) => Self::UInt32(PrimitiveArray::new(
478 take_values(array.values(), indices),
479 take_validity(array.validity(), indices),
480 )),
481 Self::UInt64(array) => Self::UInt64(PrimitiveArray::new(
482 take_values(array.values(), indices),
483 take_validity(array.validity(), indices),
484 )),
485 Self::Float32(array) => Self::Float32(PrimitiveArray::new(
486 take_values(array.values(), indices),
487 take_validity(array.validity(), indices),
488 )),
489 Self::Float64(array) => Self::Float64(PrimitiveArray::new(
490 take_values(array.values(), indices),
491 take_validity(array.validity(), indices),
492 )),
493 Self::Decimal(array) => Self::Decimal(DecimalArray::new(
494 take_values(array.values(), indices),
495 take_validity(array.validity(), indices),
496 array.precision(),
497 array.scale(),
498 )),
499 Self::Timestamp(array) => Self::Timestamp(TimestampArray::new(
500 take_values(array.values(), indices),
501 take_validity(array.validity(), indices),
502 array.unit(),
503 array.timezone().clone(),
504 )),
505 Self::Utf8(array) => Self::Utf8(Utf8Array::new(
506 take_values(array.values(), indices),
507 take_validity(array.validity(), indices),
508 )),
509 Self::Categorical(array) => Self::Categorical(Utf8Array::new(
510 take_values(array.values(), indices),
511 take_validity(array.validity(), indices),
512 )),
513 Self::Binary(array) => Self::Binary(BinaryArray::new(
514 take_values(array.values(), indices),
515 take_validity(array.validity(), indices),
516 )),
517 Self::FixedBinary(array) => Self::FixedBinary(BinaryArray::new(
518 take_values(array.values(), indices),
519 take_validity(array.validity(), indices),
520 )),
521 Self::Date32(array) => Self::Date32(PrimitiveArray::new(
522 take_values(array.values(), indices),
523 take_validity(array.validity(), indices),
524 )),
525 }
526 }
527}
528
529fn slice_validity(validity: Option<&[bool]>, start: usize, end: usize) -> Option<Vec<bool>> {
530 validity.map(|validity| validity[start..end].to_vec())
531}
532
533fn take_values<T: Clone>(values: &[T], indices: &[usize]) -> Vec<T> {
534 indices.iter().map(|&index| values[index].clone()).collect()
535}
536
537fn take_validity(validity: Option<&[bool]>, indices: &[usize]) -> Option<Vec<bool>> {
538 validity.map(|validity| indices.iter().map(|&index| validity[index]).collect())
539}
540
541#[derive(Debug, Clone, Copy, PartialEq)]
543#[non_exhaustive]
544pub enum ScalarValue<'a> {
545 Bool(bool),
546 Int8(i8),
547 Int16(i16),
548 Int32(i32),
549 Int64(i64),
550 UInt8(u8),
551 UInt16(u16),
552 UInt32(u32),
553 UInt64(u64),
554 Float32(f32),
555 Float64(f64),
556 Decimal {
557 unscaled: i64,
558 precision: u16,
559 scale: i16,
560 },
561 Timestamp {
562 value: i64,
563 unit: TimeUnit,
564 timezone: &'a TimeZone,
565 },
566 Utf8(&'a str),
567 Categorical(&'a str),
568 Binary(&'a [u8]),
569 FixedBinary(&'a [u8]),
570 Date32(i32),
571}
572
573#[cfg(test)]
574mod tests {
575 use super::*;
576
577 fn mixed_int32() -> Array {
578 Array::Int32(PrimitiveArray::new(
579 vec![10, 20, 30, 40],
580 Some(vec![true, false, true, false]),
581 ))
582 }
583
584 #[test]
585 fn a_slice_keeps_the_values_of_its_row_range() {
586 let sliced = mixed_int32().slice(1, 3);
587
588 assert_eq!(sliced.value_at(1), Some(ScalarValue::Int32(30)));
589 }
590
591 #[test]
592 fn a_slice_keeps_the_nulls_of_its_row_range() {
593 let sliced = mixed_int32().slice(1, 3);
594
595 assert_eq!((sliced.is_null(0), sliced.is_null(1)), (true, false));
596 }
597
598 #[test]
599 fn a_slice_of_a_bitmap_free_array_stays_bitmap_free() {
600 let array = Array::Int64(PrimitiveArray::new(vec![1, 2, 3], None));
601
602 let Array::Int64(sliced) = array.slice(0, 2) else {
603 panic!("slicing preserves the logical type");
604 };
605 assert_eq!(sliced.validity(), None);
606 }
607
608 #[test]
609 fn an_empty_slice_has_no_rows() {
610 assert_eq!(mixed_int32().slice(2, 2).len(), 0);
611 }
612
613 #[test]
614 fn a_variable_width_slice_keeps_its_own_values() {
615 let array = Array::Utf8(Utf8Array::new(
616 vec!["a".into(), "bb".into(), "ccc".into()],
617 None,
618 ));
619
620 assert_eq!(array.slice(1, 3).value_at(0), Some(ScalarValue::Utf8("bb")));
621 }
622
623 #[test]
624 fn a_decimal_slice_keeps_its_precision_and_scale() {
625 let array = Array::Decimal(DecimalArray::new(vec![1, 2, 3], None, 10, 2));
626
627 assert_eq!(
628 array.slice(0, 1).value_at(0),
629 Some(ScalarValue::Decimal {
630 unscaled: 1,
631 precision: 10,
632 scale: 2
633 })
634 );
635 }
636
637 #[test]
638 fn a_timestamp_slice_keeps_its_unit_and_timezone() {
639 let array = Array::Timestamp(TimestampArray::new(
640 vec![5, 6],
641 None,
642 TimeUnit::Millisecond,
643 TimeZone::Utc,
644 ));
645
646 assert_eq!(
647 array.slice(1, 2).value_at(0),
648 Some(ScalarValue::Timestamp {
649 value: 6,
650 unit: TimeUnit::Millisecond,
651 timezone: &TimeZone::Utc
652 })
653 );
654 }
655
656 #[test]
657 fn every_logical_type_survives_a_slice() {
658 let arrays = [
659 Array::Bool(BooleanArray::new(vec![true, false], None)),
660 Array::Int8(PrimitiveArray::new(vec![1, 2], None)),
661 Array::Int16(PrimitiveArray::new(vec![1, 2], None)),
662 Array::Int32(PrimitiveArray::new(vec![1, 2], None)),
663 Array::Int64(PrimitiveArray::new(vec![1, 2], None)),
664 Array::UInt8(PrimitiveArray::new(vec![1, 2], None)),
665 Array::UInt16(PrimitiveArray::new(vec![1, 2], None)),
666 Array::UInt32(PrimitiveArray::new(vec![1, 2], None)),
667 Array::UInt64(PrimitiveArray::new(vec![1, 2], None)),
668 Array::Float32(PrimitiveArray::new(vec![1.0, 2.0], None)),
669 Array::Float64(PrimitiveArray::new(vec![1.0, 2.0], None)),
670 Array::Decimal(DecimalArray::new(vec![1, 2], None, 5, 1)),
671 Array::Timestamp(TimestampArray::new(
672 vec![1, 2],
673 None,
674 TimeUnit::Second,
675 TimeZone::Naive,
676 )),
677 Array::Utf8(Utf8Array::new(vec!["a".into(), "b".into()], None)),
678 Array::Categorical(Utf8Array::new(vec!["a".into(), "b".into()], None)),
679 Array::Binary(BinaryArray::new(vec![vec![1], vec![2]], None)),
680 Array::FixedBinary(BinaryArray::new(vec![vec![1], vec![2]], None)),
681 Array::Date32(PrimitiveArray::new(vec![1, 2], None)),
682 ];
683
684 for array in arrays {
685 let sliced = array.slice(1, 2);
686 assert_eq!(sliced.len(), 1, "{array:?} lost its row");
687 assert_eq!(sliced.value_at(0), array.value_at(1), "{array:?}");
688 }
689 }
690}