polars_core/frame/row/
av_buffer.rs

1#![deny(unsafe_op_in_unsafe_fn)]
2
3use std::hint::unreachable_unchecked;
4
5use polars_arrow::array::builder::ShareStrategy;
6use polars_arrow::bitmap::BitmapBuilder;
7#[cfg(feature = "dtype-decimal")]
8use polars_compute::decimal::DecimalFmtBuffer;
9#[cfg(feature = "dtype-struct")]
10use polars_utils::pl_str::PlSmallStr;
11
12use super::*;
13use crate::chunked_array::builder::NullChunkedBuilder;
14use crate::chunked_array::flags::StatisticsFlags;
15#[cfg(feature = "dtype-struct")]
16use crate::prelude::any_value::arr_to_any_value;
17use crate::series::builder::SeriesBuilder;
18
19#[derive(Clone)]
20pub enum AnyValueBuffer<'a> {
21    Boolean(BooleanChunkedBuilder),
22    #[cfg(feature = "dtype-i8")]
23    Int8(PrimitiveChunkedBuilder<Int8Type>),
24    #[cfg(feature = "dtype-i16")]
25    Int16(PrimitiveChunkedBuilder<Int16Type>),
26    Int32(PrimitiveChunkedBuilder<Int32Type>),
27    Int64(PrimitiveChunkedBuilder<Int64Type>),
28    #[cfg(feature = "dtype-u8")]
29    UInt8(PrimitiveChunkedBuilder<UInt8Type>),
30    #[cfg(feature = "dtype-u16")]
31    UInt16(PrimitiveChunkedBuilder<UInt16Type>),
32    UInt32(PrimitiveChunkedBuilder<UInt32Type>),
33    UInt64(PrimitiveChunkedBuilder<UInt64Type>),
34    #[cfg(feature = "dtype-date")]
35    Date(PrimitiveChunkedBuilder<Int32Type>),
36    #[cfg(feature = "dtype-datetime")]
37    Datetime(
38        PrimitiveChunkedBuilder<Int64Type>,
39        TimeUnit,
40        Option<TimeZone>,
41    ),
42    #[cfg(feature = "dtype-duration")]
43    Duration(PrimitiveChunkedBuilder<Int64Type>, TimeUnit),
44    #[cfg(feature = "dtype-time")]
45    Time(PrimitiveChunkedBuilder<Int64Type>),
46    Float32(PrimitiveChunkedBuilder<Float32Type>),
47    Float64(PrimitiveChunkedBuilder<Float64Type>),
48    String(StringChunkedBuilder),
49    Null(NullChunkedBuilder),
50    All(DataType, Vec<AnyValue<'a>>),
51}
52
53impl<'a> AnyValueBuffer<'a> {
54    #[inline]
55    pub fn add(&mut self, val: AnyValue<'_>) -> Option<()> {
56        use AnyValueBuffer::*;
57        match (self, val) {
58            (Boolean(builder), AnyValue::Null) => builder.append_null(),
59            (Boolean(builder), AnyValue::Boolean(v)) => builder.append_value(v),
60            (Boolean(builder), val) => {
61                let v = val.extract::<u8>()?;
62                builder.append_value(v == 1)
63            },
64            (Int32(builder), AnyValue::Null) => builder.append_null(),
65            (Int32(builder), val) => builder.append_value(val.extract()?),
66            (Int64(builder), AnyValue::Null) => builder.append_null(),
67            (Int64(builder), val) => builder.append_value(val.extract()?),
68            (UInt32(builder), AnyValue::Null) => builder.append_null(),
69            (UInt32(builder), val) => builder.append_value(val.extract()?),
70            (UInt64(builder), AnyValue::Null) => builder.append_null(),
71            (UInt64(builder), val) => builder.append_value(val.extract()?),
72            (Float32(builder), AnyValue::Null) => builder.append_null(),
73            (Float64(builder), AnyValue::Null) => builder.append_null(),
74            (Float32(builder), val) => builder.append_value(val.extract()?),
75            (Float64(builder), val) => builder.append_value(val.extract()?),
76            (String(builder), AnyValue::String(v)) => builder.append_value(v),
77            (String(builder), AnyValue::StringOwned(v)) => builder.append_value(v.as_str()),
78            (String(builder), AnyValue::Null) => builder.append_null(),
79            #[cfg(feature = "dtype-i8")]
80            (Int8(builder), AnyValue::Null) => builder.append_null(),
81            #[cfg(feature = "dtype-i8")]
82            (Int8(builder), val) => builder.append_value(val.extract()?),
83            #[cfg(feature = "dtype-i16")]
84            (Int16(builder), AnyValue::Null) => builder.append_null(),
85            #[cfg(feature = "dtype-i16")]
86            (Int16(builder), val) => builder.append_value(val.extract()?),
87            #[cfg(feature = "dtype-u8")]
88            (UInt8(builder), AnyValue::Null) => builder.append_null(),
89            #[cfg(feature = "dtype-u8")]
90            (UInt8(builder), val) => builder.append_value(val.extract()?),
91            #[cfg(feature = "dtype-u16")]
92            (UInt16(builder), AnyValue::Null) => builder.append_null(),
93            #[cfg(feature = "dtype-u16")]
94            (UInt16(builder), val) => builder.append_value(val.extract()?),
95            #[cfg(feature = "dtype-date")]
96            (Date(builder), AnyValue::Null) => builder.append_null(),
97            #[cfg(feature = "dtype-date")]
98            (Date(builder), AnyValue::Date(v)) => builder.append_value(v),
99            #[cfg(feature = "dtype-date")]
100            (Date(builder), val) if val.is_primitive_numeric() => {
101                builder.append_value(val.extract()?)
102            },
103            #[cfg(feature = "dtype-datetime")]
104            (Datetime(builder, _, _), AnyValue::Null) => builder.append_null(),
105            #[cfg(feature = "dtype-datetime")]
106            (
107                Datetime(builder, tu_l, _),
108                AnyValue::Datetime(v, tu_r, _) | AnyValue::DatetimeOwned(v, tu_r, _),
109            ) => {
110                // we convert right tu to left tu
111                // so we swap.
112                let v = crate::datatypes::time_unit::convert_time_units(v, tu_r, *tu_l);
113                builder.append_value(v)
114            },
115            #[cfg(feature = "dtype-datetime")]
116            (Datetime(builder, _, _), val) if val.is_primitive_numeric() => {
117                builder.append_value(val.extract()?)
118            },
119            #[cfg(feature = "dtype-duration")]
120            (Duration(builder, _), AnyValue::Null) => builder.append_null(),
121            #[cfg(feature = "dtype-duration")]
122            (Duration(builder, tu_l), AnyValue::Duration(v, tu_r)) => {
123                let v = crate::datatypes::time_unit::convert_time_units(v, tu_r, *tu_l);
124                builder.append_value(v)
125            },
126            #[cfg(feature = "dtype-duration")]
127            (Duration(builder, _), val) if val.is_primitive_numeric() => {
128                builder.append_value(val.extract()?)
129            },
130            #[cfg(feature = "dtype-time")]
131            (Time(builder), AnyValue::Time(v)) => builder.append_value(v),
132            #[cfg(feature = "dtype-time")]
133            (Time(builder), AnyValue::Null) => builder.append_null(),
134            #[cfg(feature = "dtype-time")]
135            (Time(builder), val) if val.is_primitive_numeric() => {
136                builder.append_value(val.extract()?)
137            },
138            (Null(builder), AnyValue::Null) => builder.append_null(),
139            // Struct and List can be recursive so use AnyValues for that
140            (All(_, vals), v) => vals.push(v.into_static()),
141
142            // dynamic types
143            (String(builder), av) => match av {
144                AnyValue::Int64(v) => builder.append_value(format!("{v}")),
145                AnyValue::Float64(v) => builder.append_value(format!("{v}")),
146                AnyValue::Boolean(true) => builder.append_value("true"),
147                AnyValue::Boolean(false) => builder.append_value("false"),
148                #[cfg(feature = "dtype-decimal")]
149                AnyValue::Decimal(v, _p, s) => {
150                    let mut fmt = DecimalFmtBuffer::new();
151                    builder.append_value(fmt.format_dec128(v, s, false, false));
152                },
153                _ => return None,
154            },
155            _ => return None,
156        };
157        Some(())
158    }
159
160    pub fn add_fallible(&mut self, val: &AnyValue<'a>) -> PolarsResult<()> {
161        self.add(val.as_borrowed()).ok_or_else(|| {
162            polars_err!(
163                ComputeError: "could not append value: {} of type: {} to the builder; make sure that all rows \
164                have the same schema or consider increasing `infer_schema_length`\n\
165                \n\
166                it might also be that a value overflows the data-type's capacity", val, val.dtype()
167            )
168        })
169    }
170
171    pub fn reset(&mut self, capacity: usize, strict: bool) -> PolarsResult<Series> {
172        use AnyValueBuffer::*;
173        let out = match self {
174            Boolean(b) => {
175                let mut new = BooleanChunkedBuilder::new(b.field.name().clone(), capacity);
176                std::mem::swap(&mut new, b);
177                new.finish().into_series()
178            },
179            Int32(b) => {
180                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
181                std::mem::swap(&mut new, b);
182                new.finish().into_series()
183            },
184            Int64(b) => {
185                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
186                std::mem::swap(&mut new, b);
187                new.finish().into_series()
188            },
189            UInt32(b) => {
190                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
191                std::mem::swap(&mut new, b);
192                new.finish().into_series()
193            },
194            UInt64(b) => {
195                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
196                std::mem::swap(&mut new, b);
197                new.finish().into_series()
198            },
199            #[cfg(feature = "dtype-date")]
200            Date(b) => {
201                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
202                std::mem::swap(&mut new, b);
203                new.finish().into_date().into_series()
204            },
205            #[cfg(feature = "dtype-datetime")]
206            Datetime(b, tu, tz) => {
207                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
208                std::mem::swap(&mut new, b);
209                let tz = if capacity > 0 {
210                    tz.clone()
211                } else {
212                    std::mem::take(tz)
213                };
214                new.finish().into_datetime(*tu, tz).into_series()
215            },
216            #[cfg(feature = "dtype-duration")]
217            Duration(b, tu) => {
218                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
219                std::mem::swap(&mut new, b);
220                new.finish().into_duration(*tu).into_series()
221            },
222            #[cfg(feature = "dtype-time")]
223            Time(b) => {
224                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
225                std::mem::swap(&mut new, b);
226                new.finish().into_time().into_series()
227            },
228            Float32(b) => {
229                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
230                std::mem::swap(&mut new, b);
231                new.finish().into_series()
232            },
233            Float64(b) => {
234                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
235                std::mem::swap(&mut new, b);
236                new.finish().into_series()
237            },
238            String(b) => {
239                let mut new = StringChunkedBuilder::new(b.field.name().clone(), capacity);
240                std::mem::swap(&mut new, b);
241                new.finish().into_series()
242            },
243            #[cfg(feature = "dtype-i8")]
244            Int8(b) => {
245                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
246                std::mem::swap(&mut new, b);
247                new.finish().into_series()
248            },
249            #[cfg(feature = "dtype-i16")]
250            Int16(b) => {
251                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
252                std::mem::swap(&mut new, b);
253                new.finish().into_series()
254            },
255            #[cfg(feature = "dtype-u8")]
256            UInt8(b) => {
257                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
258                std::mem::swap(&mut new, b);
259                new.finish().into_series()
260            },
261            #[cfg(feature = "dtype-u16")]
262            UInt16(b) => {
263                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
264                std::mem::swap(&mut new, b);
265                new.finish().into_series()
266            },
267            Null(b) => {
268                let mut new = NullChunkedBuilder::new(b.field.name().clone(), 0);
269                std::mem::swap(&mut new, b);
270                new.finish().into_series()
271            },
272            All(dtype, vals) => {
273                let out =
274                    Series::from_any_values_and_dtype(PlSmallStr::EMPTY, vals, dtype, strict)?;
275                let mut new = Vec::with_capacity(capacity);
276                std::mem::swap(&mut new, vals);
277                out
278            },
279        };
280        Ok(out)
281    }
282
283    pub fn into_series(mut self) -> PolarsResult<Series> {
284        self.reset(0, false)
285    }
286
287    pub fn new(dtype: &DataType, capacity: usize) -> AnyValueBuffer<'a> {
288        (dtype, capacity).into()
289    }
290}
291
292// datatype and length
293impl From<(&DataType, usize)> for AnyValueBuffer<'_> {
294    fn from(a: (&DataType, usize)) -> Self {
295        let (dt, len) = a;
296        use DataType::*;
297        match dt {
298            Boolean => AnyValueBuffer::Boolean(BooleanChunkedBuilder::new(PlSmallStr::EMPTY, len)),
299            Int32 => AnyValueBuffer::Int32(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len)),
300            Int64 => AnyValueBuffer::Int64(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len)),
301            UInt32 => AnyValueBuffer::UInt32(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len)),
302            UInt64 => AnyValueBuffer::UInt64(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len)),
303            #[cfg(feature = "dtype-i8")]
304            Int8 => AnyValueBuffer::Int8(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len)),
305            #[cfg(feature = "dtype-i16")]
306            Int16 => AnyValueBuffer::Int16(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len)),
307            #[cfg(feature = "dtype-u8")]
308            UInt8 => AnyValueBuffer::UInt8(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len)),
309            #[cfg(feature = "dtype-u16")]
310            UInt16 => AnyValueBuffer::UInt16(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len)),
311            #[cfg(feature = "dtype-date")]
312            Date => AnyValueBuffer::Date(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len)),
313            #[cfg(feature = "dtype-datetime")]
314            Datetime(tu, tz) => AnyValueBuffer::Datetime(
315                PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len),
316                *tu,
317                tz.clone(),
318            ),
319            #[cfg(feature = "dtype-duration")]
320            Duration(tu) => {
321                AnyValueBuffer::Duration(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len), *tu)
322            },
323            #[cfg(feature = "dtype-time")]
324            Time => AnyValueBuffer::Time(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len)),
325            Float32 => {
326                AnyValueBuffer::Float32(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
327            },
328            Float64 => {
329                AnyValueBuffer::Float64(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
330            },
331            String => AnyValueBuffer::String(StringChunkedBuilder::new(PlSmallStr::EMPTY, len)),
332            Null => AnyValueBuffer::Null(NullChunkedBuilder::new(PlSmallStr::EMPTY, 0)),
333            // Struct and List can be recursive so use AnyValues for that
334            dt => AnyValueBuffer::All(dt.clone(), Vec::with_capacity(len)),
335        }
336    }
337}
338
339/// An [`AnyValueBuffer`] that converts nested values in batches (as each value retains
340/// a [`Series`] until finally converted, this can prevent excessive peak memory use).
341pub struct AnyValueBufferBatched<'a> {
342    values: AnyValueBuffer<'a>,
343    converted: Option<SeriesBuilder>,
344    fast_explode: bool,
345}
346
347impl<'a> AnyValueBufferBatched<'a> {
348    const BATCH_SIZE: usize = 1024;
349
350    pub fn new(dtype: &DataType, capacity: usize) -> Self {
351        if dtype.is_nested() && Self::can_batch(dtype) {
352            let mut converted = SeriesBuilder::new(dtype.clone());
353            converted.reserve(capacity);
354            Self {
355                values: AnyValueBuffer::new(dtype, capacity.min(Self::BATCH_SIZE)),
356                converted: Some(converted),
357                fast_explode: dtype.is_list(),
358            }
359        } else {
360            Self {
361                values: AnyValueBuffer::new(dtype, capacity),
362                converted: None,
363                fast_explode: false,
364            }
365        }
366    }
367
368    /// Whether values can convert in batches.
369    fn can_batch(dtype: &DataType) -> bool {
370        match dtype {
371            DataType::List(inner) => Self::can_batch(inner),
372            #[cfg(feature = "dtype-array")]
373            DataType::Array(inner, _) => Self::can_batch(inner),
374            #[cfg(feature = "dtype-struct")]
375            DataType::Struct(fields) => fields.iter().all(|f| Self::can_batch(f.dtype())),
376            _ => {
377                !(dtype.is_nested()
378                    || dtype.is_extension()
379                    || dtype.is_null()
380                    || dtype.is_unknown()
381                    || dtype.is_object())
382            },
383        }
384    }
385
386    #[inline]
387    pub fn add(&mut self, val: AnyValue<'_>) -> Option<()> {
388        self.values.add(val)
389    }
390
391    pub fn add_fallible(&mut self, val: &AnyValue<'a>) -> PolarsResult<()> {
392        self.values.add_fallible(val)
393    }
394
395    /// Whether a full batch of nested values is ready to [`flush`](Self::flush).
396    #[inline]
397    pub fn needs_flush(&self) -> bool {
398        self.converted.is_some()
399            && matches!(&self.values, AnyValueBuffer::All(_, vals) if vals.len() == Self::BATCH_SIZE)
400    }
401
402    /// Converts the buffered nested values, extending the converted [`SeriesBuilder`].
403    /// Frees the memory used by the buffered values.
404    pub fn flush(&mut self) -> PolarsResult<()> {
405        if let Some(converted) = &mut self.converted {
406            let batch = self.values.reset(Self::BATCH_SIZE, false)?;
407            self.fast_explode &= batch.is_empty() || batch.get_flags().can_fast_explode_list();
408            converted.extend(&batch, ShareStrategy::Never);
409        }
410        Ok(())
411    }
412
413    pub fn into_series(mut self) -> PolarsResult<Series> {
414        if self.converted.as_ref().is_none_or(SeriesBuilder::is_empty) {
415            return self.values.into_series();
416        }
417        self.flush()?;
418        let mut series = self.converted.unwrap().freeze(PlSmallStr::EMPTY);
419        if self.fast_explode {
420            // `freeze` drops the flags set by the batches' list builders; put them back
421            series.set_flags(series.get_flags() | StatisticsFlags::CAN_FAST_EXPLODE_LIST);
422        }
423        Ok(series)
424    }
425
426    /// Convert into a column of `height` rows.
427    pub fn into_column(self, name: PlSmallStr, height: usize) -> PolarsResult<Column> {
428        let series = self.into_series()?;
429        Ok(if series.is_empty() {
430            Column::full_null(name, height, series.dtype())
431        } else {
432            series.with_name(name).into_column()
433        })
434    }
435}
436
437/// An [`AnyValueBuffer`] that should be used when we trust the builder.
438#[derive(Clone)]
439pub enum AnyValueBufferTrusted<'a> {
440    Boolean(BooleanChunkedBuilder),
441    #[cfg(feature = "dtype-i8")]
442    Int8(PrimitiveChunkedBuilder<Int8Type>),
443    #[cfg(feature = "dtype-i16")]
444    Int16(PrimitiveChunkedBuilder<Int16Type>),
445    Int32(PrimitiveChunkedBuilder<Int32Type>),
446    Int64(PrimitiveChunkedBuilder<Int64Type>),
447    #[cfg(feature = "dtype-u8")]
448    UInt8(PrimitiveChunkedBuilder<UInt8Type>),
449    #[cfg(feature = "dtype-u16")]
450    UInt16(PrimitiveChunkedBuilder<UInt16Type>),
451    UInt32(PrimitiveChunkedBuilder<UInt32Type>),
452    UInt64(PrimitiveChunkedBuilder<UInt64Type>),
453    Float32(PrimitiveChunkedBuilder<Float32Type>),
454    Float64(PrimitiveChunkedBuilder<Float64Type>),
455    String(StringChunkedBuilder),
456    #[cfg(feature = "dtype-struct")]
457    // not the trusted variant!
458    Struct(BitmapBuilder, Vec<(AnyValueBuffer<'a>, PlSmallStr)>),
459    Null(NullChunkedBuilder),
460    All(DataType, Vec<AnyValue<'a>>),
461}
462
463impl<'a> AnyValueBufferTrusted<'a> {
464    pub fn new(dtype: &DataType, len: usize) -> Self {
465        (dtype, len).into()
466    }
467
468    #[inline]
469    fn add_null(&mut self) {
470        use AnyValueBufferTrusted::*;
471        match self {
472            Boolean(builder) => builder.append_null(),
473            #[cfg(feature = "dtype-i8")]
474            Int8(builder) => builder.append_null(),
475            #[cfg(feature = "dtype-i16")]
476            Int16(builder) => builder.append_null(),
477            Int32(builder) => builder.append_null(),
478            Int64(builder) => builder.append_null(),
479            #[cfg(feature = "dtype-u8")]
480            UInt8(builder) => builder.append_null(),
481            #[cfg(feature = "dtype-u16")]
482            UInt16(builder) => builder.append_null(),
483            UInt32(builder) => builder.append_null(),
484            UInt64(builder) => builder.append_null(),
485            Float32(builder) => builder.append_null(),
486            Float64(builder) => builder.append_null(),
487            String(builder) => builder.append_null(),
488            #[cfg(feature = "dtype-struct")]
489            Struct(outer_validity, builders) => {
490                outer_validity.push(false);
491                for (b, _) in builders.iter_mut() {
492                    b.add(AnyValue::Null);
493                }
494            },
495            Null(builder) => builder.append_null(),
496            All(_, vals) => vals.push(AnyValue::Null),
497        }
498    }
499
500    /// # Safety
501    /// The caller must ensure that the [`AnyValue`] type exactly matches the `Buffer` type.
502    #[inline]
503    unsafe fn add_physical(&mut self, val: &AnyValue<'_>) {
504        // SAFETY: All unsafe blocks rely directly on the function contract.
505
506        use AnyValueBufferTrusted::*;
507        match self {
508            Boolean(builder) => {
509                let AnyValue::Boolean(v) = val else {
510                    unsafe { unreachable_unchecked() }
511                };
512                builder.append_value(*v)
513            },
514            #[cfg(feature = "dtype-i8")]
515            Int8(builder) => {
516                let AnyValue::Int8(v) = val else {
517                    unsafe { unreachable_unchecked() }
518                };
519                builder.append_value(*v)
520            },
521            #[cfg(feature = "dtype-i16")]
522            Int16(builder) => {
523                let AnyValue::Int16(v) = val else {
524                    unsafe { unreachable_unchecked() }
525                };
526                builder.append_value(*v)
527            },
528            Int32(builder) => {
529                let AnyValue::Int32(v) = val else {
530                    unsafe { unreachable_unchecked() }
531                };
532                builder.append_value(*v)
533            },
534            Int64(builder) => {
535                let AnyValue::Int64(v) = val else {
536                    unsafe { unreachable_unchecked() }
537                };
538                builder.append_value(*v)
539            },
540            #[cfg(feature = "dtype-u8")]
541            UInt8(builder) => {
542                let AnyValue::UInt8(v) = val else {
543                    unsafe { unreachable_unchecked() }
544                };
545                builder.append_value(*v)
546            },
547            #[cfg(feature = "dtype-u16")]
548            UInt16(builder) => {
549                let AnyValue::UInt16(v) = val else {
550                    unsafe { unreachable_unchecked() }
551                };
552                builder.append_value(*v)
553            },
554            UInt32(builder) => {
555                let AnyValue::UInt32(v) = val else {
556                    unsafe { unreachable_unchecked() }
557                };
558                builder.append_value(*v)
559            },
560            UInt64(builder) => {
561                let AnyValue::UInt64(v) = val else {
562                    unsafe { unreachable_unchecked() }
563                };
564                builder.append_value(*v)
565            },
566            Float32(builder) => {
567                let AnyValue::Float32(v) = val else {
568                    unsafe { unreachable_unchecked() }
569                };
570                builder.append_value(*v)
571            },
572            Float64(builder) => {
573                let AnyValue::Float64(v) = val else {
574                    unsafe { unreachable_unchecked() }
575                };
576                builder.append_value(*v)
577            },
578            Null(builder) => {
579                let AnyValue::Null = val else {
580                    unsafe { unreachable_unchecked() }
581                };
582                builder.append_null()
583            },
584            _ => unreachable!(),
585        }
586    }
587
588    /// Will add the [`AnyValue`] into [`Self`] and unpack as the physical type belonging to
589    /// [`Self`]. This should only be used with physical buffers
590    ///
591    /// If a type is not primitive or String, the AnyValues will be converted to static
592    ///
593    /// # Safety
594    /// The caller must ensure that the [`AnyValue`] type exactly matches the `Buffer` type and is
595    /// owned.
596    #[inline]
597    pub unsafe fn add_unchecked_owned_physical(&mut self, val: &AnyValue<'a>) {
598        use AnyValueBufferTrusted::*;
599        match val {
600            AnyValue::Null => self.add_null(),
601            _ => {
602                match self {
603                    String(builder) => {
604                        let AnyValue::StringOwned(v) = val else {
605                            // SAFETY: Function contract.
606                            unsafe { unreachable_unchecked() }
607                        };
608                        builder.append_value(v.as_str())
609                    },
610                    #[cfg(feature = "dtype-struct")]
611                    Struct(outer_validity, builders) => {
612                        let AnyValue::StructOwned(payload) = val else {
613                            // SAFETY: Function contract.
614                            unsafe { unreachable_unchecked() }
615                        };
616                        let avs = &*payload.0;
617
618                        debug_assert_eq!(builders.len(), avs.len());
619                        for ((builder, _), av) in builders.iter_mut().zip(avs.iter().cloned()) {
620                            builder.add(av);
621                        }
622                        outer_validity.push(true);
623                    },
624                    All(_, vals) => vals.push(val.clone().into_static()),
625                    // SAFETY: Function contract.
626                    _ => unsafe { self.add_physical(val) },
627                }
628            },
629        }
630    }
631
632    /// # Safety
633    /// The caller must ensure that the [`AnyValue`] type exactly matches the `Buffer` type and is
634    /// borrowed and if `val` is a `AnyValue::Struct` that the values are internally consistent.
635    #[inline]
636    pub unsafe fn add_unchecked_borrowed_physical(&mut self, val: &AnyValue<'a>) {
637        use AnyValueBufferTrusted::*;
638        match val {
639            AnyValue::Null => self.add_null(),
640            _ => {
641                match self {
642                    String(builder) => {
643                        let AnyValue::String(v) = val else {
644                            unsafe { unreachable_unchecked() }
645                        };
646                        builder.append_value(v)
647                    },
648                    #[cfg(feature = "dtype-struct")]
649                    Struct(outer_validity, builders) => {
650                        let AnyValue::Struct(idx, arr, fields) = *val else {
651                            // SAFETY: Function contract.
652                            unsafe { unreachable_unchecked() }
653                        };
654                        let arrays = arr.values();
655                        debug_assert_eq!(builders.len(), arrays.len());
656                        debug_assert_eq!(fields.len(), arrays.len());
657                        for ((field, array), (builder, _)) in
658                            fields.iter().zip(arrays).zip(builders.iter_mut())
659                        {
660                            // SAFETY: The values inside `val` need to be consistent, `idx` MUST be
661                            // in-bounds for all `arr.values()` and `field.dtype` correct.
662                            let av_new = unsafe { arr_to_any_value(&**array, idx, &field.dtype) };
663                            builder.add(av_new);
664                        }
665                        outer_validity.push(true);
666                    },
667                    All(_, vals) => vals.push(val.clone().into_static()),
668                    // SAFETY: Function contract.
669                    _ => unsafe { self.add_physical(val) },
670                }
671            },
672        }
673    }
674
675    /// Clear `self` and give `capacity`, returning the old contents as a [`Series`].
676    pub fn reset(&mut self, capacity: usize, strict: bool) -> PolarsResult<Series> {
677        use AnyValueBufferTrusted::*;
678        let out = match self {
679            Boolean(b) => {
680                let mut new = BooleanChunkedBuilder::new(b.field.name().clone(), capacity);
681                std::mem::swap(&mut new, b);
682                new.finish().into_series()
683            },
684            Int32(b) => {
685                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
686                std::mem::swap(&mut new, b);
687                new.finish().into_series()
688            },
689            Int64(b) => {
690                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
691                std::mem::swap(&mut new, b);
692                new.finish().into_series()
693            },
694            UInt32(b) => {
695                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
696                std::mem::swap(&mut new, b);
697                new.finish().into_series()
698            },
699            UInt64(b) => {
700                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
701                std::mem::swap(&mut new, b);
702                new.finish().into_series()
703            },
704            Float32(b) => {
705                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
706                std::mem::swap(&mut new, b);
707                new.finish().into_series()
708            },
709            Float64(b) => {
710                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
711                std::mem::swap(&mut new, b);
712                new.finish().into_series()
713            },
714            String(b) => {
715                let mut new = StringChunkedBuilder::new(b.field.name().clone(), capacity);
716                std::mem::swap(&mut new, b);
717                new.finish().into_series()
718            },
719            #[cfg(feature = "dtype-i8")]
720            Int8(b) => {
721                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
722                std::mem::swap(&mut new, b);
723                new.finish().into_series()
724            },
725            #[cfg(feature = "dtype-i16")]
726            Int16(b) => {
727                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
728                std::mem::swap(&mut new, b);
729                new.finish().into_series()
730            },
731            #[cfg(feature = "dtype-u8")]
732            UInt8(b) => {
733                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
734                std::mem::swap(&mut new, b);
735                new.finish().into_series()
736            },
737            #[cfg(feature = "dtype-u16")]
738            UInt16(b) => {
739                let mut new = PrimitiveChunkedBuilder::new(b.field.name().clone(), capacity);
740                std::mem::swap(&mut new, b);
741                new.finish().into_series()
742            },
743            #[cfg(feature = "dtype-struct")]
744            Struct(outer_validity, b) => {
745                // @Q? Maybe we need to add a length parameter here for ZFS's. I am not very happy
746                // with just setting the length to zero for that case.
747                if b.is_empty() {
748                    return Ok(
749                        StructChunked::from_series(PlSmallStr::EMPTY, 0, [].iter())?.into_series()
750                    );
751                }
752
753                let mut min_len = usize::MAX;
754                let mut max_len = usize::MIN;
755
756                let v = b
757                    .iter_mut()
758                    .map(|(b, name)| {
759                        let mut s = b.reset(capacity, strict)?;
760
761                        min_len = min_len.min(s.len());
762                        max_len = max_len.max(s.len());
763
764                        s.rename(name.clone());
765                        Ok(s)
766                    })
767                    .collect::<PolarsResult<Vec<_>>>()?;
768
769                let length = if min_len == 0 { 0 } else { max_len };
770
771                let old_outer_validity = core::mem::take(outer_validity);
772                outer_validity.reserve(capacity);
773
774                StructChunked::from_series(PlSmallStr::EMPTY, length, v.iter())?
775                    .with_outer_validity(Some(old_outer_validity.freeze()))
776                    .into_series()
777            },
778            Null(b) => {
779                let mut new = NullChunkedBuilder::new(b.field.name().clone(), 0);
780                std::mem::swap(&mut new, b);
781                new.finish().into_series()
782            },
783            All(dtype, vals) => {
784                let mut swap_vals = Vec::with_capacity(capacity);
785                std::mem::swap(vals, &mut swap_vals);
786                Series::from_any_values_and_dtype(PlSmallStr::EMPTY, &swap_vals, dtype, false)?
787            },
788        };
789
790        Ok(out)
791    }
792
793    pub fn into_series(mut self) -> Series {
794        // unwrap: non-strict does not error.
795        self.reset(0, false).unwrap()
796    }
797}
798
799impl From<(&DataType, usize)> for AnyValueBufferTrusted<'_> {
800    fn from(a: (&DataType, usize)) -> Self {
801        let (dt, len) = a;
802        use DataType::*;
803        match dt {
804            Boolean => {
805                AnyValueBufferTrusted::Boolean(BooleanChunkedBuilder::new(PlSmallStr::EMPTY, len))
806            },
807            Int32 => {
808                AnyValueBufferTrusted::Int32(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
809            },
810            Int64 => {
811                AnyValueBufferTrusted::Int64(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
812            },
813            UInt32 => {
814                AnyValueBufferTrusted::UInt32(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
815            },
816            UInt64 => {
817                AnyValueBufferTrusted::UInt64(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
818            },
819            #[cfg(feature = "dtype-i8")]
820            Int8 => {
821                AnyValueBufferTrusted::Int8(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
822            },
823            #[cfg(feature = "dtype-i16")]
824            Int16 => {
825                AnyValueBufferTrusted::Int16(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
826            },
827            #[cfg(feature = "dtype-u8")]
828            UInt8 => {
829                AnyValueBufferTrusted::UInt8(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
830            },
831            #[cfg(feature = "dtype-u16")]
832            UInt16 => {
833                AnyValueBufferTrusted::UInt16(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
834            },
835            Float32 => {
836                AnyValueBufferTrusted::Float32(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
837            },
838            Float64 => {
839                AnyValueBufferTrusted::Float64(PrimitiveChunkedBuilder::new(PlSmallStr::EMPTY, len))
840            },
841            String => {
842                AnyValueBufferTrusted::String(StringChunkedBuilder::new(PlSmallStr::EMPTY, len))
843            },
844            #[cfg(feature = "dtype-struct")]
845            Struct(fields) => {
846                let outer_validity = BitmapBuilder::with_capacity(len);
847                let buffers = fields
848                    .iter()
849                    .map(|field| {
850                        let dtype = field.dtype().to_physical();
851                        let buffer: AnyValueBuffer = (&dtype, len).into();
852                        (buffer, field.name.clone())
853                    })
854                    .collect::<Vec<_>>();
855                AnyValueBufferTrusted::Struct(outer_validity, buffers)
856            },
857            // List can be recursive so use AnyValues for that
858            dt => AnyValueBufferTrusted::All(dt.clone(), Vec::with_capacity(len)),
859        }
860    }
861}
polars_core/frame/row/av_buffer.rs

polars_core/frame/row/
av_buffer.rs
