Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-compute/src/gather/binary.rs
6939 views
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements. See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership. The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License. You may obtain a copy of the License at
8
//
9
// http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied. See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
use arrow::array::{Array, BinaryArray, PrimitiveArray};
19
use arrow::offset::Offset;
20
21
use super::Index;
22
use super::generic_binary::*;
23
24
/// `take` implementation for utf8 arrays
25
/// # Safety
26
/// The indices must be in-bounds.
27
pub unsafe fn take_unchecked<O: Offset, I: Index>(
28
values: &BinaryArray<O>,
29
indices: &PrimitiveArray<I>,
30
) -> BinaryArray<O> {
31
let dtype = values.dtype().clone();
32
let indices_has_validity = indices.null_count() > 0;
33
let values_has_validity = values.null_count() > 0;
34
35
let (offsets, values, validity) = match (values_has_validity, indices_has_validity) {
36
(false, false) => {
37
take_no_validity_unchecked::<O, I>(values.offsets(), values.values(), indices.values())
38
},
39
(true, false) => take_values_validity(values, indices.values()),
40
(false, true) => take_indices_validity(values.offsets(), values.values(), indices),
41
(true, true) => take_values_indices_validity(values, indices),
42
};
43
BinaryArray::<O>::new_unchecked(dtype, offsets, values, validity)
44
}
45
46