polars_core/frame/row/
av_buffer.rs

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

polars_core/frame/row/
av_buffer.rs
