From eb9f1e81ef1995364a9c9f0ee09d50e45e40399e Mon Sep 17 00:00:00 2001 From: Peter Bower <37089506+pbower@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:41:41 +0100 Subject: [PATCH] Gather throughput enhancements 1. Index gathers match the array variant once for the whole column 2. The gather's null bits skip the per-element bounds checks 3. Index gathers prefetch the read a stride ahead 4. A padded index gather nulls its sentinel positions in one pass --- src/structs/views/array_view.rs | 599 ++++++++++++++++++++------------ src/structs/views/table_view.rs | 357 +------------------ 2 files changed, 386 insertions(+), 570 deletions(-) diff --git a/src/structs/views/array_view.rs b/src/structs/views/array_view.rs index 5315a92..52a3888 100644 --- a/src/structs/views/array_view.rs +++ b/src/structs/views/array_view.rs @@ -53,6 +53,21 @@ use crate::traits::selection::{DataSelector, RowSelection}; use crate::traits::shape::Shape; use crate::{Array, Bitmask, BitmaskV, FieldArray, MaskedArray, TextArray}; +/// Keeps a stride of gathered loads in flight so their memory latency +/// overlaps rather than serialising, since gather indices land anywhere +/// in the window. Both index gathers hint the read this many indices +/// ahead. A no-op off x86-64. +const PREFETCH_AHEAD: usize = 16; + +#[inline(always)] +#[allow(unused_variables)] +fn prefetch_read(ptr: *const T) { + #[cfg(target_arch = "x86_64")] + unsafe { + core::arch::x86_64::_mm_prefetch(ptr as *const i8, core::arch::x86_64::_MM_HINT_T0); + } +} + /// # ArrayView /// /// Logical, windowed view over an `Array`. @@ -393,290 +408,428 @@ impl ArrayV { } /// Gather specific indices from this view into a new materialised Array. - /// Indices are relative to this view's window. + /// Indices are relative to this view's window and must lie within it. + /// + /// The array variant is matched once for the whole gather, so each + /// element costs one indexed copy from the typed buffer rather than a + /// per-element downcast and null test. The output carries a null mask + /// only when the source does, with each gathered position's bit read + /// from the source mask. The value under a null position is + /// unspecified, matching the mask-driven gathers. #[cfg(feature = "select")] pub fn gather_indices(&self, indices: &[usize]) -> Array { use crate::{ BooleanArray, CategoricalArray, FloatArray, IntegerArray, NumericArray, StringArray, - TextArray, + TextArray, Vec64, }; #[cfg(feature = "datetime")] use crate::{DatetimeArray, TemporalArray}; - match &self.array { - Array::Null => Array::Null, - Array::NumericArray(num_arr) => match num_arr { - NumericArray::Int32(_) => { - let mut new_arr = IntegerArray::::with_capacity(indices.len(), true); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } + // Gathers a primitive-typed window into (Vec64, Option) + // with one indexed copy per element, the read a stride ahead + // prefetched, and null bits read from the source mask at each + // gathered position. + macro_rules! gather_idx_prim { + ($self_:expr, $arr:expr, $indices:expr, $T:ty) => {{ + let offset = $self_.offset; + let view_len = $self_.len(); + let data = &$arr.data.as_slice()[offset..offset + view_len]; + let src_mask = $arr.null_mask.as_ref(); + let mut out = Vec64::<$T>::with_capacity($indices.len()); + for (i, &idx) in $indices.iter().enumerate() { + if let Some(&ahead) = $indices.get(i + PREFETCH_AHEAD) { + // wrapping_add keeps a contract-violating index + // defined here, and the hint itself cannot fault. + prefetch_read(data.as_ptr().wrapping_add(ahead)); } - Array::from_int32(new_arr) - } - NumericArray::Int64(_) => { - let mut new_arr = IntegerArray::::with_capacity(indices.len(), true); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); + debug_assert!(idx < data.len(), "gather index outside the window"); + // Safety: the window contract puts every index inside + // `data`, asserted above in debug builds. + out.push(unsafe { *data.get_unchecked(idx) }); + } + let out_mask = src_mask.map(|sm| { + let mut m = Bitmask::new_set_all($indices.len(), true); + for (i, &idx) in $indices.iter().enumerate() { + // Safety: the value loop's contract puts every index + // inside the window, the mask spans the backing + // array, and `i` counts within the output mask's + // own length. + unsafe { + if !sm.get_unchecked(offset + idx) { + m.set_unchecked(i, false); + } } } - Array::from_int64(new_arr) - } - NumericArray::Float32(_) => { - let mut new_arr = FloatArray::::with_capacity(indices.len(), true); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); + m + }); + (out, out_mask) + }}; + } + + // Gathers a string-family window through `get_str`, recording null + // positions to restore after construction. + macro_rules! gather_idx_str { + ($self_:expr, $indices:expr, $ArrTy:ty, $from:path) => {{ + let mut values: Vec<&str> = Vec::with_capacity($indices.len()); + let mut null_at: Vec = Vec::new(); + for &idx in $indices { + match $self_.get_str(idx) { + Some(v) => values.push(v), + None => { + null_at.push(values.len()); + values.push(""); } } - Array::from_float32(new_arr) - } - NumericArray::Float64(_) => { - let mut new_arr = FloatArray::::with_capacity(indices.len(), true); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } + } + let mut new_arr = <$ArrTy>::from_vec(values, None); + for &i in &null_at { + new_arr.set_null(i); + } + $from(new_arr) + }}; + } + + match &self.array { + Array::Null => Array::Null, + Array::NumericArray(num_arr) => match num_arr { + NumericArray::Int32(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, i32); + Array::from_int32(IntegerArray::new(d, m)) + } + NumericArray::Int64(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, i64); + Array::from_int64(IntegerArray::new(d, m)) + } + NumericArray::Float32(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, f32); + Array::from_float32(FloatArray::new(d, m)) + } + NumericArray::Float64(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, f64); + Array::from_float64(FloatArray::new(d, m)) + } + NumericArray::UInt32(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, u32); + Array::from_uint32(IntegerArray::new(d, m)) + } + NumericArray::UInt64(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, u64); + Array::from_uint64(IntegerArray::new(d, m)) + } + #[cfg(feature = "extended_numeric_types")] + NumericArray::Int8(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, i8); + Array::from_int8(IntegerArray::new(d, m)) + } + #[cfg(feature = "extended_numeric_types")] + NumericArray::Int16(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, i16); + Array::from_int16(IntegerArray::new(d, m)) + } + #[cfg(feature = "extended_numeric_types")] + NumericArray::UInt8(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, u8); + Array::from_uint8(IntegerArray::new(d, m)) + } + #[cfg(feature = "extended_numeric_types")] + NumericArray::UInt16(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, u16); + Array::from_uint16(IntegerArray::new(d, m)) + } + NumericArray::Null => Array::Null, + }, + Array::TextArray(text_arr) => match text_arr { + TextArray::String32(_) => { + gather_idx_str!(self, indices, StringArray, Array::from_string32) + } + #[cfg(feature = "large_string")] + TextArray::String64(_) => { + gather_idx_str!(self, indices, StringArray, Array::from_string64) + } + #[cfg(any( + not(feature = "default_categorical_8"), + feature = "extended_categorical" + ))] + TextArray::Categorical32(_) => { + gather_idx_str!( + self, + indices, + CategoricalArray, + Array::from_categorical32 + ) + } + #[cfg(feature = "default_categorical_8")] + TextArray::Categorical8(_) => { + gather_idx_str!( + self, + indices, + CategoricalArray, + Array::from_categorical8 + ) + } + #[cfg(feature = "extended_categorical")] + TextArray::Categorical16(_) => { + gather_idx_str!( + self, + indices, + CategoricalArray, + Array::from_categorical16 + ) + } + #[cfg(feature = "extended_categorical")] + TextArray::Categorical64(_) => { + gather_idx_str!( + self, + indices, + CategoricalArray, + Array::from_categorical64 + ) + } + TextArray::Null => Array::Null, + }, + Array::BooleanArray(arr) => { + let offset = self.offset; + let src_mask = arr.null_mask.as_ref(); + // Both start at length zero - appends set the final length. + let mut bits = Bitmask::new_set_all(0, false); + let mut out_mask = src_mask.map(|_| Bitmask::new_set_all(0, false)); + for &idx in indices { + bits.extend_from_bitmask_range(&arr.data, offset + idx, 1); + if let (Some(om), Some(sm)) = (out_mask.as_mut(), src_mask) { + om.extend_from_bitmask_range(sm, offset + idx, 1); } - Array::from_float64(new_arr) - } - NumericArray::UInt32(_) => { - let mut new_arr = IntegerArray::::with_capacity(indices.len(), true); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); + } + Array::from_bool(BooleanArray::new(bits, out_mask)) + } + #[cfg(feature = "datetime")] + Array::TemporalArray(temp_arr) => match temp_arr { + TemporalArray::Datetime32(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, i32); + Array::from_datetime_i32(DatetimeArray::new(d, m, Some(arr.time_unit))) + } + TemporalArray::Datetime64(arr) => { + let (d, m) = gather_idx_prim!(self, arr, indices, i64); + Array::from_datetime_i64(DatetimeArray::new(d, m, Some(arr.time_unit))) + } + TemporalArray::Null => Array::Null, + }, + } + } + + /// Gather specific indices from this view with a pad sentinel: + /// every index equal to `pad` produces a null at its output + /// position, and every other index copies as + /// [`gather_indices`](Self::gather_indices) does. + /// + /// This is the one-pass form of a padded gather. Callers such as the + /// preserved-side joins otherwise gather through a placeholder index + /// and clear the padded bits afterwards, which costs a second index + /// buffer and a mask combine that this function does not. + #[cfg(feature = "select")] + pub fn gather_indices_padded(&self, indices: &[usize], pad: usize) -> Array { + use crate::{ + BooleanArray, CategoricalArray, FloatArray, IntegerArray, NumericArray, StringArray, + TextArray, Vec64, + }; + #[cfg(feature = "datetime")] + use crate::{DatetimeArray, TemporalArray}; + + // Gathers a primitive-typed window into (Vec64, Bitmask), with + // the pad sentinel clearing its output bit and contributing an + // unspecified value, and the source mask's bit carrying through + // at every real index. + macro_rules! gather_pad_prim { + ($self_:expr, $arr:expr, $indices:expr, $pad:expr, $T:ty) => {{ + let offset = $self_.offset; + let view_len = $self_.len(); + let data = &$arr.data.as_slice()[offset..offset + view_len]; + let src_mask = $arr.null_mask.as_ref(); + let mut out = Vec64::<$T>::with_capacity($indices.len()); + let mut mask = Bitmask::new_set_all($indices.len(), true); + for (i, &idx) in $indices.iter().enumerate() { + if let Some(&ahead) = $indices.get(i + PREFETCH_AHEAD) { + if ahead != $pad { + prefetch_read(data.as_ptr().wrapping_add(ahead)); } } - Array::from_uint32(new_arr) - } - NumericArray::UInt64(_) => { - let mut new_arr = IntegerArray::::with_capacity(indices.len(), true); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); + if idx == $pad { + out.push(<$T>::default()); + unsafe { mask.set_unchecked(i, false) }; + } else { + debug_assert!(idx < data.len(), "gather index outside the window"); + // Safety: the window contract puts every real index + // inside `data`, asserted above in debug builds, the + // mask spans the backing array, and `i` counts + // within the output mask's own length. + unsafe { + out.push(*data.get_unchecked(idx)); + if let Some(sm) = src_mask { + if !sm.get_unchecked(offset + idx) { + mask.set_unchecked(i, false); + } + } } } - Array::from_uint64(new_arr) } - #[cfg(feature = "extended_numeric_types")] - NumericArray::Int8(_) => { - let mut new_arr = IntegerArray::::with_capacity(indices.len(), true); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); + (out, Some(mask)) + }}; + } + + // Gathers a string-family window through `get_str`, with the pad + // sentinel recorded as a null position. + macro_rules! gather_pad_str { + ($self_:expr, $indices:expr, $pad:expr, $ArrTy:ty, $from:path) => {{ + let mut values: Vec<&str> = Vec::with_capacity($indices.len()); + let mut null_at: Vec = Vec::new(); + for &idx in $indices { + let val = if idx == $pad { None } else { $self_.get_str(idx) }; + match val { + Some(v) => values.push(v), + None => { + null_at.push(values.len()); + values.push(""); } } - Array::from_int8(new_arr) + } + let mut new_arr = <$ArrTy>::from_vec(values, None); + for &i in &null_at { + new_arr.set_null(i); + } + $from(new_arr) + }}; + } + + match &self.array { + Array::Null => Array::Null, + Array::NumericArray(num_arr) => match num_arr { + NumericArray::Int32(arr) => { + let (d, m) = gather_pad_prim!(self, arr, indices, pad, i32); + Array::from_int32(IntegerArray::new(d, m)) + } + NumericArray::Int64(arr) => { + let (d, m) = gather_pad_prim!(self, arr, indices, pad, i64); + Array::from_int64(IntegerArray::new(d, m)) + } + NumericArray::Float32(arr) => { + let (d, m) = gather_pad_prim!(self, arr, indices, pad, f32); + Array::from_float32(FloatArray::new(d, m)) + } + NumericArray::Float64(arr) => { + let (d, m) = gather_pad_prim!(self, arr, indices, pad, f64); + Array::from_float64(FloatArray::new(d, m)) + } + NumericArray::UInt32(arr) => { + let (d, m) = gather_pad_prim!(self, arr, indices, pad, u32); + Array::from_uint32(IntegerArray::new(d, m)) + } + NumericArray::UInt64(arr) => { + let (d, m) = gather_pad_prim!(self, arr, indices, pad, u64); + Array::from_uint64(IntegerArray::new(d, m)) } #[cfg(feature = "extended_numeric_types")] - NumericArray::Int16(_) => { - let mut new_arr = IntegerArray::::with_capacity(indices.len(), true); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_int16(new_arr) + NumericArray::Int8(arr) => { + let (d, m) = gather_pad_prim!(self, arr, indices, pad, i8); + Array::from_int8(IntegerArray::new(d, m)) } #[cfg(feature = "extended_numeric_types")] - NumericArray::UInt8(_) => { - let mut new_arr = IntegerArray::::with_capacity(indices.len(), true); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_uint8(new_arr) + NumericArray::Int16(arr) => { + let (d, m) = gather_pad_prim!(self, arr, indices, pad, i16); + Array::from_int16(IntegerArray::new(d, m)) } #[cfg(feature = "extended_numeric_types")] - NumericArray::UInt16(_) => { - let mut new_arr = IntegerArray::::with_capacity(indices.len(), true); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_uint16(new_arr) + NumericArray::UInt8(arr) => { + let (d, m) = gather_pad_prim!(self, arr, indices, pad, u8); + Array::from_uint8(IntegerArray::new(d, m)) + } + #[cfg(feature = "extended_numeric_types")] + NumericArray::UInt16(arr) => { + let (d, m) = gather_pad_prim!(self, arr, indices, pad, u16); + Array::from_uint16(IntegerArray::new(d, m)) } NumericArray::Null => Array::Null, }, Array::TextArray(text_arr) => match text_arr { TextArray::String32(_) => { - let mut values: Vec<&str> = Vec::with_capacity(indices.len()); - for &idx in indices { - if let Some(val) = self.get_str(idx) { - values.push(val); - } else { - values.push(""); - } - } - let mut new_arr = StringArray::::from_vec(values, None); - for (i, &idx) in indices.iter().enumerate() { - if self.get_str(idx).is_none() { - new_arr.set_null(i); - } - } - Array::from_string32(new_arr) + gather_pad_str!(self, indices, pad, StringArray, Array::from_string32) } #[cfg(feature = "large_string")] TextArray::String64(_) => { - let mut values: Vec<&str> = Vec::with_capacity(indices.len()); - for &idx in indices { - if let Some(val) = self.get_str(idx) { - values.push(val); - } else { - values.push(""); - } - } - let mut new_arr = StringArray::::from_vec(values, None); - for (i, &idx) in indices.iter().enumerate() { - if self.get_str(idx).is_none() { - new_arr.set_null(i); - } - } - Array::from_string64(new_arr) + gather_pad_str!(self, indices, pad, StringArray, Array::from_string64) } #[cfg(any( not(feature = "default_categorical_8"), feature = "extended_categorical" ))] TextArray::Categorical32(_) => { - let mut values: Vec<&str> = Vec::with_capacity(indices.len()); - for &idx in indices { - if let Some(val) = self.get_str(idx) { - values.push(val); - } else { - values.push(""); - } - } - let mut new_arr = CategoricalArray::::from_vec(values, None); - for (i, &idx) in indices.iter().enumerate() { - if self.get_str(idx).is_none() { - new_arr.set_null(i); - } - } - Array::from_categorical32(new_arr) + gather_pad_str!( + self, + indices, + pad, + CategoricalArray, + Array::from_categorical32 + ) } #[cfg(feature = "default_categorical_8")] TextArray::Categorical8(_) => { - let mut values: Vec<&str> = Vec::with_capacity(indices.len()); - for &idx in indices { - if let Some(val) = self.get_str(idx) { - values.push(val); - } else { - values.push(""); - } - } - let mut new_arr = CategoricalArray::::from_vec(values, None); - for (i, &idx) in indices.iter().enumerate() { - if self.get_str(idx).is_none() { - new_arr.set_null(i); - } - } - Array::from_categorical8(new_arr) + gather_pad_str!( + self, + indices, + pad, + CategoricalArray, + Array::from_categorical8 + ) } #[cfg(feature = "extended_categorical")] TextArray::Categorical16(_) => { - let mut values: Vec<&str> = Vec::with_capacity(indices.len()); - for &idx in indices { - if let Some(val) = self.get_str(idx) { - values.push(val); - } else { - values.push(""); - } - } - let mut new_arr = CategoricalArray::::from_vec(values, None); - for (i, &idx) in indices.iter().enumerate() { - if self.get_str(idx).is_none() { - new_arr.set_null(i); - } - } - Array::from_categorical16(new_arr) + gather_pad_str!( + self, + indices, + pad, + CategoricalArray, + Array::from_categorical16 + ) } #[cfg(feature = "extended_categorical")] TextArray::Categorical64(_) => { - let mut values: Vec<&str> = Vec::with_capacity(indices.len()); - for &idx in indices { - if let Some(val) = self.get_str(idx) { - values.push(val); - } else { - values.push(""); - } - } - let mut new_arr = CategoricalArray::::from_vec(values, None); - for (i, &idx) in indices.iter().enumerate() { - if self.get_str(idx).is_none() { - new_arr.set_null(i); - } - } - Array::from_categorical64(new_arr) + gather_pad_str!( + self, + indices, + pad, + CategoricalArray, + Array::from_categorical64 + ) } TextArray::Null => Array::Null, }, - Array::BooleanArray(_) => { - let mut new_arr = BooleanArray::with_capacity(indices.len(), true); + Array::BooleanArray(arr) => { + let offset = self.offset; + let src_mask = arr.null_mask.as_ref(); + // Both start at length zero - appends set the final length. + let mut bits = Bitmask::new_set_all(0, false); + let mut mask = Bitmask::new_set_all(0, false); for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); + if idx == pad { + bits.push_bits(false, 1); + mask.push_bits(false, 1); } else { - new_arr.push_null(); + bits.extend_from_bitmask_range(&arr.data, offset + idx, 1); + match src_mask { + Some(sm) => mask.extend_from_bitmask_range(sm, offset + idx, 1), + None => mask.push_bits(true, 1), + } } } - Array::from_bool(new_arr) + Array::from_bool(BooleanArray::new(bits, Some(mask))) } #[cfg(feature = "datetime")] Array::TemporalArray(temp_arr) => match temp_arr { TemporalArray::Datetime32(arr) => { - let mut new_arr = DatetimeArray::::with_capacity( - indices.len(), - true, - Some(arr.time_unit), - ); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_datetime_i32(new_arr) + let (d, m) = gather_pad_prim!(self, arr, indices, pad, i32); + Array::from_datetime_i32(DatetimeArray::new(d, m, Some(arr.time_unit))) } TemporalArray::Datetime64(arr) => { - let mut new_arr = DatetimeArray::::with_capacity( - indices.len(), - true, - Some(arr.time_unit), - ); - for &idx in indices { - if let Some(val) = self.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_datetime_i64(new_arr) + let (d, m) = gather_pad_prim!(self, arr, indices, pad, i64); + Array::from_datetime_i64(DatetimeArray::new(d, m, Some(arr.time_unit))) } TemporalArray::Null => Array::Null, }, @@ -766,7 +919,9 @@ impl ArrayV { } }, |idx| { - out.push(data[idx]); + debug_assert!(idx < data.len()); + // Safety: the walk bounds `idx` by the window length. + out.push(unsafe { *data.get_unchecked(idx) }); if let (Some(om), Some(sm)) = (out_mask.as_mut(), src_mask) { om.extend_from_bitmask_range(sm, offset + idx, 1); } diff --git a/src/structs/views/table_view.rs b/src/structs/views/table_view.rs index 3d65881..a0d3fb7 100644 --- a/src/structs/views/table_view.rs +++ b/src/structs/views/table_view.rs @@ -448,356 +448,17 @@ impl TableV { Ok(Table::new(self.name.clone(), Some(cols))) } - /// Gather specific rows from an ArrayV window + /// Gather specific rows from an ArrayV window. + /// + /// Delegates to [`ArrayV::gather_indices`], whose typed per-column + /// loops carry the cost, and reads its `Array::Null` result as a + /// column with nothing to gather. #[cfg(feature = "select")] fn gather_rows_from_window(&self, window: &ArrayV, row_indices: &[usize]) -> Option { - use crate::{ - Array, BooleanArray, CategoricalArray, FloatArray, IntegerArray, MaskedArray, - NumericArray, StringArray, TextArray, - }; - #[cfg(feature = "datetime")] - use crate::{DatetimeArray, TemporalArray}; - - let result = match &window.array { - Array::Null => return None, - Array::NumericArray(num_arr) => match num_arr { - NumericArray::Int32(_) => { - let mut new_arr = IntegerArray::::with_capacity(row_indices.len(), true); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_int32(new_arr) - } - NumericArray::Int64(_) => { - let mut new_arr = IntegerArray::::with_capacity(row_indices.len(), true); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_int64(new_arr) - } - NumericArray::UInt32(_) => { - let mut new_arr = IntegerArray::::with_capacity(row_indices.len(), true); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_uint32(new_arr) - } - NumericArray::UInt64(_) => { - let mut new_arr = IntegerArray::::with_capacity(row_indices.len(), true); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_uint64(new_arr) - } - NumericArray::Float32(_) => { - let mut new_arr = FloatArray::::with_capacity(row_indices.len(), true); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_float32(new_arr) - } - NumericArray::Float64(_) => { - let mut new_arr = FloatArray::::with_capacity(row_indices.len(), true); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_float64(new_arr) - } - #[cfg(feature = "extended_numeric_types")] - NumericArray::Int8(_) => { - let mut new_arr = IntegerArray::::with_capacity(row_indices.len(), true); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_int8(new_arr) - } - #[cfg(feature = "extended_numeric_types")] - NumericArray::Int16(_) => { - let mut new_arr = IntegerArray::::with_capacity(row_indices.len(), true); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_int16(new_arr) - } - #[cfg(feature = "extended_numeric_types")] - NumericArray::UInt8(_) => { - let mut new_arr = IntegerArray::::with_capacity(row_indices.len(), true); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_uint8(new_arr) - } - #[cfg(feature = "extended_numeric_types")] - NumericArray::UInt16(_) => { - let mut new_arr = IntegerArray::::with_capacity(row_indices.len(), true); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_uint16(new_arr) - } - NumericArray::Null => return None, - }, - Array::TextArray(text_arr) => match text_arr { - TextArray::String32(_) => { - let mut new_arr = StringArray::::default(); - for &idx in row_indices { - if let Some(val) = window.get_str(idx) { - new_arr.push(val.to_string()); - } else { - new_arr.push_null(); - } - } - Array::from_string32(new_arr) - } - #[cfg(feature = "large_string")] - TextArray::String64(_) => { - let mut new_arr = StringArray::::default(); - for &idx in row_indices { - if let Some(val) = window.get_str(idx) { - new_arr.push(val.to_string()); - } else { - new_arr.push_null(); - } - } - Array::from_string64(new_arr) - } - #[cfg(any( - not(feature = "default_categorical_8"), - feature = "extended_categorical" - ))] - TextArray::Categorical32(_) => { - use crate::{Bitmask, Vec64}; - use std::collections::HashMap; - - let mut codes = Vec64::::with_capacity(row_indices.len()); - let mut value_map = HashMap::::new(); - let mut mask = Bitmask::new_set_all(row_indices.len(), true); - - for (i, &idx) in row_indices.iter().enumerate() { - if let Some(val) = window.get_str(idx) { - let code = if let Some(&existing_code) = value_map.get(val) { - existing_code - } else { - let new_code = value_map.len() as u32; - value_map.insert(val.to_string(), new_code); - new_code - }; - codes.push(code); - } else { - codes.push(0); - mask.set_false(i); - } - } - - let mut unique_values = Vec64::::with_capacity(value_map.len()); - unique_values.resize(value_map.len(), String::new()); - for (val, code) in value_map { - unique_values[code as usize] = val; - } - - let null_mask = if mask.all_set() { None } else { Some(mask) }; - - let new_arr = CategoricalArray::::new(codes, unique_values, null_mask); - Array::from_categorical32(new_arr) - } - #[cfg(feature = "default_categorical_8")] - TextArray::Categorical8(_) => { - use crate::{Bitmask, Vec64}; - use std::collections::HashMap; - - let mut codes = Vec64::::with_capacity(row_indices.len()); - let mut value_map = HashMap::::new(); - let mut mask = Bitmask::new_set_all(row_indices.len(), true); - - for (i, &idx) in row_indices.iter().enumerate() { - if let Some(val) = window.get_str(idx) { - let code = if let Some(&existing_code) = value_map.get(val) { - existing_code - } else { - let new_code = value_map.len() as u8; - value_map.insert(val.to_string(), new_code); - new_code - }; - codes.push(code); - } else { - codes.push(0); - mask.set_false(i); - } - } - - let mut unique_values = Vec64::::with_capacity(value_map.len()); - unique_values.resize(value_map.len(), String::new()); - for (val, code) in value_map { - unique_values[code as usize] = val; - } - - let null_mask = if mask.all_set() { None } else { Some(mask) }; - - let new_arr = CategoricalArray::::new(codes, unique_values, null_mask); - Array::from_categorical8(new_arr) - } - #[cfg(feature = "extended_categorical")] - TextArray::Categorical16(_) => { - use crate::{Bitmask, Vec64}; - use std::collections::HashMap; - - let mut codes = Vec64::::with_capacity(row_indices.len()); - let mut value_map = HashMap::::new(); - let mut mask = Bitmask::new_set_all(row_indices.len(), true); - - for (i, &idx) in row_indices.iter().enumerate() { - if let Some(val) = window.get_str(idx) { - let code = if let Some(&existing_code) = value_map.get(val) { - existing_code - } else { - let new_code = value_map.len() as u16; - value_map.insert(val.to_string(), new_code); - new_code - }; - codes.push(code); - } else { - codes.push(0); - mask.set_false(i); - } - } - - let mut unique_values = Vec64::::with_capacity(value_map.len()); - unique_values.resize(value_map.len(), String::new()); - for (val, code) in value_map { - unique_values[code as usize] = val; - } - - let null_mask = if mask.all_set() { None } else { Some(mask) }; - - let new_arr = CategoricalArray::::new(codes, unique_values, null_mask); - Array::from_categorical16(new_arr) - } - #[cfg(feature = "extended_categorical")] - TextArray::Categorical64(_) => { - use crate::{Bitmask, Vec64}; - use std::collections::HashMap; - - let mut codes = Vec64::::with_capacity(row_indices.len()); - let mut value_map = HashMap::::new(); - let mut mask = Bitmask::new_set_all(row_indices.len(), true); - - for (i, &idx) in row_indices.iter().enumerate() { - if let Some(val) = window.get_str(idx) { - let code = if let Some(&existing_code) = value_map.get(val) { - existing_code - } else { - let new_code = value_map.len() as u64; - value_map.insert(val.to_string(), new_code); - new_code - }; - codes.push(code); - } else { - codes.push(0); - mask.set_false(i); - } - } - - let mut unique_values = Vec64::::with_capacity(value_map.len()); - unique_values.resize(value_map.len(), String::new()); - for (val, code) in value_map { - unique_values[code as usize] = val; - } - - let null_mask = if mask.all_set() { None } else { Some(mask) }; - - let new_arr = CategoricalArray::::new(codes, unique_values, null_mask); - Array::from_categorical64(new_arr) - } - TextArray::Null => return None, - }, - Array::BooleanArray(_) => { - let mut new_arr = BooleanArray::with_capacity(row_indices.len(), true); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_bool(new_arr) - } - #[cfg(feature = "datetime")] - Array::TemporalArray(temp_arr) => match temp_arr { - TemporalArray::Datetime32(arr) => { - let mut new_arr = DatetimeArray::::with_capacity( - row_indices.len(), - true, - Some(arr.time_unit), - ); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_datetime_i32(new_arr) - } - TemporalArray::Datetime64(arr) => { - let mut new_arr = DatetimeArray::::with_capacity( - row_indices.len(), - true, - Some(arr.time_unit), - ); - for &idx in row_indices { - if let Some(val) = window.get::>(idx) { - new_arr.push(val); - } else { - new_arr.push_null(); - } - } - Array::from_datetime_i64(new_arr) - } - TemporalArray::Null => return None, - }, - }; - - Some(result) + match window.gather_indices(row_indices) { + Array::Null => None, + arr => Some(arr), + } } /// Converts a column window into an owned `FieldArray`, slicing the array and copying data.