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

polars_core/frame/row/
av_buffer.rs
