polars_core/chunked_array/
collect.rsuse std::sync::Arc;
use arrow::trusted_len::TrustedLen;
use polars_utils::pl_str::PlSmallStr;
use crate::chunked_array::ChunkedArray;
use crate::datatypes::{
ArrayCollectIterExt, ArrayFromIter, ArrayFromIterDtype, DataType, Field, PolarsDataType,
};
use crate::prelude::CompatLevel;
pub trait ChunkedCollectIterExt<T: PolarsDataType>: Iterator + Sized {
#[inline]
fn collect_ca_with_dtype(self, name: PlSmallStr, dtype: DataType) -> ChunkedArray<T>
where
T::Array: ArrayFromIterDtype<Self::Item>,
{
let field = Arc::new(Field::new(name, dtype.clone()));
let arr = self.collect_arr_with_dtype(field.dtype.to_arrow(CompatLevel::newest()));
ChunkedArray::from_chunk_iter_and_field(field, [arr])
}
#[inline]
fn collect_ca_like(self, name_dtype_src: &ChunkedArray<T>) -> ChunkedArray<T>
where
T::Array: ArrayFromIterDtype<Self::Item>,
{
let field = Arc::clone(&name_dtype_src.field);
let arr = self.collect_arr_with_dtype(field.dtype.to_arrow(CompatLevel::newest()));
ChunkedArray::from_chunk_iter_and_field(field, [arr])
}
#[inline]
fn collect_ca_trusted_with_dtype(self, name: PlSmallStr, dtype: DataType) -> ChunkedArray<T>
where
T::Array: ArrayFromIterDtype<Self::Item>,
Self: TrustedLen,
{
let field = Arc::new(Field::new(name, dtype.clone()));
let arr = self.collect_arr_trusted_with_dtype(field.dtype.to_arrow(CompatLevel::newest()));
ChunkedArray::from_chunk_iter_and_field(field, [arr])
}
#[inline]
fn collect_ca_trusted_like(self, name_dtype_src: &ChunkedArray<T>) -> ChunkedArray<T>
where
T::Array: ArrayFromIterDtype<Self::Item>,
Self: TrustedLen,
{
let field = Arc::clone(&name_dtype_src.field);
let arr = self.collect_arr_trusted_with_dtype(field.dtype.to_arrow(CompatLevel::newest()));
ChunkedArray::from_chunk_iter_and_field(field, [arr])
}
#[inline]
fn try_collect_ca_with_dtype<U, E>(
self,
name: PlSmallStr,
dtype: DataType,
) -> Result<ChunkedArray<T>, E>
where
T::Array: ArrayFromIterDtype<U>,
Self: Iterator<Item = Result<U, E>>,
{
let field = Arc::new(Field::new(name, dtype.clone()));
let arr = self.try_collect_arr_with_dtype(field.dtype.to_arrow(CompatLevel::newest()))?;
Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
}
#[inline]
fn try_collect_ca_like<U, E>(
self,
name_dtype_src: &ChunkedArray<T>,
) -> Result<ChunkedArray<T>, E>
where
T::Array: ArrayFromIterDtype<U>,
Self: Iterator<Item = Result<U, E>>,
{
let field = Arc::clone(&name_dtype_src.field);
let arr = self.try_collect_arr_with_dtype(field.dtype.to_arrow(CompatLevel::newest()))?;
Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
}
#[inline]
fn try_collect_ca_trusted_with_dtype<U, E>(
self,
name: PlSmallStr,
dtype: DataType,
) -> Result<ChunkedArray<T>, E>
where
T::Array: ArrayFromIterDtype<U>,
Self: Iterator<Item = Result<U, E>> + TrustedLen,
{
let field = Arc::new(Field::new(name, dtype.clone()));
let arr =
self.try_collect_arr_trusted_with_dtype(field.dtype.to_arrow(CompatLevel::newest()))?;
Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
}
#[inline]
fn try_collect_ca_trusted_like<U, E>(
self,
name_dtype_src: &ChunkedArray<T>,
) -> Result<ChunkedArray<T>, E>
where
T::Array: ArrayFromIterDtype<U>,
Self: Iterator<Item = Result<U, E>> + TrustedLen,
{
let field = Arc::clone(&name_dtype_src.field);
let arr =
self.try_collect_arr_trusted_with_dtype(field.dtype.to_arrow(CompatLevel::newest()))?;
Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
}
}
impl<T: PolarsDataType, I: Iterator> ChunkedCollectIterExt<T> for I {}
pub trait ChunkedCollectInferIterExt<T: PolarsDataType>: Iterator + Sized {
#[inline]
fn collect_ca(self, name: PlSmallStr) -> ChunkedArray<T>
where
T::Array: ArrayFromIter<Self::Item>,
{
let field = Arc::new(Field::new(name, T::get_dtype()));
let arr = self.collect_arr();
ChunkedArray::from_chunk_iter_and_field(field, [arr])
}
#[inline]
fn collect_ca_trusted(self, name: PlSmallStr) -> ChunkedArray<T>
where
T::Array: ArrayFromIter<Self::Item>,
Self: TrustedLen,
{
let field = Arc::new(Field::new(name, T::get_dtype()));
let arr = self.collect_arr_trusted();
ChunkedArray::from_chunk_iter_and_field(field, [arr])
}
#[inline]
fn try_collect_ca<U, E>(self, name: PlSmallStr) -> Result<ChunkedArray<T>, E>
where
T::Array: ArrayFromIter<U>,
Self: Iterator<Item = Result<U, E>>,
{
let field = Arc::new(Field::new(name, T::get_dtype()));
let arr = self.try_collect_arr()?;
Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
}
#[inline]
fn try_collect_ca_trusted<U, E>(self, name: PlSmallStr) -> Result<ChunkedArray<T>, E>
where
T::Array: ArrayFromIter<U>,
Self: Iterator<Item = Result<U, E>> + TrustedLen,
{
let field = Arc::new(Field::new(name, T::get_dtype()));
let arr = self.try_collect_arr_trusted()?;
Ok(ChunkedArray::from_chunk_iter_and_field(field, [arr]))
}
}
impl<T: PolarsDataType, I: Iterator> ChunkedCollectInferIterExt<T> for I {}