datafusion_functions/encoding/
inner.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
18//! Encoding expressions
19
20use arrow::{
21    array::{
22        Array, ArrayRef, AsArray, BinaryArrayType, GenericBinaryArray,
23        GenericStringArray, OffsetSizeTrait,
24    },
25    datatypes::DataType,
26};
27use arrow_buffer::{Buffer, OffsetBuffer};
28use base64::{
29    Engine as _,
30    engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig},
31};
32use datafusion_common::{
33    DataFusionError, Result, ScalarValue, exec_datafusion_err, exec_err, internal_err,
34    not_impl_err, plan_err,
35    types::{NativeType, logical_string},
36    utils::{
37        hex::{HexCase, encode_bytes as encode_hex, encode_bytes_to_slice},
38        take_function_args,
39    },
40};
41use datafusion_expr::{
42    Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
43    TypeSignatureClass, Volatility,
44};
45use datafusion_macros::user_doc;
46use std::fmt;
47use std::sync::Arc;
48
49// Allow padding characters, but don't require them, and don't generate them.
50const BASE64_ENGINE: GeneralPurpose = GeneralPurpose::new(
51    &base64::alphabet::STANDARD,
52    GeneralPurposeConfig::new()
53        .with_encode_padding(false)
54        .with_decode_padding_mode(DecodePaddingMode::Indifferent),
55);
56
57// Generate padding characters when encoding
58const BASE64_ENGINE_PADDED: GeneralPurpose = GeneralPurpose::new(
59    &base64::alphabet::STANDARD,
60    GeneralPurposeConfig::new().with_encode_padding(true),
61);
62
63#[user_doc(
64    doc_section(label = "Binary String Functions"),
65    description = "Encode binary data into a textual representation.",
66    syntax_example = "encode(expression, format)",
67    argument(
68        name = "expression",
69        description = "Expression containing string or binary data"
70    ),
71    argument(
72        name = "format",
73        description = "Supported formats are: `base64`, `base64pad`, `hex`"
74    ),
75    related_udf(name = "decode")
76)]
77#[derive(Debug, PartialEq, Eq, Hash)]
78pub struct EncodeFunc {
79    signature: Signature,
80}
81
82impl Default for EncodeFunc {
83    fn default() -> Self {
84        Self::new()
85    }
86}
87
88impl EncodeFunc {
89    pub fn new() -> Self {
90        Self {
91            signature: Signature::coercible(
92                vec![
93                    Coercion::new_implicit(
94                        TypeSignatureClass::Binary,
95                        vec![TypeSignatureClass::Native(logical_string())],
96                        NativeType::Binary,
97                    ),
98                    Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
99                ],
100                Volatility::Immutable,
101            ),
102        }
103    }
104}
105
106impl ScalarUDFImpl for EncodeFunc {
107    fn name(&self) -> &str {
108        "encode"
109    }
110
111    fn signature(&self) -> &Signature {
112        &self.signature
113    }
114
115    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
116        match &arg_types[0] {
117            DataType::LargeBinary => Ok(DataType::LargeUtf8),
118            _ => Ok(DataType::Utf8),
119        }
120    }
121
122    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
123        let [expression, encoding] = take_function_args("encode", &args.args)?;
124        let encoding = Encoding::try_from(encoding)?;
125        match expression {
126            _ if expression.data_type().is_null() => {
127                Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None)))
128            }
129            ColumnarValue::Array(array) => encode_array(array, encoding),
130            ColumnarValue::Scalar(scalar) => encode_scalar(scalar, encoding),
131        }
132    }
133
134    fn documentation(&self) -> Option<&Documentation> {
135        self.doc()
136    }
137}
138
139#[user_doc(
140    doc_section(label = "Binary String Functions"),
141    description = "Decode binary data from textual representation in string.",
142    syntax_example = "decode(expression, format)",
143    argument(
144        name = "expression",
145        description = "Expression containing encoded string data"
146    ),
147    argument(name = "format", description = "Same arguments as [encode](#encode)"),
148    related_udf(name = "encode")
149)]
150#[derive(Debug, PartialEq, Eq, Hash)]
151pub struct DecodeFunc {
152    signature: Signature,
153}
154
155impl Default for DecodeFunc {
156    fn default() -> Self {
157        Self::new()
158    }
159}
160
161impl DecodeFunc {
162    pub fn new() -> Self {
163        Self {
164            signature: Signature::coercible(
165                vec![
166                    Coercion::new_implicit(
167                        TypeSignatureClass::Binary,
168                        vec![TypeSignatureClass::Native(logical_string())],
169                        NativeType::Binary,
170                    ),
171                    Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
172                ],
173                Volatility::Immutable,
174            ),
175        }
176    }
177}
178
179impl ScalarUDFImpl for DecodeFunc {
180    fn name(&self) -> &str {
181        "decode"
182    }
183
184    fn signature(&self) -> &Signature {
185        &self.signature
186    }
187
188    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
189        match &arg_types[0] {
190            DataType::LargeBinary => Ok(DataType::LargeBinary),
191            _ => Ok(DataType::Binary),
192        }
193    }
194
195    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
196        let [expression, encoding] = take_function_args("decode", &args.args)?;
197        let encoding = Encoding::try_from(encoding)?;
198        match expression {
199            _ if expression.data_type().is_null() => {
200                Ok(ColumnarValue::Scalar(ScalarValue::Binary(None)))
201            }
202            ColumnarValue::Array(array) => decode_array(array, encoding),
203            ColumnarValue::Scalar(scalar) => decode_scalar(scalar, encoding),
204        }
205    }
206
207    fn documentation(&self) -> Option<&Documentation> {
208        self.doc()
209    }
210}
211
212fn encode_scalar(value: &ScalarValue, encoding: Encoding) -> Result<ColumnarValue> {
213    match value {
214        ScalarValue::Binary(maybe_bytes)
215        | ScalarValue::BinaryView(maybe_bytes)
216        | ScalarValue::FixedSizeBinary(_, maybe_bytes) => {
217            Ok(ColumnarValue::Scalar(ScalarValue::Utf8(
218                maybe_bytes
219                    .as_ref()
220                    .map(|bytes| encoding.encode_bytes(bytes)),
221            )))
222        }
223        ScalarValue::LargeBinary(maybe_bytes) => {
224            Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(
225                maybe_bytes
226                    .as_ref()
227                    .map(|bytes| encoding.encode_bytes(bytes)),
228            )))
229        }
230        v => internal_err!("Unexpected value for encode: {v}"),
231    }
232}
233
234fn encode_array(array: &ArrayRef, encoding: Encoding) -> Result<ColumnarValue> {
235    let array = match array.data_type() {
236        DataType::Binary => encoding.encode_array::<_, i32>(&array.as_binary::<i32>()),
237        DataType::BinaryView => encoding.encode_array::<_, i32>(&array.as_binary_view()),
238        DataType::LargeBinary => {
239            encoding.encode_array::<_, i64>(&array.as_binary::<i64>())
240        }
241        DataType::FixedSizeBinary(_) => {
242            encoding.encode_array::<_, i32>(&array.as_fixed_size_binary())
243        }
244        dt => {
245            internal_err!("Unexpected data type for encode: {dt}")
246        }
247    };
248    array.map(ColumnarValue::Array)
249}
250
251fn decode_scalar(value: &ScalarValue, encoding: Encoding) -> Result<ColumnarValue> {
252    match value {
253        ScalarValue::Binary(maybe_bytes)
254        | ScalarValue::BinaryView(maybe_bytes)
255        | ScalarValue::FixedSizeBinary(_, maybe_bytes) => {
256            Ok(ColumnarValue::Scalar(ScalarValue::Binary(
257                maybe_bytes
258                    .as_ref()
259                    .map(|x| encoding.decode_bytes(x))
260                    .transpose()?,
261            )))
262        }
263        ScalarValue::LargeBinary(maybe_bytes) => {
264            Ok(ColumnarValue::Scalar(ScalarValue::LargeBinary(
265                maybe_bytes
266                    .as_ref()
267                    .map(|x| encoding.decode_bytes(x))
268                    .transpose()?,
269            )))
270        }
271        v => internal_err!("Unexpected value for decode: {v}"),
272    }
273}
274
275/// Estimate how many bytes are actually represented by the array; in case the
276/// the array slices it's internal buffer, this returns the byte size of that slice
277/// but not the byte size of the entire buffer.
278///
279/// This is an estimation only as it can estimate higher if null slots are non-zero
280/// sized.
281fn estimate_byte_data_size<O: OffsetSizeTrait>(array: &GenericBinaryArray<O>) -> usize {
282    let offsets = array.value_offsets();
283    // Unwraps are safe as should always have 1 element in offset buffer
284    let start = *offsets.first().unwrap();
285    let end = *offsets.last().unwrap();
286    let data_size = end - start;
287    data_size.as_usize()
288}
289
290fn decode_array(array: &ArrayRef, encoding: Encoding) -> Result<ColumnarValue> {
291    let array = match array.data_type() {
292        DataType::Binary => {
293            let array = array.as_binary::<i32>();
294            encoding.decode_array::<_, i32>(&array, estimate_byte_data_size(array))
295        }
296        DataType::BinaryView => {
297            let array = array.as_binary_view();
298            encoding.decode_array::<_, i32>(
299                &array,
300                array.lengths().map(|l| l as usize).sum::<usize>(),
301            )
302        }
303        DataType::LargeBinary => {
304            let array = array.as_binary::<i64>();
305            encoding.decode_array::<_, i64>(&array, estimate_byte_data_size(array))
306        }
307        DataType::FixedSizeBinary(size) => {
308            let array = array.as_fixed_size_binary();
309            // TODO: could we be more conservative by accounting for nulls?
310            let estimate = array.len().saturating_mul(*size as usize);
311            encoding.decode_array::<_, i32>(&array, estimate)
312        }
313        dt => {
314            internal_err!("Unexpected data type for decode: {dt}")
315        }
316    };
317    array.map(ColumnarValue::Array)
318}
319
320#[derive(Debug, Copy, Clone)]
321enum Encoding {
322    Base64,
323    Base64Padded,
324    Hex,
325}
326
327impl fmt::Display for Encoding {
328    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
329        let name = match self {
330            Self::Base64 => "base64",
331            Self::Base64Padded => "base64pad",
332            Self::Hex => "hex",
333        };
334        write!(f, "{name}")
335    }
336}
337
338impl TryFrom<&ColumnarValue> for Encoding {
339    type Error = DataFusionError;
340
341    fn try_from(encoding: &ColumnarValue) -> Result<Self> {
342        let encoding = match encoding {
343            ColumnarValue::Scalar(encoding) => match encoding.try_as_str().flatten() {
344                Some(encoding) => encoding,
345                _ => return exec_err!("Encoding must be a non-null string"),
346            },
347            ColumnarValue::Array(_) => {
348                return not_impl_err!(
349                    "Encoding must be a scalar; array specified encoding is not yet supported"
350                );
351            }
352        };
353        match encoding {
354            "base64" => Ok(Self::Base64),
355            "base64pad" => Ok(Self::Base64Padded),
356            "hex" => Ok(Self::Hex),
357            _ => {
358                let options = [Self::Base64, Self::Base64Padded, Self::Hex]
359                    .iter()
360                    .map(|i| i.to_string())
361                    .collect::<Vec<_>>()
362                    .join(", ");
363                plan_err!(
364                    "There is no built-in encoding named '{encoding}', currently supported encodings are: {options}"
365                )
366            }
367        }
368    }
369}
370
371impl Encoding {
372    fn encode_bytes(self, value: &[u8]) -> String {
373        match self {
374            Self::Base64 => BASE64_ENGINE.encode(value),
375            Self::Base64Padded => BASE64_ENGINE_PADDED.encode(value),
376            Self::Hex => encode_hex(value, HexCase::Lower),
377        }
378    }
379
380    fn decode_bytes(self, value: &[u8]) -> Result<Vec<u8>> {
381        match self {
382            Self::Base64 | Self::Base64Padded => {
383                BASE64_ENGINE.decode(value).map_err(|e| {
384                    exec_datafusion_err!("Failed to decode value using {self}: {e}")
385                })
386            }
387            Self::Hex => hex::decode(value).map_err(|e| {
388                exec_datafusion_err!("Failed to decode value using hex: {e}")
389            }),
390        }
391    }
392
393    // OutputOffset important to ensure Large types output Large arrays
394    fn encode_array<'a, InputBinaryArray, OutputOffset>(
395        self,
396        array: &InputBinaryArray,
397    ) -> Result<ArrayRef>
398    where
399        InputBinaryArray: BinaryArrayType<'a>,
400        OutputOffset: OffsetSizeTrait,
401    {
402        match self {
403            Self::Base64 => {
404                let array: GenericStringArray<OutputOffset> = array
405                    .iter()
406                    .map(|x| x.map(|x| BASE64_ENGINE.encode(x)))
407                    .collect();
408                Ok(Arc::new(array))
409            }
410            Self::Base64Padded => {
411                let array: GenericStringArray<OutputOffset> = array
412                    .iter()
413                    .map(|x| x.map(|x| BASE64_ENGINE_PADDED.encode(x)))
414                    .collect();
415                Ok(Arc::new(array))
416            }
417            Self::Hex => hex_encode_array::<_, OutputOffset>(array),
418        }
419    }
420
421    // OutputOffset important to ensure Large types output Large arrays
422    fn decode_array<'a, InputBinaryArray, OutputOffset>(
423        self,
424        value: &InputBinaryArray,
425        approx_data_size: usize,
426    ) -> Result<ArrayRef>
427    where
428        InputBinaryArray: BinaryArrayType<'a>,
429        OutputOffset: OffsetSizeTrait,
430    {
431        fn hex_decode(input: &[u8], buf: &mut [u8]) -> Result<usize> {
432            // only write input / 2 bytes to buf
433            let out_len = input.len() / 2;
434            let buf = &mut buf[..out_len];
435            hex::decode_to_slice(input, buf)
436                .map_err(|e| exec_datafusion_err!("Failed to decode from hex: {e}"))?;
437            Ok(out_len)
438        }
439
440        fn base64_decode(input: &[u8], buf: &mut [u8]) -> Result<usize> {
441            BASE64_ENGINE
442                .decode_slice(input, buf)
443                .map_err(|e| exec_datafusion_err!("Failed to decode from base64: {e}"))
444        }
445
446        match self {
447            Self::Base64 | Self::Base64Padded => {
448                let upper_bound = base64::decoded_len_estimate(approx_data_size);
449                delegated_decode::<_, _, OutputOffset>(base64_decode, value, upper_bound)
450            }
451            Self::Hex => {
452                // Calculate the upper bound for decoded byte size
453                // For hex encoding, each pair of hex characters (2 bytes) represents 1 byte when decoded
454                // So the upper bound is half the length of the input values.
455                let upper_bound = approx_data_size / 2;
456                delegated_decode::<_, _, OutputOffset>(hex_decode, value, upper_bound)
457            }
458        }
459    }
460}
461
462/// Hex-encode a binary array into a string array, writing the lowercase hex
463/// digits directly into a single pre-sized value buffer. Each input byte maps
464/// to exactly two hex characters, so the output size is known up front and no
465/// per-element `String` is allocated.
466fn hex_encode_array<'a, InputBinaryArray, OutputOffset>(
467    array: &InputBinaryArray,
468) -> Result<ArrayRef>
469where
470    InputBinaryArray: BinaryArrayType<'a>,
471    OutputOffset: OffsetSizeTrait,
472{
473    let total_input_bytes: usize = array.iter().flatten().map(|v| v.len()).sum();
474
475    let mut values = vec![0u8; total_input_bytes * 2];
476    let mut offsets = Vec::<OutputOffset>::with_capacity(array.len() + 1);
477    offsets.push(OutputOffset::zero());
478
479    let mut pos = 0usize;
480    for v in array.iter() {
481        if let Some(v) = v {
482            let out_len = v.len() * 2;
483            encode_bytes_to_slice(v, HexCase::Lower, &mut values[pos..pos + out_len])?;
484            pos += out_len;
485        }
486        offsets.push(OutputOffset::usize_as(pos));
487    }
488
489    let array = GenericStringArray::<OutputOffset>::try_new(
490        OffsetBuffer::new(offsets.into()),
491        Buffer::from_vec(values),
492        array.nulls().cloned(),
493    )?;
494    Ok(Arc::new(array))
495}
496
497fn delegated_decode<'a, DecodeFunction, InputBinaryArray, OutputOffset>(
498    decode: DecodeFunction,
499    input: &InputBinaryArray,
500    conservative_upper_bound_size: usize,
501) -> Result<ArrayRef>
502where
503    DecodeFunction: Fn(&[u8], &mut [u8]) -> Result<usize>,
504    InputBinaryArray: BinaryArrayType<'a>,
505    OutputOffset: OffsetSizeTrait,
506{
507    let mut values = vec![0; conservative_upper_bound_size];
508    let mut offsets = Vec::<OutputOffset>::with_capacity(input.len() + 1);
509    offsets.push(OutputOffset::zero());
510    let mut total_bytes_decoded = 0;
511    for v in input.iter() {
512        if let Some(v) = v {
513            let cursor = &mut values[total_bytes_decoded..];
514            let decoded = decode(v, cursor)?;
515            total_bytes_decoded += decoded;
516        }
517        offsets.push(OutputOffset::usize_as(total_bytes_decoded));
518    }
519    // We reserved an upper bound size for the values buffer, but we only use the actual size
520    values.truncate(total_bytes_decoded);
521    let binary_array = GenericBinaryArray::<OutputOffset>::try_new(
522        OffsetBuffer::new(offsets.into()),
523        Buffer::from_vec(values),
524        input.nulls().cloned(),
525    )?;
526    Ok(Arc::new(binary_array))
527}
528
529#[cfg(test)]
530mod tests {
531    use arrow::array::{ArrayBuilder, BinaryArray, BinaryViewBuilder};
532    use arrow_buffer::OffsetBuffer;
533
534    use super::*;
535
536    #[test]
537    fn test_estimate_byte_data_size() {
538        // Offsets starting at 0, but don't count entire data buffer size
539        let array = BinaryArray::new(
540            OffsetBuffer::new(vec![0, 5, 10, 15].into()),
541            vec![0; 100].into(),
542            None,
543        );
544        let size = estimate_byte_data_size(&array);
545        assert_eq!(size, 15);
546
547        // Offsets starting at 0, but don't count entire data buffer size
548        let array = BinaryArray::new(
549            OffsetBuffer::new(vec![50, 51, 51, 60, 80, 81].into()),
550            vec![0; 100].into(),
551            Some(vec![true, false, false, true, true].into()),
552        );
553        let size = estimate_byte_data_size(&array);
554        assert_eq!(size, 31);
555    }
556
557    #[test]
558    fn test_estimate_view_size() {
559        let mut builder = BinaryViewBuilder::new().with_deduplicate_strings();
560        for _ in 0..1000 {
561            builder.append_value([65u8; 64]);
562        }
563        let arr = ArrayBuilder::finish(&mut builder);
564        decode_array(&arr, Encoding::Base64).unwrap();
565    }
566}
datafusion_functions/encoding/inner.rs

datafusion_functions/encoding/
inner.rs
