polars_core/frame/row/
av_buffer.rs

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

polars_core/frame/row/
av_buffer.rs
