polars_core/frame/row/
mod.rs

1mod av_buffer;
2mod dataframe;
3mod transpose;
4
5use std::borrow::Borrow;
6use std::fmt::Debug;
7#[cfg(feature = "object")]
8use std::hash::{Hash, Hasher};
9
10pub use av_buffer::*;
11pub use dataframe::check_row_width;
12use polars_arrow::bitmap::Bitmap;
13use polars_utils::format_pl_smallstr;
14#[cfg(feature = "object")]
15use polars_utils::total_ord::TotalHash;
16use rayon::prelude::*;
17
18use crate::prelude::*;
19use crate::runtime::RAYON;
20use crate::utils::{dtypes_to_schema, dtypes_to_supertype, try_get_supertype};
21
22#[cfg(feature = "object")]
23pub(crate) struct AnyValueRows<'a> {
24    vals: Vec<AnyValue<'a>>,
25    width: usize,
26}
27
28#[cfg(feature = "object")]
29pub(crate) struct AnyValueRow<'a>(&'a [AnyValue<'a>]);
30
31#[cfg(feature = "object")]
32impl<'a> AnyValueRows<'a> {
33    pub(crate) fn get(&'a self, i: usize) -> AnyValueRow<'a> {
34        let start = i * self.width;
35        let end = (i + 1) * self.width;
36        AnyValueRow(&self.vals[start..end])
37    }
38}
39
40#[cfg(feature = "object")]
41impl TotalEq for AnyValueRow<'_> {
42    fn tot_eq(&self, other: &Self) -> bool {
43        let lhs = self.0;
44        let rhs = other.0;
45
46        // Should only be used in that context.
47        debug_assert_eq!(lhs.len(), rhs.len());
48        lhs.iter().zip(rhs.iter()).all(|(l, r)| l == r)
49    }
50}
51
52#[cfg(feature = "object")]
53impl TotalHash for AnyValueRow<'_> {
54    fn tot_hash<H>(&self, state: &mut H)
55    where
56        H: Hasher,
57    {
58        self.0.iter().for_each(|av| av.hash(state))
59    }
60}
61
62impl DataFrame {
63    #[cfg(feature = "object")]
64    #[allow(clippy::wrong_self_convention)]
65    // Create indexable rows in a single allocation.
66    pub(crate) fn to_av_rows(&mut self) -> AnyValueRows<'_> {
67        let width = self.width();
68        let size = width * self.height();
69        let mut buf = vec![AnyValue::Null; size];
70        for (col_i, s) in self.materialized_column_iter().enumerate() {
71            for (row_i, av) in s.iter().enumerate() {
72                buf[row_i * width + col_i] = av
73            }
74        }
75        AnyValueRows { vals: buf, width }
76    }
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Default)]
80pub struct Row<'a>(pub Vec<AnyValue<'a>>);
81
82impl<'a> Row<'a> {
83    pub fn new(values: Vec<AnyValue<'a>>) -> Self {
84        Row(values)
85    }
86}
87
88type Tracker = PlIndexMap<PlSmallStr, PlHashSet<DataType>>;
89
90pub fn infer_schema(
91    iter: impl Iterator<Item = Vec<(impl Into<PlSmallStr>, impl Into<DataType>)>>,
92    infer_schema_length: usize,
93) -> Schema {
94    let mut values: Tracker = Tracker::default();
95    let len = iter.size_hint().1.unwrap_or(infer_schema_length);
96
97    let max_infer = std::cmp::min(len, infer_schema_length);
98    for inner in iter.take(max_infer) {
99        for (key, value) in inner {
100            add_or_insert(&mut values, key.into(), value.into());
101        }
102    }
103    Schema::from_iter(resolve_fields(values))
104}
105
106fn add_or_insert(values: &mut Tracker, key: PlSmallStr, dtype: DataType) {
107    if values.contains_key(&key) {
108        let x = values.get_mut(&key).unwrap();
109        x.insert(dtype);
110    } else {
111        // create hashset and add value type
112        let mut hs = PlHashSet::new();
113        hs.insert(dtype);
114        values.insert(key, hs);
115    }
116}
117
118fn resolve_fields(spec: Tracker) -> Vec<Field> {
119    spec.iter()
120        .map(|(k, hs)| {
121            let v: Vec<&DataType> = hs.iter().collect();
122            Field::new(k.clone(), coerce_dtype(&v))
123        })
124        .collect()
125}
126
127/// Coerces a slice of datatypes into a single supertype.
128pub fn coerce_dtype<A: Borrow<DataType>>(datatypes: &[A]) -> DataType {
129    use DataType::*;
130
131    let are_all_equal = datatypes.windows(2).all(|w| w[0].borrow() == w[1].borrow());
132
133    if are_all_equal {
134        return datatypes[0].borrow().clone();
135    }
136    if datatypes.len() > 2 {
137        return String;
138    }
139
140    let (lhs, rhs) = (datatypes[0].borrow(), datatypes[1].borrow());
141    try_get_supertype(lhs, rhs).unwrap_or(String)
142}
143
144/// Infer the schema of rows by determining the supertype of the values.
145///
146/// Field names are set as `column_0`, `column_1`, and so on.
147pub fn rows_to_schema_supertypes(
148    rows: &[Row],
149    infer_schema_length: Option<usize>,
150) -> PolarsResult<Schema> {
151    let dtypes = rows_to_supertypes(rows, infer_schema_length)?;
152    let schema = dtypes_to_schema(dtypes);
153    Ok(schema)
154}
155
156/// Infer the schema data types of rows by determining the supertype of the values.
157pub fn rows_to_supertypes(
158    rows: &[Row],
159    infer_schema_length: Option<usize>,
160) -> PolarsResult<Vec<DataType>> {
161    polars_ensure!(!rows.is_empty(), NoData: "no rows, cannot infer schema");
162
163    let max_infer = infer_schema_length.unwrap_or(rows.len());
164
165    let mut dtypes: Vec<PlIndexSet<DataType>> = vec![PlIndexSet::new(); rows[0].0.len()];
166    for row in rows.iter().take(max_infer) {
167        for (val, dtypes_set) in row.0.iter().zip(dtypes.iter_mut()) {
168            dtypes_set.insert(val.into());
169        }
170    }
171
172    dtypes
173        .into_iter()
174        .map(|dtypes_set| dtypes_to_supertype(&dtypes_set))
175        .collect()
176}
177
178/// Infer schema from rows and set the first no null type as column data type.
179pub fn rows_to_schema_first_non_null(
180    rows: &[Row],
181    infer_schema_length: Option<usize>,
182) -> PolarsResult<Schema> {
183    polars_ensure!(!rows.is_empty(), NoData: "no rows, cannot infer schema");
184
185    let max_infer = infer_schema_length.unwrap_or(rows.len());
186    let mut schema: Schema = (&rows[0]).into();
187
188    // the first row that has no nulls will be used to infer the schema.
189    // if there is a null, we check the next row and see if we can update the schema
190
191    for row in rows.iter().take(max_infer).skip(1) {
192        // for i in 1..max_infer {
193        let nulls: Vec<_> = schema
194            .iter_values()
195            .enumerate()
196            .filter_map(|(i, dtype)| {
197                // double check struct and list types
198                // nested null values can be wrongly inferred by front ends
199                match dtype {
200                    DataType::Null | DataType::List(_) => Some(i),
201                    #[cfg(feature = "dtype-struct")]
202                    DataType::Struct(_) => Some(i),
203                    _ => None,
204                }
205            })
206            .collect();
207        if nulls.is_empty() {
208            break;
209        } else {
210            for i in nulls {
211                let val = &row.0[i];
212
213                if !val.is_nested_null() {
214                    let dtype = val.into();
215                    schema.set_dtype_at_index(i, dtype).unwrap();
216                }
217            }
218        }
219    }
220    Ok(schema)
221}
222
223impl<'a> From<&AnyValue<'a>> for Field {
224    fn from(val: &AnyValue<'a>) -> Self {
225        Field::new(PlSmallStr::EMPTY, val.into())
226    }
227}
228
229impl From<&Row<'_>> for Schema {
230    fn from(row: &Row) -> Self {
231        row.0
232            .iter()
233            .enumerate()
234            .map(|(i, av)| {
235                let dtype = av.into();
236                Field::new(format_pl_smallstr!("column_{i}"), dtype)
237            })
238            .collect()
239    }
240}
polars_core/frame/row/mod.rs

polars_core/frame/row/
mod.rs
