1use std::sync::Arc;
15
16use arrow_array::builder::StringBuilder;
17use arrow_array::cast::AsArray;
18use arrow_array::types::{Decimal128Type, Int32Type};
19use arrow_array::{
20 Array, ArrayRef, BooleanArray, Decimal128Array, DictionaryArray, StringArray, UInt32Array,
21};
22use arrow_buffer::{NullBuffer, ScalarBuffer};
23use arrow_schema::DataType;
24
25use crate::series::ColumnKind;
26
27#[derive(Debug, Clone)]
30pub struct Float64Column {
31 values: ScalarBuffer<f64>,
32 all_finite: bool,
35}
36
37impl Float64Column {
38 pub fn from_values(values: Vec<f64>) -> Self {
40 Self::from_buffer(values.into())
41 }
42
43 pub(crate) fn from_buffer(values: ScalarBuffer<f64>) -> Self {
46 let all_finite = values.iter().all(|v| v.is_finite());
47 Self { values, all_finite }
48 }
49
50 pub(crate) fn buffer(&self) -> &ScalarBuffer<f64> {
52 &self.values
53 }
54
55 pub fn from_options(values: impl IntoIterator<Item = Option<f64>>) -> Self {
57 Self::from_values(values.into_iter().map(|v| v.unwrap_or(f64::NAN)).collect())
58 }
59
60 pub fn len(&self) -> usize {
62 self.values.len()
63 }
64
65 pub fn is_empty(&self) -> bool {
67 self.values.is_empty()
68 }
69
70 pub fn get(&self, i: usize) -> Option<f64> {
72 let v = self.values[i];
73 v.is_finite().then_some(v)
74 }
75
76 pub fn values(&self) -> &[f64] {
78 &self.values
79 }
80
81 pub fn all_finite(&self) -> bool {
83 self.all_finite
84 }
85
86 pub fn missing_count(&self) -> usize {
88 if self.all_finite {
89 0
90 } else {
91 self.values.iter().filter(|v| !v.is_finite()).count()
92 }
93 }
94
95 pub fn iter(&self) -> impl Iterator<Item = Option<f64>> + '_ {
97 self.values.iter().map(|&v| v.is_finite().then_some(v))
98 }
99}
100
101impl PartialEq for Float64Column {
105 fn eq(&self, other: &Self) -> bool {
106 let bits = |c: &Self| c.iter().map(|v| v.map(f64::to_bits)).collect::<Vec<_>>();
107 self.len() == other.len() && bits(self) == bits(other)
108 }
109}
110
111#[derive(Debug, Clone)]
121pub struct StringColumn {
122 repr: StringRepr,
123}
124
125#[derive(Debug, Clone)]
126enum StringRepr {
127 Plain(StringArray),
128 Dictionary(DictionaryArray<Int32Type>),
129}
130
131impl StringColumn {
132 pub fn from_options<S: AsRef<str>>(values: impl IntoIterator<Item = Option<S>>) -> Self {
134 Self::plain(values.into_iter().collect::<StringArray>())
135 }
136
137 pub(crate) fn plain(array: StringArray) -> Self {
139 Self {
140 repr: StringRepr::Plain(array),
141 }
142 }
143
144 pub(crate) fn dictionary(array: DictionaryArray<Int32Type>) -> Option<Self> {
147 let values = array.values();
148 (values.data_type() == &DataType::Utf8 && values.null_count() == 0).then_some(Self {
149 repr: StringRepr::Dictionary(array),
150 })
151 }
152
153 pub fn len(&self) -> usize {
155 match &self.repr {
156 StringRepr::Plain(a) => a.len(),
157 StringRepr::Dictionary(a) => a.len(),
158 }
159 }
160
161 pub fn is_empty(&self) -> bool {
163 self.len() == 0
164 }
165
166 pub fn get(&self, i: usize) -> Option<&str> {
168 assert!(i < self.len(), "row {i} out of range ({} rows)", self.len());
169 match &self.repr {
170 StringRepr::Plain(a) => a.is_valid(i).then(|| a.value(i)),
171 StringRepr::Dictionary(a) => a.is_valid(i).then(|| {
172 let key = a.keys().value(i) as usize;
173 a.values().as_string::<i32>().value(key)
174 }),
175 }
176 }
177
178 pub fn iter(&self) -> impl Iterator<Item = Option<&str>> + '_ {
180 (0..self.len()).map(|i| self.get(i))
181 }
182
183 pub fn missing_count(&self) -> usize {
185 self.arrow().logical_null_count()
186 }
187
188 pub fn is_dictionary(&self) -> bool {
190 matches!(self.repr, StringRepr::Dictionary(_))
191 }
192
193 pub(crate) fn arrow(&self) -> ArrayRef {
195 match &self.repr {
196 StringRepr::Plain(a) => Arc::new(a.clone()),
197 StringRepr::Dictionary(a) => Arc::new(a.clone()),
198 }
199 }
200
201 pub(crate) fn take(&self, rows: &[Option<usize>]) -> Self {
206 self.take_indices(&indices_of(rows))
207 }
208}
209
210impl StringColumn {
211 pub(crate) fn take_indices(&self, indices: &UInt32Array) -> Self {
216 let array: &dyn Array = match &self.repr {
217 StringRepr::Plain(a) => a,
218 StringRepr::Dictionary(a) => a,
219 };
220 let out = arrow_select::take::take(array, indices, None).expect("row numbers are in range");
221 match &self.repr {
222 StringRepr::Plain(_) => Self::plain(out.as_string::<i32>().clone()),
223 StringRepr::Dictionary(_) => Self {
224 repr: StringRepr::Dictionary(out.as_dictionary::<Int32Type>().clone()),
225 },
226 }
227 }
228}
229
230pub(crate) fn indices_of(rows: &[Option<usize>]) -> UInt32Array {
232 rows.iter()
233 .map(|r| r.map(|i| u32::try_from(i).expect("a text column has under 2^32 rows")))
234 .collect()
235}
236
237impl PartialEq for StringColumn {
238 fn eq(&self, other: &Self) -> bool {
239 self.len() == other.len() && self.iter().eq(other.iter())
240 }
241}
242
243#[derive(Debug, Clone, PartialEq)]
246pub struct BooleanColumn {
247 values: BooleanArray,
248}
249
250impl BooleanColumn {
251 pub fn from_options(values: impl IntoIterator<Item = Option<bool>>) -> Self {
253 Self::from_arrow(values.into_iter().collect())
254 }
255
256 pub(crate) fn from_arrow(values: BooleanArray) -> Self {
258 Self { values }
259 }
260
261 pub fn len(&self) -> usize {
263 self.values.len()
264 }
265
266 pub fn is_empty(&self) -> bool {
268 self.values.is_empty()
269 }
270
271 pub fn get(&self, i: usize) -> Option<bool> {
273 assert!(i < self.len(), "row {i} out of range ({} rows)", self.len());
274 self.values.is_valid(i).then(|| self.values.value(i))
275 }
276
277 pub fn iter(&self) -> impl Iterator<Item = Option<bool>> + '_ {
279 self.values.iter()
280 }
281
282 pub fn missing_count(&self) -> usize {
284 self.values.null_count()
285 }
286
287 pub(crate) fn arrow(&self) -> ArrayRef {
289 Arc::new(self.values.clone())
290 }
291
292 pub(crate) fn take(&self, rows: &[Option<usize>]) -> Self {
294 Self::from_options(rows.iter().map(|r| r.and_then(|i| self.get(i))))
295 }
296}
297
298#[derive(Debug, Clone, PartialEq)]
303pub enum Value {
304 Number(f64),
306 Decimal(Decimal),
308 String(String),
310 Boolean(bool),
312 Array(Vec<Scalar>),
314}
315
316impl From<f64> for Value {
317 fn from(v: f64) -> Self {
318 Value::Number(v)
319 }
320}
321
322impl From<&str> for Value {
323 fn from(v: &str) -> Self {
324 Value::String(v.to_string())
325 }
326}
327
328impl From<String> for Value {
329 fn from(v: String) -> Self {
330 Value::String(v)
331 }
332}
333
334impl From<bool> for Value {
335 fn from(v: bool) -> Self {
336 Value::Boolean(v)
337 }
338}
339
340impl From<Vec<Scalar>> for Value {
341 fn from(v: Vec<Scalar>) -> Self {
342 Value::Array(v)
343 }
344}
345
346impl From<Decimal> for Value {
347 fn from(v: Decimal) -> Self {
348 Value::Decimal(v)
349 }
350}
351
352impl Value {
353 pub fn as_f64(&self) -> Option<f64> {
355 match self {
356 Value::Number(v) => Some(*v),
357 _ => None,
358 }
359 }
360
361 pub fn as_decimal(&self) -> Option<Decimal> {
363 match self {
364 Value::Decimal(v) => Some(*v),
365 _ => None,
366 }
367 }
368
369 pub fn as_str(&self) -> Option<&str> {
386 match self {
387 Value::String(v) => Some(v),
388 _ => None,
389 }
390 }
391
392 pub fn as_bool(&self) -> Option<bool> {
394 match self {
395 Value::Boolean(v) => Some(*v),
396 _ => None,
397 }
398 }
399
400 pub fn as_array(&self) -> Option<&[Scalar]> {
402 match self {
403 Value::Array(v) => Some(v),
404 _ => None,
405 }
406 }
407}
408
409#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
414pub struct Decimal {
415 pub value: i128,
417 pub scale: i8,
419}
420
421impl Decimal {
422 pub fn to_f64(self) -> f64 {
425 const EXACT: u128 = 1 << 53;
427 if (-22..=22).contains(&self.scale) && self.value.unsigned_abs() <= EXACT {
428 let power = 10f64.powi(i32::from(self.scale.abs()));
429 return if self.scale >= 0 {
430 self.value as f64 / power
431 } else {
432 self.value as f64 * power
433 };
434 }
435 self.to_string().parse().expect("a decimal parses as f64")
437 }
438}
439
440impl std::fmt::Display for Decimal {
441 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
442 let digits = self.value.unsigned_abs().to_string();
443 let sign = if self.value < 0 { "-" } else { "" };
444 let scale = i32::from(self.scale);
445 if scale <= 0 {
446 let zeros = if self.value == 0 {
447 0
448 } else {
449 (-scale) as usize
450 };
451 return write!(f, "{sign}{digits}{}", "0".repeat(zeros));
452 }
453 let scale = scale as usize;
454 let padded = format!("{digits:0>width$}", width = scale + 1);
455 let (whole, frac) = padded.split_at(padded.len() - scale);
456 write!(f, "{sign}{whole}.{frac}")
457 }
458}
459
460#[derive(Debug, Clone, PartialEq)]
469pub struct DecimalColumn {
470 values: Decimal128Array,
471}
472
473impl DecimalColumn {
474 pub fn from_options(
478 values: impl IntoIterator<Item = Option<i128>>,
479 scale: i8,
480 ) -> crate::error::Result<Self> {
481 let array: Decimal128Array = values.into_iter().collect();
482 Self::with_type(array, 38, scale)
483 }
484
485 pub fn from_values(values: Vec<i128>, scale: i8) -> crate::error::Result<Self> {
487 Self::from_options(values.into_iter().map(Some), scale)
488 }
489
490 pub(crate) fn with_type(
491 array: Decimal128Array,
492 precision: u8,
493 scale: i8,
494 ) -> crate::error::Result<Self> {
495 let invalid = |e: arrow_schema::ArrowError| crate::error::PondError::InvalidData {
496 call: "DecimalColumn",
497 message: e.to_string(),
498 };
499 let array = array
500 .with_precision_and_scale(precision, scale)
501 .map_err(invalid)?;
502 let max = 10u128.pow(u32::from(precision)) - 1;
505 if let Some(i) =
506 (0..array.len()).find(|&i| array.is_valid(i) && array.value(i).unsigned_abs() > max)
507 {
508 return Err(crate::error::PondError::InvalidData {
509 call: "DecimalColumn",
510 message: format!(
511 "row {i}: {} has more than {precision} digits (the precision)",
512 Decimal {
513 value: array.value(i),
514 scale
515 }
516 ),
517 });
518 }
519 Ok(Self { values: array })
520 }
521
522 pub fn len(&self) -> usize {
524 self.values.len()
525 }
526
527 pub fn is_empty(&self) -> bool {
529 self.values.is_empty()
530 }
531
532 pub fn scale(&self) -> i8 {
534 self.values.scale()
535 }
536
537 pub fn precision(&self) -> u8 {
539 self.values.precision()
540 }
541
542 pub fn get(&self, i: usize) -> Option<Decimal> {
544 assert!(i < self.len(), "row {i} out of range ({} rows)", self.len());
545 self.values.is_valid(i).then(|| Decimal {
546 value: self.values.value(i),
547 scale: self.scale(),
548 })
549 }
550
551 pub fn iter(&self) -> impl Iterator<Item = Option<Decimal>> + '_ {
553 (0..self.len()).map(|i| self.get(i))
554 }
555
556 pub(crate) fn raw(&self) -> &[i128] {
559 self.values.values()
560 }
561
562 pub fn missing_count(&self) -> usize {
564 self.values.null_count()
565 }
566
567 pub(crate) fn arrow(&self) -> ArrayRef {
569 Arc::new(self.values.clone())
570 }
571
572 pub(crate) fn nulls(&self) -> Option<&NullBuffer> {
573 self.values.nulls()
574 }
575
576 pub(crate) fn to_f64s(&self) -> Vec<f64> {
578 let scale = self.scale();
579 let divisor = (0..=22)
581 .contains(&scale)
582 .then(|| 10f64.powi(i32::from(scale)));
583 const EXACT: i128 = 1 << 53;
584 let raw = self.raw();
585 (0..self.len())
586 .map(|i| {
587 let value = raw[i];
588 match divisor {
589 _ if self.values.is_null(i) => f64::NAN,
590 Some(d) if (-EXACT..=EXACT).contains(&value) => value as i64 as f64 / d,
591 _ => Decimal { value, scale }.to_f64(),
592 }
593 })
594 .collect()
595 }
596
597 pub(crate) fn rebuild(
599 &self,
600 values: impl IntoIterator<Item = Option<i128>>,
601 precision: u8,
602 ) -> Self {
603 let array: Decimal128Array = values.into_iter().collect();
604 Self {
605 values: array
606 .with_precision_and_scale(precision, self.scale())
607 .expect("this column's own scale"),
608 }
609 }
610
611 pub(crate) fn take(&self, rows: &[Option<usize>]) -> Self {
613 self.rebuild(
614 rows.iter()
615 .map(|r| r.and_then(|i| self.values.is_valid(i).then(|| self.values.value(i)))),
616 self.precision(),
617 )
618 }
619}
620
621#[derive(Debug, Clone, PartialEq)]
629pub enum Scalar {
630 Number(f64),
632 String(String),
634 Boolean(bool),
636}
637
638impl Scalar {
639 pub fn kind(&self) -> ColumnKind {
642 match self {
643 Scalar::Number(_) => ColumnKind::Number,
644 Scalar::String(_) => ColumnKind::String,
645 Scalar::Boolean(_) => ColumnKind::Boolean,
646 }
647 }
648
649 pub fn as_f64(&self) -> Option<f64> {
651 match self {
652 Scalar::Number(v) => Some(*v),
653 _ => None,
654 }
655 }
656
657 pub fn as_str(&self) -> Option<&str> {
659 match self {
660 Scalar::String(v) => Some(v),
661 _ => None,
662 }
663 }
664
665 pub fn as_bool(&self) -> Option<bool> {
667 match self {
668 Scalar::Boolean(v) => Some(*v),
669 _ => None,
670 }
671 }
672
673 pub fn to_value(&self) -> Value {
675 match self {
676 Scalar::Number(v) => Value::Number(*v),
677 Scalar::String(v) => Value::String(v.clone()),
678 Scalar::Boolean(v) => Value::Boolean(*v),
679 }
680 }
681
682 pub(crate) fn same(&self, other: &Scalar) -> bool {
684 match (self, other) {
685 (Scalar::Number(a), Scalar::Number(b)) => a.to_bits() == b.to_bits(),
686 _ => self == other,
687 }
688 }
689}
690
691impl From<f64> for Scalar {
692 fn from(v: f64) -> Self {
693 Scalar::Number(v)
694 }
695}
696
697impl From<i32> for Scalar {
698 fn from(v: i32) -> Self {
699 Scalar::Number(f64::from(v))
700 }
701}
702
703impl From<&str> for Scalar {
704 fn from(v: &str) -> Self {
705 Scalar::String(v.to_string())
706 }
707}
708
709impl From<String> for Scalar {
710 fn from(v: String) -> Self {
711 Scalar::String(v)
712 }
713}
714
715impl From<bool> for Scalar {
716 fn from(v: bool) -> Self {
717 Scalar::Boolean(v)
718 }
719}
720
721pub(crate) type Cell = Arc<[Scalar]>;
724
725#[derive(Debug, Clone)]
753pub struct ArrayColumn {
754 cells: Arc<[Option<Cell>]>,
755 offset: usize,
756 len: usize,
757}
758
759impl ArrayColumn {
760 pub fn from_options(values: impl IntoIterator<Item = Option<Vec<Scalar>>>) -> Self {
762 Self::from_cells(values.into_iter().map(|v| v.map(Cell::from)).collect())
763 }
764
765 pub(crate) fn from_cells(cells: Vec<Option<Cell>>) -> Self {
766 let len = cells.len();
767 Self {
768 cells: cells.into(),
769 offset: 0,
770 len,
771 }
772 }
773
774 pub fn len(&self) -> usize {
776 self.len
777 }
778
779 pub fn is_empty(&self) -> bool {
781 self.len == 0
782 }
783
784 fn cells(&self) -> &[Option<Cell>] {
785 &self.cells[self.offset..self.offset + self.len]
786 }
787
788 pub fn get(&self, i: usize) -> Option<&[Scalar]> {
790 self.cells()[i].as_deref()
791 }
792
793 pub(crate) fn cell(&self, i: usize) -> Option<&Cell> {
795 self.cells()[i].as_ref()
796 }
797
798 pub fn iter(&self) -> impl Iterator<Item = Option<&[Scalar]>> + '_ {
800 self.cells().iter().map(|c| c.as_deref())
801 }
802
803 pub fn missing_count(&self) -> usize {
805 self.cells().iter().filter(|c| c.is_none()).count()
806 }
807
808 pub(crate) fn first_bad_element(&self) -> Option<(usize, usize)> {
811 self.iter().enumerate().find_map(|(r, cell)| {
812 cell?
813 .iter()
814 .position(|e| matches!(e, Scalar::Number(v) if !v.is_finite()))
815 .map(|e| (r, e))
816 })
817 }
818
819 pub(crate) fn take(&self, rows: &[Option<usize>]) -> Self {
822 Self::from_cells(
823 rows.iter()
824 .map(|r| r.and_then(|i| self.cell(i).cloned()))
825 .collect(),
826 )
827 }
828}
829
830impl PartialEq for ArrayColumn {
833 fn eq(&self, other: &Self) -> bool {
834 self.len == other.len
835 && self.iter().zip(other.iter()).all(|(a, b)| match (a, b) {
836 (None, None) => true,
837 (Some(a), Some(b)) => a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.same(y)),
838 _ => false,
839 })
840 }
841}
842
843#[derive(Debug, Clone, PartialEq)]
846pub enum ColumnData {
847 Number(Float64Column),
849 Decimal(DecimalColumn),
851 String(StringColumn),
853 Boolean(BooleanColumn),
855 Array(ArrayColumn),
857}
858
859impl ColumnData {
860 pub fn len(&self) -> usize {
862 match self {
863 ColumnData::Number(c) => c.len(),
864 ColumnData::Decimal(c) => c.len(),
865 ColumnData::String(c) => c.len(),
866 ColumnData::Boolean(c) => c.len(),
867 ColumnData::Array(c) => c.len(),
868 }
869 }
870
871 pub fn is_empty(&self) -> bool {
873 self.len() == 0
874 }
875
876 pub fn missing_count(&self) -> usize {
878 match self {
879 ColumnData::Number(c) => c.missing_count(),
880 ColumnData::Decimal(c) => c.missing_count(),
881 ColumnData::String(c) => c.missing_count(),
882 ColumnData::Boolean(c) => c.missing_count(),
883 ColumnData::Array(c) => c.missing_count(),
884 }
885 }
886
887 pub(crate) fn arrow_nulls(&self) -> Option<NullBuffer> {
891 match self {
892 ColumnData::Number(_) => None,
893 ColumnData::Decimal(c) => c.nulls().cloned(),
894 ColumnData::String(c) => c.arrow().logical_nulls(),
895 ColumnData::Boolean(c) => c.values.nulls().cloned(),
896 ColumnData::Array(c) => (c.missing_count() > 0)
897 .then(|| NullBuffer::from(c.iter().map(|v| v.is_some()).collect::<Vec<bool>>())),
898 }
899 }
900
901 pub(crate) fn slice(&self, range: std::ops::Range<usize>) -> ColumnData {
903 let (offset, len) = (range.start, range.len());
904 match self {
905 ColumnData::Number(c) => ColumnData::Number(Float64Column {
906 values: c.values.slice(offset, len),
907 all_finite: c.all_finite || c.values[range].iter().all(|v| v.is_finite()),
909 }),
910 ColumnData::Decimal(c) => ColumnData::Decimal(DecimalColumn {
911 values: c.values.slice(offset, len),
912 }),
913 ColumnData::String(c) => ColumnData::String(StringColumn {
914 repr: match &c.repr {
915 StringRepr::Plain(a) => StringRepr::Plain(a.slice(offset, len)),
916 StringRepr::Dictionary(a) => StringRepr::Dictionary(a.slice(offset, len)),
917 },
918 }),
919 ColumnData::Boolean(c) => ColumnData::Boolean(BooleanColumn {
920 values: c.values.slice(offset, len),
921 }),
922 ColumnData::Array(c) => ColumnData::Array(ArrayColumn {
923 cells: c.cells.clone(),
924 offset: c.offset + offset,
925 len,
926 }),
927 }
928 }
929
930 pub(crate) fn gather(sources: &[&ColumnData], picks: &[(u32, u32)]) -> ColumnData {
933 macro_rules! each {
934 ($variant:ident) => {
935 sources
936 .iter()
937 .map(|s| match s {
938 ColumnData::$variant(c) => c,
939 _ => unreachable!("sources checked to share a kind"),
940 })
941 .collect::<Vec<_>>()
942 };
943 }
944 match sources[0] {
945 ColumnData::Number(_) => {
946 let cols = each!(Number);
947 let values: Vec<&[f64]> = cols.iter().map(|c| c.values()).collect();
948 ColumnData::from(
949 picks
950 .iter()
951 .map(|&(s, r)| values[s as usize][r as usize])
952 .collect::<Vec<f64>>(),
953 )
954 }
955 ColumnData::Decimal(first) => {
956 let cols = each!(Decimal);
957 let precision = cols.iter().map(|c| c.precision()).max().unwrap_or(38);
958 first
959 .rebuild(
960 picks
961 .iter()
962 .map(|&(s, r)| cols[s as usize].get(r as usize).map(|d| d.value)),
963 precision,
964 )
965 .into()
966 }
967 ColumnData::String(_) => {
968 let cols = each!(String);
969 let picked: Vec<Option<&str>> = picks
970 .iter()
971 .map(|&(s, r)| cols[s as usize].get(r as usize))
972 .collect();
973 let bytes = picked.iter().flatten().map(|s| s.len()).sum();
975 let mut out = StringBuilder::with_capacity(picked.len(), bytes);
976 for v in picked {
977 out.append_option(v);
978 }
979 StringColumn::plain(out.finish()).into()
980 }
981 ColumnData::Boolean(_) => {
982 let cols = each!(Boolean);
983 BooleanColumn::from_options(
984 picks.iter().map(|&(s, r)| cols[s as usize].get(r as usize)),
985 )
986 .into()
987 }
988 ColumnData::Array(_) => {
989 let cols = each!(Array);
990 ArrayColumn::from_cells(
991 picks
992 .iter()
993 .map(|&(s, r)| cols[s as usize].cell(r as usize).cloned())
994 .collect(),
995 )
996 .into()
997 }
998 }
999 }
1000
1001 pub(crate) fn take(&self, rows: &[Option<usize>]) -> ColumnData {
1003 match self {
1004 ColumnData::Number(c) => {
1005 let v = c.values();
1006 ColumnData::from(
1007 rows.iter()
1008 .map(|r| r.map_or(f64::NAN, |i| v[i]))
1009 .collect::<Vec<f64>>(),
1010 )
1011 }
1012 ColumnData::Decimal(c) => c.take(rows).into(),
1013 ColumnData::String(c) => c.take(rows).into(),
1014 ColumnData::Boolean(c) => c.take(rows).into(),
1015 ColumnData::Array(c) => c.take(rows).into(),
1016 }
1017 }
1018
1019 pub(crate) fn take_indices(&self, indices: &UInt32Array) -> ColumnData {
1022 let gather = |a: &dyn Array| {
1023 arrow_select::take::take(a, indices, None).expect("row numbers are in range")
1024 };
1025 match self {
1026 ColumnData::Number(c) => {
1027 let v = c.values();
1028 ColumnData::from(
1029 indices
1030 .iter()
1031 .map(|r| r.map_or(f64::NAN, |i| v[i as usize]))
1032 .collect::<Vec<f64>>(),
1033 )
1034 }
1035 ColumnData::Decimal(c) => DecimalColumn {
1036 values: gather(&c.values).as_primitive::<Decimal128Type>().clone(),
1037 }
1038 .into(),
1039 ColumnData::String(c) => c.take_indices(indices).into(),
1040 ColumnData::Boolean(c) => {
1041 BooleanColumn::from_arrow(gather(&c.values).as_boolean().clone()).into()
1042 }
1043 ColumnData::Array(c) => ArrayColumn::from_cells(
1044 indices
1045 .iter()
1046 .map(|r| r.and_then(|i| c.cell(i as usize).cloned()))
1047 .collect(),
1048 )
1049 .into(),
1050 }
1051 }
1052
1053 pub fn is_present(&self, i: usize) -> bool {
1055 match self {
1056 ColumnData::Number(c) => c.get(i).is_some(),
1057 ColumnData::Decimal(c) => c.get(i).is_some(),
1058 ColumnData::String(c) => c.get(i).is_some(),
1059 ColumnData::Boolean(c) => c.get(i).is_some(),
1060 ColumnData::Array(c) => c.get(i).is_some(),
1061 }
1062 }
1063
1064 pub fn kind(&self) -> ColumnKind {
1067 match self {
1068 ColumnData::Number(_) => ColumnKind::Number,
1069 ColumnData::Decimal(_) => ColumnKind::Decimal,
1070 ColumnData::String(_) => ColumnKind::String,
1071 ColumnData::Boolean(_) => ColumnKind::Boolean,
1072 ColumnData::Array(_) => ColumnKind::Array,
1073 }
1074 }
1075}
1076
1077impl From<Float64Column> for ColumnData {
1078 fn from(c: Float64Column) -> Self {
1079 ColumnData::Number(c)
1080 }
1081}
1082
1083impl From<DecimalColumn> for ColumnData {
1084 fn from(c: DecimalColumn) -> Self {
1085 ColumnData::Decimal(c)
1086 }
1087}
1088
1089impl From<StringColumn> for ColumnData {
1090 fn from(c: StringColumn) -> Self {
1091 ColumnData::String(c)
1092 }
1093}
1094
1095impl From<BooleanColumn> for ColumnData {
1096 fn from(c: BooleanColumn) -> Self {
1097 ColumnData::Boolean(c)
1098 }
1099}
1100
1101impl From<ArrayColumn> for ColumnData {
1102 fn from(c: ArrayColumn) -> Self {
1103 ColumnData::Array(c)
1104 }
1105}
1106
1107macro_rules! column_data_from {
1111 ($($item:ty => |$v:ident| $build:expr;)*) => {$(
1112 impl From<Vec<$item>> for ColumnData {
1113 fn from($v: Vec<$item>) -> Self {
1114 $build
1115 }
1116 }
1117 impl<const N: usize> From<[$item; N]> for ColumnData {
1118 fn from(values: [$item; N]) -> Self {
1119 Vec::from(values).into()
1120 }
1121 }
1122 )*};
1123}
1124
1125column_data_from! {
1126 f64 => |v| ColumnData::Number(Float64Column::from_values(v));
1127 Option<f64> => |v| ColumnData::Number(Float64Column::from_options(v));
1128 &str => |v| ColumnData::String(StringColumn::from_options(v.into_iter().map(Some)));
1129 Option<&str> => |v| ColumnData::String(StringColumn::from_options(v));
1130 String => |v| ColumnData::String(StringColumn::from_options(v.into_iter().map(Some)));
1131 Option<String> => |v| ColumnData::String(StringColumn::from_options(v));
1132 bool => |v| ColumnData::Boolean(BooleanColumn::from_options(v.into_iter().map(Some)));
1133 Option<bool> => |v| ColumnData::Boolean(BooleanColumn::from_options(v));
1134 Vec<Scalar> => |v| ColumnData::Array(ArrayColumn::from_options(v.into_iter().map(Some)));
1135 Option<Vec<Scalar>> => |v| ColumnData::Array(ArrayColumn::from_options(v));
1136}