polars_core/chunked_array/ops/
row_encode.rs

1use std::borrow::Cow;
2
3use polars_arrow::compute::utils::combine_validities_and_many;
4use polars_row::{RowEncodingContext, RowEncodingOptions, RowsEncoded, convert_columns};
5use polars_utils::itertools::Itertools;
6use rayon::prelude::*;
7
8use crate::prelude::*;
9use crate::runtime::RAYON;
10use crate::utils::_split_offsets;
11
12fn encode_rows_vertical_par(
13    by: &[Column],
14    encode: impl Fn(&[Column]) -> PolarsResult<BinaryArray<i64>> + Sync,
15) -> PolarsResult<BinaryOffsetChunked> {
16    let n_threads = RAYON.current_num_threads();
17    let len = by[0].len();
18    let splits = _split_offsets(len, n_threads);
19
20    let chunks = splits.into_par_iter().map(|(offset, len)| {
21        let sliced = by
22            .iter()
23            .map(|s| s.slice(offset as i64, len))
24            .collect::<Vec<_>>();
25        encode(&sliced)
26    });
27    let chunks = RAYON.install(|| chunks.collect::<PolarsResult<Vec<_>>>());
28
29    Ok(BinaryOffsetChunked::from_chunk_iter(
30        PlSmallStr::EMPTY,
31        chunks?,
32    ))
33}
34
35pub fn encode_rows_vertical_par_unordered(by: &[Column]) -> PolarsResult<BinaryOffsetChunked> {
36    encode_rows_vertical_par(by, |sliced| {
37        Ok(_get_rows_encoded_unordered(sliced)?.into_array())
38    })
39}
40
41pub fn encode_rows_vertical_par_ordered(
42    by: &[Column],
43    descending: &[bool],
44    nulls_last: &[bool],
45    broadcast_nulls: bool,
46) -> PolarsResult<BinaryOffsetChunked> {
47    encode_rows_vertical_par(by, |sliced| {
48        _get_rows_encoded_arr(sliced, descending, nulls_last, broadcast_nulls)
49    })
50}
51
52// Almost the same but broadcast nulls to the row-encoded array.
53pub fn encode_rows_vertical_par_unordered_broadcast_nulls(
54    by: &[Column],
55) -> PolarsResult<BinaryOffsetChunked> {
56    encode_rows_vertical_par(by, |sliced| {
57        let rows = _get_rows_encoded_unordered(sliced)?;
58
59        let validities = sliced
60            .iter()
61            .flat_map(|s| {
62                let s = s.rechunk();
63                #[allow(clippy::unnecessary_to_owned)]
64                s.as_materialized_series()
65                    .chunks()
66                    .to_vec()
67                    .into_iter()
68                    .map(|arr| arr.validity().cloned())
69            })
70            .collect::<Vec<_>>();
71
72        let validity = combine_validities_and_many(&validities);
73        Ok(rows.into_array().with_validity_typed(validity))
74    })
75}
76
77/// Get the [`RowEncodingContext`] for a certain [`DataType`].
78///
79/// This should be given the logical type in order to communicate Polars datatype information down
80/// into the row encoding / decoding.
81pub fn get_row_encoding_context(dtype: &DataType) -> Option<RowEncodingContext> {
82    match dtype {
83        DataType::Boolean
84        | DataType::UInt8
85        | DataType::UInt16
86        | DataType::UInt32
87        | DataType::UInt64
88        | DataType::UInt128
89        | DataType::Int8
90        | DataType::Int16
91        | DataType::Int32
92        | DataType::Int64
93        | DataType::Int128
94        | DataType::Float16
95        | DataType::Float32
96        | DataType::Float64
97        | DataType::String
98        | DataType::Binary
99        | DataType::BinaryOffset
100        | DataType::Null
101        | DataType::Time
102        | DataType::Date
103        | DataType::Datetime(_, _)
104        | DataType::Duration(_) => None,
105
106        #[cfg(feature = "dtype-categorical")]
107        DataType::Categorical(_, mapping) | DataType::Enum(_, mapping) => {
108            use polars_row::RowEncodingCategoricalContext;
109
110            Some(RowEncodingContext::Categorical(
111                RowEncodingCategoricalContext {
112                    is_enum: matches!(dtype, DataType::Enum(_, _)),
113                    mapping: mapping.clone(),
114                },
115            ))
116        },
117
118        DataType::Unknown(_) => panic!("Unsupported in row encoding"),
119
120        #[cfg(feature = "object")]
121        DataType::Object(_) => panic!("Unsupported in row encoding"),
122
123        #[cfg(feature = "dtype-decimal")]
124        DataType::Decimal(precision, _) => Some(RowEncodingContext::Decimal(*precision)),
125
126        #[cfg(feature = "dtype-array")]
127        DataType::Array(dtype, _) => get_row_encoding_context(dtype),
128        DataType::List(dtype) => get_row_encoding_context(dtype),
129        #[cfg(feature = "dtype-struct")]
130        DataType::Struct(fs) => {
131            let mut ctxts = Vec::new();
132
133            for (i, f) in fs.iter().enumerate() {
134                if let Some(ctxt) = get_row_encoding_context(f.dtype()) {
135                    ctxts.reserve(fs.len());
136                    ctxts.extend(std::iter::repeat_n(None, i));
137                    ctxts.push(Some(ctxt));
138                    break;
139                }
140            }
141
142            if ctxts.is_empty() {
143                return None;
144            }
145
146            ctxts.extend(
147                fs[ctxts.len()..]
148                    .iter()
149                    .map(|f| get_row_encoding_context(f.dtype())),
150            );
151
152            Some(RowEncodingContext::Struct(ctxts))
153        },
154
155        #[cfg(feature = "dtype-map")]
156        DataType::Map(key, value) => {
157            let ctxts = vec![
158                get_row_encoding_context(key),
159                get_row_encoding_context(value),
160            ];
161
162            if ctxts.iter().all(Option::is_none) {
163                return None;
164            }
165
166            Some(RowEncodingContext::Struct(ctxts))
167        },
168
169        #[cfg(feature = "dtype-extension")]
170        DataType::Extension(_, storage) => get_row_encoding_context(storage),
171    }
172}
173
174pub fn encode_rows_unordered(by: &[Column]) -> PolarsResult<BinaryOffsetChunked> {
175    let rows = _get_rows_encoded_unordered(by)?;
176    Ok(BinaryOffsetChunked::with_chunk(
177        PlSmallStr::EMPTY,
178        rows.into_array(),
179    ))
180}
181
182pub fn _get_rows_encoded_unordered(by: &[Column]) -> PolarsResult<RowsEncoded> {
183    let mut cols = Vec::with_capacity(by.len());
184    let mut opts = Vec::with_capacity(by.len());
185    let mut ctxts = Vec::with_capacity(by.len());
186
187    // Since ZFS exists, we might not actually have any arrays and need to get the length from the
188    // columns.
189    let num_rows = by.first().map_or(0, |c| c.len());
190
191    for by in by {
192        debug_assert_eq!(by.len(), num_rows);
193
194        let by = by
195            .trim_lists_to_normalized_offsets()
196            .map_or(Cow::Borrowed(by), Cow::Owned);
197        let by = by.propagate_nulls().map_or(by, Cow::Owned);
198        let by = by.as_materialized_series();
199        let arr = by.to_physical_repr().rechunk().chunks()[0].to_boxed();
200        let opt = RowEncodingOptions::new_unsorted();
201        let ctxt = get_row_encoding_context(by.dtype());
202
203        cols.push(arr);
204        opts.push(opt);
205        ctxts.push(ctxt);
206    }
207    Ok(convert_columns(num_rows, &cols, &opts, &ctxts))
208}
209
210pub fn _get_rows_encoded(
211    by: &[Column],
212    descending: &[bool],
213    nulls_last: &[bool],
214) -> PolarsResult<RowsEncoded> {
215    debug_assert_eq!(by.len(), descending.len());
216    debug_assert_eq!(by.len(), nulls_last.len());
217
218    let mut cols = Vec::with_capacity(by.len());
219    let mut opts = Vec::with_capacity(by.len());
220    let mut ctxts = Vec::with_capacity(by.len());
221
222    // Since ZFS exists, we might not actually have any arrays and need to get the length from the
223    // columns.
224    let num_rows = by.first().map_or(0, |c| c.len());
225
226    for ((by, desc), null_last) in by.iter().zip(descending).zip(nulls_last) {
227        debug_assert_eq!(by.len(), num_rows);
228
229        let by = by
230            .trim_lists_to_normalized_offsets()
231            .map_or(Cow::Borrowed(by), Cow::Owned);
232        let by = by.propagate_nulls().map_or(by, Cow::Owned);
233        let by = by.as_materialized_series();
234        let arr = by.to_physical_repr().rechunk().chunks()[0].to_boxed();
235        let opt = RowEncodingOptions::new_sorted(*desc, *null_last);
236        let ctxt = get_row_encoding_context(by.dtype());
237
238        cols.push(arr);
239        opts.push(opt);
240        ctxts.push(ctxt);
241    }
242    Ok(convert_columns(num_rows, &cols, &opts, &ctxts))
243}
244
245pub fn _get_rows_encoded_ca(
246    name: PlSmallStr,
247    by: &[Column],
248    descending: &[bool],
249    nulls_last: &[bool],
250    broadcast_nulls: bool,
251) -> PolarsResult<BinaryOffsetChunked> {
252    let mut rows_arr = _get_rows_encoded(by, descending, nulls_last)?.into_array();
253    if broadcast_nulls {
254        let validities = by
255            .iter()
256            .map(|c| c.as_materialized_series().rechunk_validity())
257            .collect_vec();
258        let combined = combine_validities_and_many(&validities);
259        rows_arr.set_validity(combined);
260    }
261    Ok(BinaryOffsetChunked::with_chunk(name, rows_arr))
262}
263
264pub fn _get_rows_encoded_arr(
265    by: &[Column],
266    descending: &[bool],
267    nulls_last: &[bool],
268    broadcast_nulls: bool,
269) -> PolarsResult<BinaryArray<i64>> {
270    let mut rows_arr = _get_rows_encoded(by, descending, nulls_last)?.into_array();
271    if broadcast_nulls {
272        let validities = by
273            .iter()
274            .map(|c| c.as_materialized_series().rechunk_validity())
275            .collect_vec();
276        let combined = combine_validities_and_many(&validities);
277        rows_arr.set_validity(combined);
278    }
279    Ok(rows_arr)
280}
281
282pub fn _get_rows_encoded_ca_unordered(
283    name: PlSmallStr,
284    by: &[Column],
285) -> PolarsResult<BinaryOffsetChunked> {
286    _get_rows_encoded_unordered(by)
287        .map(|rows| BinaryOffsetChunked::with_chunk(name, rows.into_array()))
288}
289
290#[cfg(feature = "dtype-struct")]
291pub fn row_encoding_decode(
292    ca: &BinaryOffsetChunked,
293    fields: &[Field],
294    opts: &[RowEncodingOptions],
295) -> PolarsResult<StructChunked> {
296    let (ctxts, dtypes) = fields
297        .iter()
298        .map(|f| {
299            (
300                get_row_encoding_context(f.dtype()),
301                f.dtype().to_physical().to_arrow(CompatLevel::newest()),
302            )
303        })
304        .collect::<(Vec<_>, Vec<_>)>();
305
306    let struct_arrow_dtype = ArrowDataType::Struct(
307        fields
308            .iter()
309            .map(|v| v.to_physical().to_arrow(CompatLevel::newest()))
310            .collect(),
311    );
312
313    let mut rows = Vec::new();
314    let chunks = ca
315        .downcast_iter()
316        .map(|array| {
317            let decoded_arrays = unsafe {
318                polars_row::decode::decode_rows_from_binary(array, opts, &ctxts, &dtypes, &mut rows)
319            };
320            assert_eq!(decoded_arrays.len(), fields.len());
321
322            StructArray::new(
323                struct_arrow_dtype.clone(),
324                array.len(),
325                decoded_arrays,
326                None,
327            )
328            .to_boxed()
329        })
330        .collect::<Vec<_>>();
331
332    Ok(unsafe {
333        StructChunked::from_chunks_and_dtype(
334            ca.name().clone(),
335            chunks,
336            DataType::Struct(fields.to_vec()),
337        )
338    })
339}
polars_core/chunked_array/ops/row_encode.rs

polars_core/chunked_array/ops/
row_encode.rs
