datafusion_spark/function/hash/
crc32.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 std::sync::Arc;
19
20use arrow::array::{ArrayRef, Int64Array};
21use arrow::datatypes::{DataType, Field, FieldRef};
22use crc32fast::Hasher;
23use datafusion_common::cast::{
24    as_binary_array, as_binary_view_array, as_fixed_size_binary_array,
25    as_large_binary_array,
26};
27use datafusion_common::types::{NativeType, logical_string};
28use datafusion_common::utils::take_function_args;
29use datafusion_common::{Result, internal_err};
30use datafusion_expr::{
31    Coercion, ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl,
32    Signature, TypeSignatureClass, Volatility,
33};
34use datafusion_functions::utils::make_scalar_function;
35
36/// <https://spark.apache.org/docs/latest/api/sql/index.html#crc32>
37#[derive(Debug, PartialEq, Eq, Hash)]
38pub struct SparkCrc32 {
39    signature: Signature,
40}
41
42impl Default for SparkCrc32 {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl SparkCrc32 {
49    pub fn new() -> Self {
50        Self {
51            signature: Signature::coercible(
52                vec![Coercion::new_implicit(
53                    TypeSignatureClass::Binary,
54                    vec![TypeSignatureClass::Native(logical_string())],
55                    NativeType::Binary,
56                )],
57                Volatility::Immutable,
58            ),
59        }
60    }
61}
62
63impl ScalarUDFImpl for SparkCrc32 {
64    fn name(&self) -> &str {
65        "crc32"
66    }
67
68    fn signature(&self) -> &Signature {
69        &self.signature
70    }
71
72    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
73        internal_err!("return_field_from_args should be used instead")
74    }
75
76    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
77        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
78        Ok(Arc::new(Field::new(self.name(), DataType::Int64, nullable)))
79    }
80
81    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
82        make_scalar_function(spark_crc32, vec![])(&args.args)
83    }
84}
85
86fn spark_crc32_digest(value: &[u8]) -> i64 {
87    let mut hasher = Hasher::new();
88    hasher.update(value);
89    hasher.finalize() as i64
90}
91
92fn spark_crc32_impl<'a>(input: impl Iterator<Item = Option<&'a [u8]>>) -> ArrayRef {
93    let result = input
94        .map(|value| value.map(spark_crc32_digest))
95        .collect::<Int64Array>();
96    Arc::new(result)
97}
98
99fn spark_crc32(args: &[ArrayRef]) -> Result<ArrayRef> {
100    let [input] = take_function_args("crc32", args)?;
101
102    match input.data_type() {
103        DataType::Null => Ok(Arc::new(Int64Array::new_null(input.len()))),
104        DataType::Binary => {
105            let input = as_binary_array(input)?;
106            Ok(spark_crc32_impl(input.iter()))
107        }
108        DataType::LargeBinary => {
109            let input = as_large_binary_array(input)?;
110            Ok(spark_crc32_impl(input.iter()))
111        }
112        DataType::BinaryView => {
113            let input = as_binary_view_array(input)?;
114            Ok(spark_crc32_impl(input.iter()))
115        }
116        DataType::FixedSizeBinary(_) => {
117            let input = as_fixed_size_binary_array(input)?;
118            Ok(spark_crc32_impl(input.iter()))
119        }
120        dt => {
121            internal_err!("Unsupported data type for crc32: {dt}")
122        }
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    fn test_crc32_nullability() -> Result<()> {
132        let crc32_func = SparkCrc32::new();
133
134        // non-nullable field should produce non-nullable output
135        let field_not_null = Arc::new(Field::new("data", DataType::Binary, false));
136        let result = crc32_func.return_field_from_args(ReturnFieldArgs {
137            arg_fields: std::slice::from_ref(&field_not_null),
138            scalar_arguments: &[None],
139        })?;
140        assert!(!result.is_nullable());
141        assert_eq!(result.data_type(), &DataType::Int64);
142
143        // nullable field should produce nullable output
144        let field_nullable = Arc::new(Field::new("data", DataType::Binary, true));
145        let result = crc32_func.return_field_from_args(ReturnFieldArgs {
146            arg_fields: &[field_nullable],
147            scalar_arguments: &[None],
148        })?;
149        assert!(result.is_nullable());
150        assert_eq!(result.data_type(), &DataType::Int64);
151
152        Ok(())
153    }
154}
datafusion_spark/function/hash/crc32.rs

datafusion_spark/function/hash/
crc32.rs
