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

polars_core/chunked_array/ops/
row_encode.rs
