polars_core/chunked_array/ops/
row_encode.rs

1use std::borrow::Cow;
2
3use polars_row::{
4    RowEncodingContext, RowEncodingField, RowEncodingLayout, RowEncodingOptions, RowEncodingOutput,
5    convert_columns,
6};
7use polars_utils::itertools::Itertools;
8use rayon::prelude::*;
9
10use crate::prelude::*;
11use crate::runtime::RAYON;
12use crate::utils::_split_offsets;
13
14pub fn encode_rows_vertical_par_unordered(by: &[Column]) -> PolarsResult<BinaryOffsetChunked> {
15    encode_binary_offset_par(by, &unordered_opts(by.len()), false)
16}
17
18pub fn encode_rows_vertical_par_unordered_propagate_nulls(
19    by: &[Column],
20) -> PolarsResult<BinaryOffsetChunked> {
21    encode_binary_offset_par(by, &unordered_opts(by.len()), true)
22}
23
24/// Get the [`RowEncodingContext`] for a certain [`DataType`].
25///
26/// This should be given the logical type in order to communicate Polars datatype information down
27/// into the row encoding / decoding.
28pub fn get_row_encoding_context(dtype: &DataType) -> Option<RowEncodingContext> {
29    match dtype {
30        DataType::Boolean
31        | DataType::UInt8
32        | DataType::UInt16
33        | DataType::UInt32
34        | DataType::UInt64
35        | DataType::UInt128
36        | DataType::Int8
37        | DataType::Int16
38        | DataType::Int32
39        | DataType::Int64
40        | DataType::Int128
41        | DataType::Float16
42        | DataType::Float32
43        | DataType::Float64
44        | DataType::String
45        | DataType::Binary
46        | DataType::BinaryOffset
47        | DataType::Null
48        | DataType::Time
49        | DataType::Date
50        | DataType::Datetime(_, _)
51        | DataType::Duration(_) => None,
52
53        #[cfg(feature = "dtype-categorical")]
54        DataType::Categorical(_, mapping) | DataType::Enum(_, mapping) => {
55            use polars_row::RowEncodingCategoricalContext;
56
57            Some(RowEncodingContext::Categorical(
58                RowEncodingCategoricalContext {
59                    is_enum: matches!(dtype, DataType::Enum(_, _)),
60                    mapping: mapping.clone(),
61                },
62            ))
63        },
64
65        DataType::Unknown(_) => panic!("Unsupported in row encoding"),
66
67        #[cfg(feature = "object")]
68        DataType::Object(_) => panic!("Unsupported in row encoding"),
69
70        #[cfg(feature = "dtype-decimal")]
71        DataType::Decimal(precision, _) => Some(RowEncodingContext::Decimal(*precision)),
72
73        #[cfg(feature = "dtype-array")]
74        DataType::Array(dtype, _) => get_row_encoding_context(dtype),
75        DataType::List(dtype) => get_row_encoding_context(dtype),
76        #[cfg(feature = "dtype-struct")]
77        DataType::Struct(fs) => {
78            let mut ctxts = Vec::new();
79
80            for (i, f) in fs.iter().enumerate() {
81                if let Some(ctxt) = get_row_encoding_context(f.dtype()) {
82                    ctxts.reserve(fs.len());
83                    ctxts.extend(std::iter::repeat_n(None, i));
84                    ctxts.push(Some(ctxt));
85                    break;
86                }
87            }
88
89            if ctxts.is_empty() {
90                return None;
91            }
92
93            ctxts.extend(
94                fs[ctxts.len()..]
95                    .iter()
96                    .map(|f| get_row_encoding_context(f.dtype())),
97            );
98
99            Some(RowEncodingContext::Struct(ctxts))
100        },
101
102        #[cfg(feature = "dtype-map")]
103        DataType::Map(key, value) => {
104            let ctxts = vec![
105                get_row_encoding_context(key),
106                get_row_encoding_context(value),
107            ];
108
109            if ctxts.iter().all(Option::is_none) {
110                return None;
111            }
112
113            Some(RowEncodingContext::Struct(ctxts))
114        },
115
116        #[cfg(feature = "dtype-extension")]
117        DataType::Extension(_, storage) => get_row_encoding_context(storage),
118    }
119}
120
121/// The layout for row encoding columns of these logical dtypes. Without `dtype-u128`,
122/// [`RowEncodingOutput::Optimal`] rows that would be `UInt128` are `LargeBinary` instead.
123pub fn get_row_encoding_layout(
124    dtypes: &[DataType],
125    opts: &[RowEncodingOptions],
126    output: RowEncodingOutput,
127) -> RowEncodingLayout {
128    assert_eq!(dtypes.len(), opts.len());
129    let fields = dtypes
130        .iter()
131        .zip(opts)
132        .map(|(dtype, opt)| RowEncodingField {
133            dtype: dtype.to_physical().to_arrow(CompatLevel::newest()),
134            opt: *opt,
135            ctx: get_row_encoding_context(dtype),
136        })
137        .collect_vec();
138    let layout = RowEncodingLayout::new(fields, output);
139    if *layout.dtype() == ArrowDataType::UInt128 && !cfg!(feature = "dtype-u128") {
140        return RowEncodingLayout::new(
141            layout.fields().to_vec(),
142            RowEncodingOutput::AlwaysBinaryOffset,
143        );
144    }
145    layout
146}
147
148/// The dtype of the rows of `layout`.
149pub fn rows_dtype(layout: &RowEncodingLayout) -> DataType {
150    match layout.dtype() {
151        ArrowDataType::LargeBinary => DataType::BinaryOffset,
152        dtype => DataType::from_arrow_dtype(dtype),
153    }
154}
155
156fn encode_columns(
157    by: &[Column],
158    layout: &RowEncodingLayout,
159    propagate_nulls: bool,
160) -> PolarsResult<ArrayRef> {
161    // Since ZFS exists, we might not actually have any arrays and need to get the length from the
162    // columns.
163    let num_rows = by.first().map_or(0, |c| c.len());
164
165    let mut cols = Vec::with_capacity(by.len());
166    for (by, field) in by.iter().zip(layout.fields()) {
167        polars_ensure!(
168            field.opt.encode_nulls() || propagate_nulls || !by.has_nulls(),
169            ComputeError: "row encoding: column '{}' contains nulls, but its nulls are not encoded", by.name()
170        );
171
172        let by = by
173            .trim_lists_to_normalized_offsets()
174            .map_or(Cow::Borrowed(by), Cow::Owned);
175        let by = by.propagate_nulls().map_or(by, Cow::Owned);
176        let by = by.as_materialized_series();
177        cols.push(by.to_physical_repr().rechunk().chunks()[0].to_boxed());
178    }
179    Ok(convert_columns(num_rows, &cols, layout, propagate_nulls))
180}
181
182fn unordered_opts(n: usize) -> Vec<RowEncodingOptions> {
183    vec![RowEncodingOptions::new_unsorted(); n]
184}
185
186pub(crate) fn ordered_opts(descending: &[bool], nulls_last: &[bool]) -> Vec<RowEncodingOptions> {
187    debug_assert_eq!(descending.len(), nulls_last.len());
188    descending
189        .iter()
190        .zip(nulls_last)
191        .map(|(desc, null_last)| RowEncodingOptions::new_sorted(*desc, *null_last))
192        .collect_vec()
193}
194
195/// The [`RowEncodingOutput::AlwaysBinaryOffset`] layout for these columns.
196pub fn get_binary_offset_row_encoding_layout(
197    by: &[Column],
198    opts: &[RowEncodingOptions],
199) -> RowEncodingLayout {
200    let dtypes = by.iter().map(|c| c.dtype().clone()).collect_vec();
201    get_row_encoding_layout(&dtypes, opts, RowEncodingOutput::AlwaysBinaryOffset)
202}
203
204fn encode_binary_offset(
205    by: &[Column],
206    opts: &[RowEncodingOptions],
207    propagate_nulls: bool,
208) -> PolarsResult<BinaryArray<i64>> {
209    let rows = encode_columns(
210        by,
211        &get_binary_offset_row_encoding_layout(by, opts),
212        propagate_nulls,
213    )?;
214    Ok(rows
215        .as_any()
216        .downcast_ref::<BinaryArray<i64>>()
217        .unwrap()
218        .clone())
219}
220
221fn encode_binary_offset_par(
222    by: &[Column],
223    opts: &[RowEncodingOptions],
224    propagate_nulls: bool,
225) -> PolarsResult<BinaryOffsetChunked> {
226    let rows = encode_rows_with_layout_par(
227        by,
228        &get_binary_offset_row_encoding_layout(by, opts),
229        propagate_nulls,
230    )?;
231    Ok(rows.binary_offset()?.clone())
232}
233
234/// Encodes `by` into rows of `layout`, which must be made for the dtypes of `by`. With
235/// `propagate_nulls` a row is null if any of its values is null.
236pub fn encode_rows_with_layout(
237    by: &[Column],
238    layout: &RowEncodingLayout,
239    propagate_nulls: bool,
240) -> PolarsResult<Series> {
241    let rows = encode_columns(by, layout, propagate_nulls)?;
242    // SAFETY: the rows have the dtype of the layout.
243    Ok(unsafe {
244        Series::from_chunks_and_dtype_unchecked(PlSmallStr::EMPTY, vec![rows], &rows_dtype(layout))
245    })
246}
247
248/// Below this number of rows, encoding rows in parallel isn't worth it. Low in debug builds so
249/// that tests also cover multi-chunk rows.
250const ROW_ENCODE_PAR_MIN_ROWS: usize = if cfg!(debug_assertions) { 4 } else { 1 << 16 };
251
252/// [`encode_rows_with_layout`], encoding a slice of the rows per thread if there are enough rows.
253pub fn encode_rows_with_layout_par(
254    by: &[Column],
255    layout: &RowEncodingLayout,
256    propagate_nulls: bool,
257) -> PolarsResult<Series> {
258    let len = by.first().map_or(0, |c| c.len());
259    let n_splits = RAYON.current_num_threads().min(len).max(1);
260    if len < ROW_ENCODE_PAR_MIN_ROWS || n_splits == 1 {
261        return encode_rows_with_layout(by, layout, propagate_nulls);
262    }
263    let chunks = _split_offsets(len, n_splits)
264        .into_par_iter()
265        .map(|(offset, len)| {
266            let sliced = by
267                .iter()
268                .map(|c| c.slice(offset as i64, len))
269                .collect::<Vec<_>>();
270            encode_columns(&sliced, layout, propagate_nulls)
271        });
272    let chunks = RAYON.install(|| chunks.collect::<PolarsResult<Vec<_>>>())?;
273    // SAFETY: the rows have the dtype of the layout.
274    Ok(unsafe {
275        Series::from_chunks_and_dtype_unchecked(PlSmallStr::EMPTY, chunks, &rows_dtype(layout))
276    })
277}
278
279/// The [`RowEncodingOutput::Optimal`] layout for these columns, with each column encoding nulls
280/// if it has nulls in `by` or in any of `others`.
281pub fn get_optimal_row_encoding_layout(
282    by: &[Column],
283    others: &[&[Column]],
284    opts: &[RowEncodingOptions],
285) -> RowEncodingLayout {
286    let dtypes = by.iter().map(|c| c.dtype().clone()).collect_vec();
287    let opts = opts
288        .iter()
289        .enumerate()
290        .map(|(i, opt)| {
291            let encode_nulls = by[i].has_nulls() || others.iter().any(|o| o[i].has_nulls());
292            let mut opt = *opt;
293            opt.set(RowEncodingOptions::SKIP_NULL_ENCODING, !encode_nulls);
294            opt
295        })
296        .collect_vec();
297    get_row_encoding_layout(&dtypes, &opts, RowEncodingOutput::Optimal)
298}
299
300pub fn encode_rows_unordered(by: &[Column]) -> PolarsResult<BinaryOffsetChunked> {
301    let rows = _get_rows_encoded_unordered(by, false)?;
302    Ok(BinaryOffsetChunked::with_chunk(PlSmallStr::EMPTY, rows))
303}
304
305/// With `propagate_nulls` a row is null if any of its values is null.
306pub fn _get_rows_encoded_unordered(
307    by: &[Column],
308    propagate_nulls: bool,
309) -> PolarsResult<BinaryArray<i64>> {
310    encode_binary_offset(by, &unordered_opts(by.len()), propagate_nulls)
311}
312
313pub fn _get_rows_encoded(
314    by: &[Column],
315    descending: &[bool],
316    nulls_last: &[bool],
317) -> PolarsResult<BinaryArray<i64>> {
318    _get_rows_encoded_arr(by, descending, nulls_last, false)
319}
320
321pub fn _get_rows_encoded_ca(
322    name: PlSmallStr,
323    by: &[Column],
324    descending: &[bool],
325    nulls_last: &[bool],
326    propagate_nulls: bool,
327) -> PolarsResult<BinaryOffsetChunked> {
328    let rows_arr = _get_rows_encoded_arr(by, descending, nulls_last, propagate_nulls)?;
329    Ok(BinaryOffsetChunked::with_chunk(name, rows_arr))
330}
331
332pub fn _get_rows_encoded_arr(
333    by: &[Column],
334    descending: &[bool],
335    nulls_last: &[bool],
336    propagate_nulls: bool,
337) -> PolarsResult<BinaryArray<i64>> {
338    debug_assert_eq!(by.len(), descending.len());
339    encode_binary_offset(by, &ordered_opts(descending, nulls_last), propagate_nulls)
340}
341
342pub fn _get_rows_encoded_ca_unordered(
343    name: PlSmallStr,
344    by: &[Column],
345) -> PolarsResult<BinaryOffsetChunked> {
346    _get_rows_encoded_unordered(by, false).map(|rows| BinaryOffsetChunked::with_chunk(name, rows))
347}
348
349/// Decodes rows of `layout` into a struct with these fields.
350#[cfg(feature = "dtype-struct")]
351pub fn row_encoding_decode(
352    rows: &Series,
353    fields: &[Field],
354    layout: &RowEncodingLayout,
355) -> PolarsResult<StructChunked> {
356    let expected = rows_dtype(layout);
357    polars_ensure!(
358        rows.dtype() == &expected,
359        SchemaMismatch: "row decoding expected rows of type {}, got {}", expected, rows.dtype()
360    );
361    polars_ensure!(
362        layout.is_pass_through() || !rows.has_nulls(),
363        InvalidOperation: "cannot row decode null rows"
364    );
365
366    let struct_arrow_dtype = ArrowDataType::Struct(
367        fields
368            .iter()
369            .map(|v| v.to_physical().to_arrow(CompatLevel::newest()))
370            .collect(),
371    );
372
373    let chunks = rows
374        .chunks()
375        .iter()
376        .map(|array| {
377            let decoded_arrays =
378                unsafe { polars_row::decode::decode_rows_from_array(array.as_ref(), layout) };
379            assert_eq!(decoded_arrays.len(), fields.len());
380
381            StructArray::new(
382                struct_arrow_dtype.clone(),
383                array.len(),
384                decoded_arrays,
385                None,
386            )
387            .to_boxed()
388        })
389        .collect::<Vec<_>>();
390
391    Ok(unsafe {
392        StructChunked::from_chunks_and_dtype(
393            rows.name().clone(),
394            chunks,
395            DataType::Struct(fields.to_vec()),
396        )
397    })
398}
polars_core/chunked_array/ops/row_encode.rs

polars_core/chunked_array/ops/
row_encode.rs
