datafusion_spark/function/math/
bin.rs

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
18use arrow::array::{Array, ArrayRef, AsArray, StringBuilder};
19use arrow::datatypes::{DataType, Field, FieldRef, Int64Type};
20use datafusion_common::types::{NativeType, logical_int64};
21use datafusion_common::utils::take_function_args;
22use datafusion_common::{Result, internal_err};
23use datafusion_expr::{
24    Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature,
25    TypeSignatureClass, Volatility,
26};
27use datafusion_functions::utils::make_scalar_function;
28use std::sync::Arc;
29
30/// Spark-compatible `bin` expression
31/// <https://spark.apache.org/docs/latest/api/sql/index.html#bin>
32#[derive(Debug, PartialEq, Eq, Hash)]
33pub struct SparkBin {
34    signature: Signature,
35}
36
37impl Default for SparkBin {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl SparkBin {
44    pub fn new() -> Self {
45        Self {
46            signature: Signature::one_of(
47                vec![TypeSignature::Coercible(vec![Coercion::new_implicit(
48                    TypeSignatureClass::Native(logical_int64()),
49                    vec![TypeSignatureClass::Numeric],
50                    NativeType::Int64,
51                )])],
52                Volatility::Immutable,
53            ),
54        }
55    }
56}
57
58impl ScalarUDFImpl for SparkBin {
59    fn name(&self) -> &str {
60        "bin"
61    }
62
63    fn signature(&self) -> &Signature {
64        &self.signature
65    }
66
67    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
68        internal_err!("return_field_from_args should be used instead")
69    }
70
71    fn return_field_from_args(
72        &self,
73        args: datafusion_expr::ReturnFieldArgs,
74    ) -> Result<FieldRef> {
75        Ok(Arc::new(Field::new(
76            self.name(),
77            DataType::Utf8,
78            args.arg_fields[0].is_nullable(),
79        )))
80    }
81
82    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
83        make_scalar_function(spark_bin_inner, vec![])(&args.args)
84    }
85}
86
87fn spark_bin_inner(arg: &[ArrayRef]) -> Result<ArrayRef> {
88    let [array] = take_function_args("bin", arg)?;
89    match &array.data_type() {
90        DataType::Int64 => {
91            let array = array.as_primitive::<Int64Type>();
92            let len = array.len();
93            // Most values are small, so 8 digits per row is a reasonable estimate;
94            // the buffer grows on its own for wider ones.
95            let mut builder = StringBuilder::with_capacity(len, len * 8);
96            // Digits are rendered into this stack buffer, so no row allocates.
97            let mut digits = [0u8; MAX_BIN_DIGITS];
98            for value in array.iter() {
99                match value {
100                    Some(value) => builder.append_value(spark_bin(value, &mut digits)),
101                    None => builder.append_null(),
102                }
103            }
104            Ok(Arc::new(builder.finish()))
105        }
106        data_type => {
107            internal_err!("bin does not support: {data_type}")
108        }
109    }
110}
111
112/// An `i64` renders as at most 64 binary digits.
113const MAX_BIN_DIGITS: usize = 64;
114
115/// Renders `value` as binary, right-aligned in `digits`, and returns the digits written.
116///
117/// Negative values render as their two's-complement bit pattern, matching `{:b}`.
118fn spark_bin(value: i64, digits: &mut [u8; MAX_BIN_DIGITS]) -> &str {
119    let mut pos = MAX_BIN_DIGITS;
120    let mut remaining = value as u64;
121    // `while` alone would produce an empty string for zero.
122    loop {
123        pos -= 1;
124        digits[pos] = b'0' + (remaining & 1) as u8;
125        remaining >>= 1;
126        if remaining == 0 {
127            break;
128        }
129    }
130    // SAFETY: every byte written above is an ASCII '0' or '1'.
131    unsafe { std::str::from_utf8_unchecked(&digits[pos..]) }
132}
datafusion_spark/function/math/bin.rs

datafusion_spark/function/math/
bin.rs
