polars_core/series/
builder.rs

1use arrow::array::builder::{ArrayBuilder, ShareStrategy, make_builder};
2use polars_utils::IdxSize;
3
4#[cfg(feature = "object")]
5use crate::chunked_array::object::registry::get_object_builder;
6use crate::prelude::*;
7use crate::utils::Container;
8
9/// A type-erased wrapper around ArrayBuilder.
10pub struct SeriesBuilder {
11    dtype: DataType,
12    builder: Box<dyn ArrayBuilder>,
13}
14
15impl SeriesBuilder {
16    pub fn new(dtype: DataType) -> Self {
17        // FIXME: get rid of this hack.
18        #[cfg(feature = "object")]
19        if matches!(dtype, DataType::Object(_)) {
20            let builder = get_object_builder(PlSmallStr::EMPTY, 0).as_array_builder();
21            return Self { dtype, builder };
22        }
23
24        let builder = make_builder(&dtype.to_physical().to_arrow(CompatLevel::newest()));
25        Self { dtype, builder }
26    }
27
28    #[inline(always)]
29    pub fn reserve(&mut self, additional: usize) {
30        self.builder.reserve(additional);
31    }
32
33    pub fn freeze(self, name: PlSmallStr) -> Series {
34        unsafe {
35            Series::from_chunks_and_dtype_unchecked(name, vec![self.builder.freeze()], &self.dtype)
36        }
37    }
38
39    pub fn freeze_reset(&mut self, name: PlSmallStr) -> Series {
40        unsafe {
41            Series::from_chunks_and_dtype_unchecked(
42                name,
43                vec![self.builder.freeze_reset()],
44                &self.dtype,
45            )
46        }
47    }
48
49    pub fn len(&self) -> usize {
50        self.builder.len()
51    }
52
53    pub fn is_empty(&self) -> bool {
54        self.builder.len() == 0
55    }
56
57    /// Extends this builder with the contents of the given series. May panic if
58    /// other does not match the dtype of this builder.
59    #[inline(always)]
60    pub fn extend(&mut self, other: &Series, share: ShareStrategy) {
61        self.subslice_extend(other, 0, other.len(), share);
62    }
63
64    /// Extends this builder with the contents of the given series subslice.
65    /// May panic if other does not match the dtype of this builder.
66    pub fn subslice_extend(
67        &mut self,
68        other: &Series,
69        mut start: usize,
70        mut length: usize,
71        share: ShareStrategy,
72    ) {
73        if length == 0 || other.is_empty() {
74            return;
75        }
76
77        for chunk in other.chunks() {
78            if start < chunk.len() {
79                let length_in_chunk = length.min(chunk.len() - start);
80                self.builder
81                    .subslice_extend(&**chunk, start, length_in_chunk, share);
82
83                start = 0;
84                length -= length_in_chunk;
85                if length == 0 {
86                    break;
87                }
88            } else {
89                start -= chunk.len();
90            }
91        }
92    }
93
94    pub fn subslice_extend_repeated(
95        &mut self,
96        other: &Series,
97        start: usize,
98        length: usize,
99        repeats: usize,
100        share: ShareStrategy,
101    ) {
102        if length == 0 || other.is_empty() {
103            return;
104        }
105
106        let chunks = other.chunks();
107        if chunks.len() == 1 {
108            self.builder
109                .subslice_extend_repeated(&*chunks[0], start, length, repeats, share);
110        } else {
111            for _ in 0..repeats {
112                self.subslice_extend(other, start, length, share);
113            }
114        }
115    }
116
117    /// Extends this builder with the contents of the given series at the given
118    /// indices. That is, `other[idxs[i]]` is appended to this builder in order,
119    /// for each i=0..idxs.len(). May panic if other does not match the dtype
120    /// of this builder, or if the other series is not rechunked.
121    ///
122    /// # Safety
123    /// The indices must be in-bounds.
124    pub unsafe fn gather_extend(&mut self, other: &Series, idxs: &[IdxSize], share: ShareStrategy) {
125        let chunks = other.chunks();
126        assert!(chunks.len() == 1);
127        self.builder.gather_extend(&*chunks[0], idxs, share);
128    }
129
130    pub fn opt_gather_extend(&mut self, other: &Series, idxs: &[IdxSize], share: ShareStrategy) {
131        let chunks = other.chunks();
132        assert!(chunks.len() == 1);
133        self.builder.opt_gather_extend(&*chunks[0], idxs, share);
134    }
135}
polars_core/series/builder.rs

polars_core/series/
builder.rs
