polars_core/frame/row/
av_buffer.rs

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

polars_core/frame/row/
av_buffer.rs
