polars_core/frame/column/
scalar.rs

1use std::sync::OnceLock;
2
3use polars_error::PolarsResult;
4use polars_utils::broadcast::BroadcastLength;
5use polars_utils::pl_str::PlSmallStr;
6
7use super::{AnyValue, Column, DataType, IntoColumn, Scalar, Series};
8use crate::chunked_array::cast::CastOptions;
9
10/// A [`Column`] that consists of a repeated [`Scalar`]
11///
12/// This is lazily materialized into a [`Series`].
13#[derive(Debug, Clone)]
14pub struct ScalarColumn {
15    name: PlSmallStr,
16    // The value of this scalar may be unspecified when `length == 0`.
17    scalar: Scalar,
18    length: usize,
19
20    // invariants:
21    // materialized.name() == name
22    // materialized.len() == length
23    // materialized.dtype() == value.dtype
24    // materialized[i] == value, for all 0 <= i < length
25    /// A lazily materialized [`Series`] variant of this [`ScalarColumn`]
26    materialized: OnceLock<Series>,
27}
28
29impl ScalarColumn {
30    #[inline]
31    pub fn new(name: PlSmallStr, scalar: Scalar, length: usize) -> Self {
32        Self {
33            name,
34            scalar,
35            length,
36
37            materialized: OnceLock::new(),
38        }
39    }
40
41    #[inline]
42    pub fn new_empty(name: PlSmallStr, dtype: DataType) -> Self {
43        Self {
44            name,
45            scalar: Scalar::new(dtype, AnyValue::Null),
46            length: 0,
47
48            materialized: OnceLock::new(),
49        }
50    }
51
52    pub fn full_null(name: PlSmallStr, length: usize, dtype: DataType) -> Self {
53        Self::new(name, Scalar::null(dtype), length)
54    }
55
56    pub fn name(&self) -> &PlSmallStr {
57        &self.name
58    }
59
60    pub fn scalar(&self) -> &Scalar {
61        &self.scalar
62    }
63
64    pub fn dtype(&self) -> &DataType {
65        self.scalar.dtype()
66    }
67
68    #[inline]
69    pub fn len(&self) -> usize {
70        self.length
71    }
72
73    pub fn is_empty(&self) -> bool {
74        self.length == 0
75    }
76
77    pub fn is_full_null(&self) -> bool {
78        self.scalar.is_null()
79    }
80
81    fn _to_series(name: PlSmallStr, value: Scalar, length: usize) -> Series {
82        let series = if length == 0 {
83            Series::new_empty(name, value.dtype())
84        } else {
85            value.into_series(name).new_from_index(0, length)
86        };
87
88        debug_assert_eq!(series.len(), length);
89
90        series
91    }
92
93    /// Materialize the [`ScalarColumn`] into a [`Series`].
94    pub fn to_series(&self) -> Series {
95        Self::_to_series(self.name.clone(), self.scalar.clone(), self.length)
96    }
97
98    /// Get the [`ScalarColumn`] as [`Series`] if it was already materialized.
99    pub fn lazy_as_materialized_series(&self) -> Option<&Series> {
100        self.materialized.get()
101    }
102
103    /// Get the [`ScalarColumn`] as [`Series`]
104    ///
105    /// This needs to materialize upon the first call. Afterwards, this is cached.
106    pub fn as_materialized_series(&self) -> &Series {
107        self.materialized.get_or_init(|| self.to_series())
108    }
109
110    /// Take the [`ScalarColumn`] and materialize as a [`Series`] if not already done.
111    pub fn take_materialized_series(self) -> Series {
112        self.materialized
113            .into_inner()
114            .unwrap_or_else(|| Self::_to_series(self.name, self.scalar, self.length))
115    }
116
117    /// Estimated size of the column, see [`Series::estimated_size`].
118    ///
119    /// If the column is not materialized and `expanded` is false, only the scalar value is
120    /// counted. Otherwise, this is the size of the materialized [`Series`], computed without
121    /// materializing it.
122    pub fn estimated_size(&self, expanded: bool) -> usize {
123        if let Some(s) = self.materialized.get() {
124            return s.estimated_size();
125        }
126        let length = if expanded { self.length } else { 1 };
127        estimated_repeated_size(self.dtype(), self.scalar.value(), length)
128            .unwrap_or_else(|| self.as_single_value_series().estimated_size() * length)
129    }
130
131    /// Take the [`ScalarColumn`] as a series with a single value.
132    ///
133    /// If the [`ScalarColumn`] has `length=0` the resulting `Series` will also have `length=0`.
134    pub fn as_single_value_series(&self) -> Series {
135        self.as_n_values_series(1)
136    }
137
138    /// Take the [`ScalarColumn`] as a series with a `n` values.
139    ///
140    /// If the [`ScalarColumn`] has `length=0` the resulting `Series` will also have `length=0`.
141    pub fn as_n_values_series(&self, n: usize) -> Series {
142        let length = usize::min(n, self.length);
143
144        match self.materialized.get() {
145            // Don't take a refcount if we only want length-1 (or empty) - the materialized series
146            // could be extremely large.
147            Some(s) if length == self.length || length > 1 => s.head(Some(length)),
148            _ => Self::_to_series(self.name.clone(), self.scalar.clone(), length),
149        }
150    }
151
152    /// Create a new [`ScalarColumn`] from a `length=1` Series and expand it `length`.
153    ///
154    /// This will panic if the value cannot be made static or if the series has length `0`.
155    #[inline]
156    pub fn unit_scalar_from_series(series: Series) -> Self {
157        assert_eq!(series.len(), 1);
158        // SAFETY: We just did the bounds check
159        let value = unsafe { series.get_unchecked(0) };
160        let value = value.into_static();
161        let value = Scalar::new(series.dtype().clone(), value);
162        let mut sc = ScalarColumn::new(series.name().clone(), value, 1);
163        sc.materialized = OnceLock::from(series);
164        sc
165    }
166
167    /// Create a new [`ScalarColumn`] from a `length<=1` Series and expand it `length`.
168    ///
169    /// If `series` is empty and `length` is non-zero, a full-NULL column of `length` will be returned.
170    ///
171    /// This will panic if the value cannot be made static.
172    pub fn from_single_value_series(series: Series, length: usize) -> Self {
173        debug_assert!(series.len() <= 1);
174
175        let value = if series.is_empty() {
176            AnyValue::Null
177        } else {
178            unsafe { series.get_unchecked(0) }.into_static()
179        };
180        let value = Scalar::new(series.dtype().clone(), value);
181        ScalarColumn::new(series.name().clone(), value, length)
182    }
183
184    /// Resize the [`ScalarColumn`] to new `length`.
185    ///
186    /// This reuses the materialized [`Series`], if `length <= self.length`.
187    pub fn resize(&self, length: usize) -> ScalarColumn {
188        if self.length == length {
189            return self.clone();
190        }
191
192        // This is violates an invariant if this triggers, the scalar value is undefined if the
193        // self.length == 0 so therefore we should never resize using that value.
194        debug_assert!(length == 0 || self.length > 0);
195
196        let mut resized = Self {
197            name: self.name.clone(),
198            scalar: self.scalar.clone(),
199            length,
200            materialized: OnceLock::new(),
201        };
202
203        if length == self.length || (length < self.length && length > 1) {
204            if let Some(materialized) = self.materialized.get() {
205                resized.materialized = OnceLock::from(materialized.head(Some(length)));
206                debug_assert_eq!(resized.materialized.get().unwrap().len(), length);
207            }
208        }
209
210        resized
211    }
212
213    /// Append `other` to `self`, keeping both unmaterialized.
214    ///
215    /// Returns whether that was possible. `self` is left untouched when it was not.
216    pub fn try_append(&mut self, other: &Self) -> bool {
217        // Unequal dtypes either need a cast or must raise.
218        if self.dtype() != other.dtype() {
219            return false;
220        }
221
222        // The value of a length-0 column is unspecified, so it takes on the other's.
223        if other.is_empty() {
224            return true;
225        }
226        if self.is_empty() {
227            let name = std::mem::take(&mut self.name);
228            *self = other.clone();
229            self.rename(name);
230            return true;
231        }
232
233        if !is_same_value(self.dtype(), self.scalar.value(), other.scalar.value()) {
234            return false;
235        }
236
237        self.length += other.length;
238        self.materialized.take();
239        true
240    }
241
242    pub fn cast_with_options(&self, dtype: &DataType, options: CastOptions) -> PolarsResult<Self> {
243        // @NOTE: We expect that when casting the materialized series mostly does not need change
244        // the physical array. Therefore, we try to cast the entire materialized array if it is
245        // available.
246
247        match self.materialized.get() {
248            Some(s) => {
249                let materialized = s.cast_with_options(dtype, options)?;
250                assert_eq!(self.length, materialized.len());
251
252                let mut casted = if materialized.is_empty() {
253                    Self::new_empty(materialized.name().clone(), materialized.dtype().clone())
254                } else {
255                    // SAFETY: Just did bounds check
256                    let scalar = unsafe { materialized.get_unchecked(0) }.into_static();
257                    Self::new(
258                        materialized.name().clone(),
259                        Scalar::new(materialized.dtype().clone(), scalar),
260                        self.length,
261                    )
262                };
263                casted.materialized = OnceLock::from(materialized);
264                Ok(casted)
265            },
266            None => {
267                let s = self
268                    .as_single_value_series()
269                    .cast_with_options(dtype, options)?;
270
271                if self.length == 0 {
272                    Ok(Self::new_empty(s.name().clone(), s.dtype().clone()))
273                } else {
274                    assert_eq!(1, s.len());
275                    Ok(Self::from_single_value_series(s, self.length))
276                }
277            },
278        }
279    }
280
281    pub fn strict_cast(&self, dtype: &DataType) -> PolarsResult<Self> {
282        self.cast_with_options(dtype, CastOptions::Strict)
283    }
284    pub fn cast(&self, dtype: &DataType) -> PolarsResult<Self> {
285        self.cast_with_options(dtype, CastOptions::NonStrict)
286    }
287    /// # Safety
288    ///
289    /// This can lead to invalid memory access in downstream code.
290    pub unsafe fn cast_unchecked(&self, dtype: &DataType) -> PolarsResult<Self> {
291        // @NOTE: We expect that when casting the materialized series mostly does not need change
292        // the physical array. Therefore, we try to cast the entire materialized array if it is
293        // available.
294
295        match self.materialized.get() {
296            Some(s) => {
297                let materialized = s.cast_unchecked(dtype)?;
298                assert_eq!(self.length, materialized.len());
299
300                let mut casted = if materialized.is_empty() {
301                    Self::new_empty(materialized.name().clone(), materialized.dtype().clone())
302                } else {
303                    // SAFETY: Just did bounds check
304                    let scalar = unsafe { materialized.get_unchecked(0) }.into_static();
305                    Self::new(
306                        materialized.name().clone(),
307                        Scalar::new(materialized.dtype().clone(), scalar),
308                        self.length,
309                    )
310                };
311                casted.materialized = OnceLock::from(materialized);
312                Ok(casted)
313            },
314            None => {
315                let s = self.as_single_value_series().cast_unchecked(dtype)?;
316                assert_eq!(1, s.len());
317
318                if self.length == 0 {
319                    Ok(Self::new_empty(s.name().clone(), s.dtype().clone()))
320                } else {
321                    Ok(Self::from_single_value_series(s, self.length))
322                }
323            },
324        }
325    }
326
327    pub fn rename(&mut self, name: PlSmallStr) -> &mut Self {
328        if let Some(series) = self.materialized.get_mut() {
329            series.rename(name.clone());
330        }
331
332        self.name = name;
333        self
334    }
335
336    pub fn has_nulls(&self) -> bool {
337        self.length != 0 && self.scalar.is_null()
338    }
339
340    pub fn drop_nulls(&self) -> Self {
341        if self.scalar.is_null() {
342            self.resize(0)
343        } else {
344            self.clone()
345        }
346    }
347
348    pub fn into_nulls(mut self) -> Self {
349        self.scalar.update(AnyValue::Null);
350        self
351    }
352
353    /// Packs every element into a single-element list.
354    pub fn to_unit_list(&self) -> Self {
355        let mut slf = self.clone();
356        slf.map_scalar(|s| Scalar::new_list(s.into_series(PlSmallStr::EMPTY)));
357        slf
358    }
359
360    pub fn map_scalar(&mut self, map_scalar: impl Fn(Scalar) -> Scalar) {
361        self.scalar = map_scalar(std::mem::take(&mut self.scalar));
362        self.materialized.take();
363    }
364    pub fn with_value(&mut self, value: AnyValue<'static>) -> &mut Self {
365        self.scalar.update(value);
366        self.materialized.take();
367        self
368    }
369}
370
371/// Whether `l` and `r`, both of `dtype`, are the same value and not only equal ones.
372///
373/// For floats `0.0 == -0.0`, and for objects `==` calls into Python, so those are only
374/// treated as the same when that is cheap to prove.
375fn is_same_value(dtype: &DataType, l: &AnyValue, r: &AnyValue) -> bool {
376    match (l, r) {
377        (AnyValue::Null, AnyValue::Null) => true,
378        (AnyValue::Float16(l), AnyValue::Float16(r)) => l.to_bits() == r.to_bits(),
379        (AnyValue::Float32(l), AnyValue::Float32(r)) => l.to_bits() == r.to_bits(),
380        (AnyValue::Float64(l), AnyValue::Float64(r)) => l.to_bits() == r.to_bits(),
381        _ => {
382            let mut eq_is_same = true;
383            dtype.visit_with(|dtype| eq_is_same &= !(dtype.is_float() || dtype.is_object()));
384            eq_is_same && l == r
385        },
386    }
387}
388
389/// [`Series::estimated_size`] of `value` repeated `length` times, as built by
390/// [`Series::new_from_index`].
391///
392/// Returns `None` if the size cannot be derived from `dtype` and `value`.
393fn estimated_repeated_size(dtype: &DataType, value: &AnyValue, length: usize) -> Option<usize> {
394    let is_null = value.is_null();
395    let validity_size = if is_null { length.div_ceil(8) } else { 0 };
396    let size = match dtype {
397        DataType::Null => 0,
398        DataType::Boolean => length.div_ceil(8) + validity_size,
399        // Views and validity are not counted for these types, see `estimated_bytes_size`.
400        DataType::String | DataType::Binary if is_null => 0,
401        DataType::String => length * value.extract_str()?.len(),
402        DataType::Binary => length * value.extract_bytes()?.len(),
403        DataType::List(_) => match value {
404            AnyValue::List(s) => {
405                estimated_repeated_series_size(s, length) + length * size_of::<i64>()
406            },
407            _ if is_null => length * size_of::<i64>() + validity_size,
408            _ => return None,
409        },
410        #[cfg(feature = "dtype-array")]
411        DataType::Array(inner, width) => match value {
412            AnyValue::Array(s, _) => estimated_repeated_series_size(s, length),
413            _ if is_null => {
414                estimated_repeated_size(inner, &AnyValue::Null, length * width)? + validity_size
415            },
416            _ => return None,
417        },
418        #[cfg(feature = "dtype-struct")]
419        DataType::Struct(fields) => match value {
420            AnyValue::StructOwned(payload) => payload
421                .0
422                .iter()
423                .zip(fields)
424                .map(|(value, field)| estimated_repeated_size(field.dtype(), value, length))
425                .sum::<Option<usize>>()?,
426            _ if is_null => {
427                fields
428                    .iter()
429                    .map(|field| estimated_repeated_size(field.dtype(), &AnyValue::Null, length))
430                    .sum::<Option<usize>>()?
431                    + validity_size
432            },
433            _ => return None,
434        },
435        #[cfg(feature = "dtype-extension")]
436        DataType::Extension(_, storage) => return estimated_repeated_size(storage, value, length),
437        _ => length * dtype.byte_width()? as usize + validity_size,
438    };
439    Some(size)
440}
441
442/// [`Series::estimated_size`] of `s` repeated `n` times.
443fn estimated_repeated_series_size(s: &Series, n: usize) -> usize {
444    let dtype = s.dtype();
445    match dtype.byte_width() {
446        // Bitmaps round up to whole bytes, so scaling the size of `s` would count too much.
447        Some(width) if !dtype.is_nested() => {
448            let len = n * s.len();
449            let validity_size = if s.has_nulls() { len.div_ceil(8) } else { 0 };
450            (len as f64 * width).ceil() as usize + validity_size
451        },
452        _ => n * s.estimated_size(),
453    }
454}
455
456impl IntoColumn for ScalarColumn {
457    #[inline(always)]
458    fn into_column(self) -> Column {
459        self.into()
460    }
461}
462
463impl From<ScalarColumn> for Column {
464    #[inline]
465    fn from(value: ScalarColumn) -> Self {
466        Self::Scalar(value)
467    }
468}
469
470impl BroadcastLength for ScalarColumn {
471    fn _broadcast_len(&self) -> usize {
472        self.len()
473    }
474
475    fn _column_name(&self) -> Option<&str> {
476        Some(self.name())
477    }
478}
479
480#[cfg(feature = "dsl-schema")]
481impl schemars::JsonSchema for ScalarColumn {
482    fn schema_name() -> std::borrow::Cow<'static, str> {
483        "ScalarColumn".into()
484    }
485
486    fn schema_id() -> std::borrow::Cow<'static, str> {
487        std::borrow::Cow::Borrowed(concat!(module_path!(), "::", "ScalarColumn"))
488    }
489
490    fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
491        serde_impl::SerializeWrap::json_schema(generator)
492    }
493}
494
495#[cfg(feature = "serde")]
496mod serde_impl {
497    use std::sync::OnceLock;
498
499    use polars_error::PolarsError;
500    use polars_utils::pl_str::PlSmallStr;
501
502    use super::ScalarColumn;
503    use crate::frame::{Scalar, Series};
504
505    #[derive(serde::Serialize, serde::Deserialize)]
506    #[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
507    pub struct SerializeWrap {
508        name: PlSmallStr,
509        /// Unit-length series for dispatching to IPC serialize
510        unit_series: Series,
511        length: usize,
512    }
513
514    impl From<&ScalarColumn> for SerializeWrap {
515        fn from(value: &ScalarColumn) -> Self {
516            Self {
517                name: value.name.clone(),
518                unit_series: value.scalar.clone().into_series(PlSmallStr::EMPTY),
519                length: value.length,
520            }
521        }
522    }
523
524    impl TryFrom<SerializeWrap> for ScalarColumn {
525        type Error = PolarsError;
526
527        fn try_from(value: SerializeWrap) -> Result<Self, Self::Error> {
528            let slf = Self {
529                name: value.name,
530                scalar: Scalar::new(
531                    value.unit_series.dtype().clone(),
532                    value.unit_series.get(0)?.into_static(),
533                ),
534                length: value.length,
535                materialized: OnceLock::new(),
536            };
537
538            Ok(slf)
539        }
540    }
541
542    impl serde::ser::Serialize for ScalarColumn {
543        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
544        where
545            S: serde::Serializer,
546        {
547            SerializeWrap::from(self).serialize(serializer)
548        }
549    }
550
551    impl<'de> serde::de::Deserialize<'de> for ScalarColumn {
552        fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
553        where
554            D: serde::Deserializer<'de>,
555        {
556            use serde::de::Error;
557
558            SerializeWrap::deserialize(deserializer)
559                .and_then(|x| ScalarColumn::try_from(x).map_err(D::Error::custom))
560        }
561    }
562}
563
564#[cfg(test)]
565mod test {
566    use super::*;
567    use crate::prelude::*;
568
569    fn check_estimated_size(s: Series) {
570        let scalar = Scalar::new(s.dtype().clone(), s.get(0).unwrap().into_static());
571        let unit_size = ScalarColumn::new(s.name().clone(), scalar.clone(), 1)
572            .to_series()
573            .estimated_size();
574        let sc = ScalarColumn::new(s.name().clone(), scalar, 1001);
575        assert_eq!(sc.estimated_size(true), sc.to_series().estimated_size());
576        assert_eq!(sc.estimated_size(false), unit_size);
577
578        let materialized_size = sc.as_materialized_series().estimated_size();
579        assert_eq!(sc.estimated_size(false), materialized_size);
580    }
581
582    #[test]
583    fn test_estimated_size() -> PolarsResult<()> {
584        let list = Series::new("a".into(), [Series::new("".into(), [Some(1i32), None])]);
585        let mut series = vec![
586            Series::new("a".into(), [1i64]),
587            Series::new("a".into(), [true]),
588            Series::new("a".into(), ["a".repeat(20)]),
589            Series::new("a".into(), [b"ab".as_slice()]),
590            list.clone(),
591        ];
592        #[cfg(feature = "dtype-array")]
593        series.push(list.cast(&DataType::Array(Box::new(DataType::Int32), 2))?);
594        #[cfg(feature = "dtype-struct")]
595        series.push(
596            StructChunked::from_series(
597                "a".into(),
598                1,
599                [Series::new("x".into(), [1i8]), list.clone()].iter(),
600            )?
601            .into_series(),
602        );
603
604        for s in series {
605            check_estimated_size(Series::full_null("a".into(), 1, s.dtype()));
606            check_estimated_size(s);
607        }
608        Ok(())
609    }
610}
polars_core/frame/column/scalar.rs

polars_core/frame/column/
scalar.rs
