partition/
function.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! `iceberg_hash`, the hash allowed on the column side of a partition rule.
16//!
17//! The hash decides which region a row is written to and is persisted in
18//! partition metadata, so its result for a given input must never change across
19//! versions. A different algorithm needs a new type.
20//!
21//! Rule simplification and coverage checks treat the hash as an independent
22//! dimension, which holds because it is defined for every value of the accepted
23//! column types.
24
25use std::fmt::{Display, Formatter};
26use std::sync::Arc;
27
28use datafusion_common::arrow::array::{Array, ArrayRef, AsArray, Int32Array};
29use datafusion_common::arrow::compute::take;
30use datafusion_common::arrow::datatypes::{
31    DataType, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type,
32    UInt64Type,
33};
34use datafusion_common::{DataFusionError, ScalarValue};
35use datafusion_expr::{
36    ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
37};
38use datatypes::prelude::ConcreteDataType;
39use datatypes::value::Value;
40use serde::{Deserialize, Serialize};
41
42use crate::error::{self, Result};
43
44/// `iceberg_hash(column)`: the 32-bit Murmur3 hash defined by the Iceberg table
45/// spec (appendix B) over a string or integer column, as an `INT`.
46///
47/// Integers of every width hash as Iceberg longs. `UINT64` values hash their bits
48/// as a long, so values up to `i64::MAX` match Iceberg. `NULL` hashes to `NULL`.
49/// The field is the hashed column.
50#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
51pub struct IcebergHash(String);
52
53impl IcebergHash {
54    pub const NAME: &str = "iceberg_hash";
55
56    pub fn new(column: impl Into<String>) -> Self {
57        Self(column.into())
58    }
59
60    pub fn column(&self) -> &str {
61        &self.0
62    }
63
64    pub fn accepts(data_type: &ConcreteDataType) -> bool {
65        data_type.is_string() || data_type.is_unsigned() || is_signed_integer(data_type)
66    }
67
68    /// The type of the result, which is also the type of the literals it is compared with.
69    pub fn output_type() -> ConcreteDataType {
70        ConcreteDataType::int32_datatype()
71    }
72
73    pub fn evaluate_value(value: &Value) -> Result<Value> {
74        let hash = match value {
75            Value::Null => return Ok(Value::Null),
76            Value::String(v) => murmur3_32(v.as_utf8().as_bytes()),
77            Value::Int8(v) => hash_long(*v as i64),
78            Value::Int16(v) => hash_long(*v as i64),
79            Value::Int32(v) => hash_long(*v as i64),
80            Value::Int64(v) => hash_long(*v),
81            Value::UInt8(v) => hash_long(*v as i64),
82            Value::UInt16(v) => hash_long(*v as i64),
83            Value::UInt32(v) => hash_long(*v as i64),
84            Value::UInt64(v) => hash_long(*v as i64),
85            other => {
86                return error::UnexpectedSnafu {
87                    err_msg: format!("Unsupported value for {}: {other:?}", Self::NAME),
88                }
89                .fail();
90            }
91        };
92        Ok(Value::Int32(hash))
93    }
94
95    pub fn evaluate_array(array: &dyn Array) -> datafusion_common::Result<ArrayRef> {
96        iceberg_hash_array(array)
97    }
98
99    pub fn to_logical_expr(&self) -> Expr {
100        let column = Expr::Column(datafusion_common::Column::from_name(self.column()));
101        ScalarUDF::new_from_impl(IcebergHashUdf::default()).call(vec![column])
102    }
103}
104
105impl Display for IcebergHash {
106    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
107        write!(f, "{}({})", Self::NAME, self.column())
108    }
109}
110
111fn is_signed_integer(data_type: &ConcreteDataType) -> bool {
112    matches!(
113        data_type,
114        ConcreteDataType::Int8(_)
115            | ConcreteDataType::Int16(_)
116            | ConcreteDataType::Int32(_)
117            | ConcreteDataType::Int64(_)
118    )
119}
120
121fn iceberg_hash_array(array: &dyn Array) -> datafusion_common::Result<ArrayRef> {
122    let hashes: Int32Array = match array.data_type() {
123        DataType::Utf8 => array
124            .as_string::<i32>()
125            .iter()
126            .map(|v| v.map(|v| murmur3_32(v.as_bytes())))
127            .collect(),
128        DataType::LargeUtf8 => array
129            .as_string::<i64>()
130            .iter()
131            .map(|v| v.map(|v| murmur3_32(v.as_bytes())))
132            .collect(),
133        DataType::Utf8View => array
134            .as_string_view()
135            .iter()
136            .map(|v| v.map(|v| murmur3_32(v.as_bytes())))
137            .collect(),
138        DataType::Int8 => array
139            .as_primitive::<Int8Type>()
140            .unary(|v| hash_long(v as i64)),
141        DataType::Int16 => array
142            .as_primitive::<Int16Type>()
143            .unary(|v| hash_long(v as i64)),
144        DataType::Int32 => array
145            .as_primitive::<Int32Type>()
146            .unary(|v| hash_long(v as i64)),
147        DataType::Int64 => array.as_primitive::<Int64Type>().unary(hash_long),
148        DataType::UInt8 => array
149            .as_primitive::<UInt8Type>()
150            .unary(|v| hash_long(v as i64)),
151        DataType::UInt16 => array
152            .as_primitive::<UInt16Type>()
153            .unary(|v| hash_long(v as i64)),
154        DataType::UInt32 => array
155            .as_primitive::<UInt32Type>()
156            .unary(|v| hash_long(v as i64)),
157        DataType::UInt64 => array
158            .as_primitive::<UInt64Type>()
159            .unary(|v| hash_long(v as i64)),
160        DataType::Dictionary(_, _) => {
161            let dictionary = array.as_any_dictionary();
162            let values = iceberg_hash_array(dictionary.values().as_ref())?;
163            return Ok(take(values.as_ref(), dictionary.keys(), None)?);
164        }
165        other => {
166            return Err(DataFusionError::Execution(format!(
167                "{} does not support type {other}",
168                IcebergHash::NAME
169            )));
170        }
171    };
172    Ok(Arc::new(hashes))
173}
174
175fn hash_long(value: i64) -> i32 {
176    murmur3_32(&value.to_le_bytes())
177}
178
179/// MurmurHash3 x86 32-bit with seed 0, returned as a signed integer like Iceberg.
180fn murmur3_32(data: &[u8]) -> i32 {
181    const C1: u32 = 0xcc9e_2d51;
182    const C2: u32 = 0x1b87_3593;
183
184    fn mix_k(k: u32) -> u32 {
185        k.wrapping_mul(C1).rotate_left(15).wrapping_mul(C2)
186    }
187
188    let mut h = 0u32;
189    let (blocks, tail) = data.as_chunks::<4>();
190    for block in blocks {
191        h = (h ^ mix_k(u32::from_le_bytes(*block)))
192            .rotate_left(13)
193            .wrapping_mul(5)
194            .wrapping_add(0xe654_6b64);
195    }
196    if !tail.is_empty() {
197        let k = tail
198            .iter()
199            .rev()
200            .fold(0u32, |k, byte| (k << 8) | *byte as u32);
201        h ^= mix_k(k);
202    }
203
204    // Murmur3 mixes in the length modulo 2^32.
205    h ^= data.len() as u32;
206    h ^= h >> 16;
207    h = h.wrapping_mul(0x85eb_ca6b);
208    h ^= h >> 13;
209    h = h.wrapping_mul(0xc2b2_ae35);
210    h ^= h >> 16;
211    h as i32
212}
213
214/// DataFusion form of [`IcebergHash`], so storage engines
215/// can evaluate a region's partition expression as a filter.
216#[derive(Debug, PartialEq, Eq, Hash)]
217struct IcebergHashUdf {
218    signature: Signature,
219}
220
221impl Default for IcebergHashUdf {
222    fn default() -> Self {
223        Self {
224            signature: Signature::any(1, Volatility::Immutable),
225        }
226    }
227}
228
229impl ScalarUDFImpl for IcebergHashUdf {
230    fn name(&self) -> &str {
231        IcebergHash::NAME
232    }
233
234    fn signature(&self) -> &Signature {
235        &self.signature
236    }
237
238    fn return_type(&self, _arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
239        Ok(DataType::Int32)
240    }
241
242    fn invoke_with_args(
243        &self,
244        args: ScalarFunctionArgs,
245    ) -> datafusion_common::Result<ColumnarValue> {
246        let [arg] = args.args.as_slice() else {
247            return Err(DataFusionError::Internal(format!(
248                "{} expects one argument, got {}",
249                IcebergHash::NAME,
250                args.args.len()
251            )));
252        };
253        match arg {
254            ColumnarValue::Array(array) => {
255                Ok(ColumnarValue::Array(iceberg_hash_array(array.as_ref())?))
256            }
257            ColumnarValue::Scalar(scalar) => {
258                let hashes = iceberg_hash_array(scalar.to_array()?.as_ref())?;
259                Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(
260                    &hashes, 0,
261                )?))
262            }
263        }
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use datafusion_common::arrow::array::{
270        Int64Array, StringArray, StringDictionaryBuilder, StringViewArray, UInt64Array,
271    };
272    use datafusion_common::arrow::datatypes::UInt16Type;
273
274    use super::*;
275
276    #[test]
277    fn test_iceberg_spec_vectors() {
278        // Iceberg table spec, appendix B.
279        assert_eq!(hash_long(34), 2017239379);
280        // date 2017-11-16 hashes as int 17486.
281        assert_eq!(hash_long(17486), -653330422);
282        // time 22:31:08 hashes as long 81068000000 micros.
283        assert_eq!(hash_long(81068000000), -662762989);
284        assert_eq!(murmur3_32("iceberg".as_bytes()), 1210000089);
285        assert_eq!(murmur3_32(&[0, 1, 2, 3]), -188683207);
286    }
287
288    #[test]
289    fn test_unsigned_policy() {
290        assert_eq!(
291            IcebergHash::evaluate_value(&Value::UInt64(34)).unwrap(),
292            Value::Int32(2017239379)
293        );
294        assert_eq!(
295            IcebergHash::evaluate_value(&Value::UInt64(u64::MAX)).unwrap(),
296            Value::Int32(hash_long(-1))
297        );
298    }
299
300    #[test]
301    fn test_array_matches_value() {
302        let strings = [Some("iceberg"), None, Some(""), Some("中🙂")];
303        let mut dictionary = StringDictionaryBuilder::<UInt16Type>::new();
304        for v in strings {
305            dictionary.append_option(v);
306        }
307        let string_arrays: Vec<ArrayRef> = vec![
308            Arc::new(StringArray::from(strings.to_vec())),
309            Arc::new(StringViewArray::from(strings.to_vec())),
310            Arc::new(dictionary.finish()),
311        ];
312        let string_values = strings.map(|v| v.map_or(Value::Null, Value::from));
313
314        let ints = [Some(34), None, Some(i64::MIN), Some(-1)];
315        let unsigned = [Some(34), None, Some(u64::MAX), Some(0)];
316        let cases = string_arrays
317            .into_iter()
318            .map(|array| (array, string_values.to_vec()))
319            .chain([
320                (
321                    Arc::new(Int64Array::from(ints.to_vec())) as ArrayRef,
322                    ints.map(|v| v.map_or(Value::Null, Value::Int64)).to_vec(),
323                ),
324                (
325                    Arc::new(UInt64Array::from(unsigned.to_vec())) as ArrayRef,
326                    unsigned
327                        .map(|v| v.map_or(Value::Null, Value::UInt64))
328                        .to_vec(),
329                ),
330            ]);
331        for (array, values) in cases {
332            let hashes = IcebergHash::evaluate_array(array.as_ref()).unwrap();
333            let hashes = hashes.as_primitive::<Int32Type>();
334            for (row, value) in values.iter().enumerate() {
335                let expected = IcebergHash::evaluate_value(value).unwrap();
336                let actual = if hashes.is_null(row) {
337                    Value::Null
338                } else {
339                    Value::Int32(hashes.value(row))
340                };
341                assert_eq!(actual, expected, "{:?} row {row}", array.data_type());
342            }
343        }
344    }
345}
partition/function.rs

partition/
function.rs
