From 8ae7768a5056d3745fb521e208976311537bb1e7 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Wed, 11 Nov 2026 19:49:50 +0700 Subject: [PATCH] feat: add a native kernel for Delta's InterleaveBits Adds delta_interleave_bits, a scalar function that builds Delta Lake's Z-order key for one or two Int32 inputs: bits are taken from the most significant down, alternating between inputs with the first input in the higher position, a null input counts as zero, and the result is four bytes per input and never null. The tests check the output against values produced by Delta's own InterleaveBits on both its default and fast code paths, including signed extrema and nulls. Nothing calls the function from Spark yet; the Delta serde that uses it comes separately. --- native/spark-expr/src/comet_scalar_funcs.rs | 9 +- .../src/delta_funcs/interleave_bits.rs | 421 ++++++++++++++++++ native/spark-expr/src/delta_funcs/mod.rs | 22 + native/spark-expr/src/lib.rs | 2 + 4 files changed, 450 insertions(+), 4 deletions(-) create mode 100644 native/spark-expr/src/delta_funcs/interleave_bits.rs create mode 100644 native/spark-expr/src/delta_funcs/mod.rs diff --git a/native/spark-expr/src/comet_scalar_funcs.rs b/native/spark-expr/src/comet_scalar_funcs.rs index 0086538b276..e7f968166be 100644 --- a/native/spark-expr/src/comet_scalar_funcs.rs +++ b/native/spark-expr/src/comet_scalar_funcs.rs @@ -32,10 +32,10 @@ use crate::{ spark_ceil, spark_day_name, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_isnan, spark_lpad, spark_make_decimal, spark_month_name, spark_read_side_padding, spark_round, spark_rpad, spark_sequence, spark_sequence_recomputed, spark_to_time, spark_unhex, - spark_unscaled_value, EvalMode, SparkArrayExtrema, SparkArrayPositionFunc, SparkArrayRemove, - SparkArraySetOp, SparkArraySlice, SparkArraySort, SparkArraysOverlap, SparkContains, - SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkDayOfWeek, SparkFlatten, - SparkFloatArrayContains, SparkIcebergBucket, SparkIcebergTemporalTransform, + spark_unscaled_value, DeltaInterleaveBits, EvalMode, SparkArrayExtrema, SparkArrayPositionFunc, + SparkArrayRemove, SparkArraySetOp, SparkArraySlice, SparkArraySort, SparkArraysOverlap, + SparkContains, SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkDayOfWeek, + SparkFlatten, SparkFloatArrayContains, SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate, SparkMakeDate, SparkMakeInterval, SparkMakeTime, SparkMapExtract, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc, SparkSortArray, SparkWeekDay, }; @@ -365,6 +365,7 @@ fn all_scalar_functions() -> Vec> { Arc::new(ScalarUDF::new_from_impl(SparkDateTrunc::default())), Arc::new(ScalarUDF::new_from_impl(SparkDayOfWeek::default())), Arc::new(ScalarUDF::new_from_impl(SparkWeekDay::default())), + Arc::new(ScalarUDF::new_from_impl(DeltaInterleaveBits::default())), Arc::new(ScalarUDF::new_from_impl(SparkFlatten::default())), Arc::new(ScalarUDF::new_from_impl(SparkIcebergBucket::default())), Arc::new(ScalarUDF::new_from_impl(SparkIcebergTruncate::default())), diff --git a/native/spark-expr/src/delta_funcs/interleave_bits.rs b/native/spark-expr/src/delta_funcs/interleave_bits.rs new file mode 100644 index 00000000000..3480a13a249 --- /dev/null +++ b/native/spark-expr/src/delta_funcs/interleave_bits.rs @@ -0,0 +1,421 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{Array, AsArray, BinaryArray, BinaryBuilder, Int32Array}; +use arrow::datatypes::{DataType, Int32Type}; +use datafusion::common::{plan_err, Result, ScalarValue}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, +}; +use std::sync::Arc; + +const NAME: &str = "delta_interleave_bits"; + +/// Largest number of inputs the kernel accepts; two 32-bit inputs fill the `u64` key. +const MAX_INPUTS: usize = 2; + +const BITS_PER_INPUT: u32 = i32::BITS; +const BYTES_PER_INPUT: usize = 4; + +/// Delta Lake's `InterleaveBits`, the Z-order key `OPTIMIZE ... ZORDER BY` sorts on. +/// +/// Takes one or two `Int32` inputs and returns `4 * n` bytes: bits 31 down to 0 of each input's +/// two's-complement value, round-robin in argument order, so the first input supplies the higher +/// bit of each group. A null input counts as 0 and the result is never null. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct DeltaInterleaveBits { + signature: Signature, +} + +impl DeltaInterleaveBits { + pub fn new() -> Self { + Self { + signature: Signature::one_of( + vec![ + TypeSignature::Exact(vec![DataType::Int32]), + TypeSignature::Exact(vec![DataType::Int32, DataType::Int32]), + ], + Volatility::Immutable, + ), + } + } +} + +impl Default for DeltaInterleaveBits { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for DeltaInterleaveBits { + fn name(&self) -> &str { + NAME + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + check_arg_types(arg_types)?; + Ok(DataType::Binary) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let number_rows = args.number_rows; + let args = args.args; + check_arg_types(&args.iter().map(|a| a.data_type()).collect::>())?; + let all_scalar = args.iter().all(|a| matches!(a, ColumnarValue::Scalar(_))); + let rows = if all_scalar { 1 } else { number_rows }; + let arrays = args + .iter() + .map(|a| a.to_array(rows)) + .collect::>>()?; + let columns = arrays + .iter() + .map(|a| a.as_primitive::()) + .collect::>(); + let result = interleave_columns(&columns, rows); + if all_scalar { + Ok(ColumnarValue::Scalar(ScalarValue::try_from_array( + &result, 0, + )?)) + } else { + Ok(ColumnarValue::Array(Arc::new(result))) + } + } +} + +fn check_arg_types(arg_types: &[DataType]) -> Result<()> { + if !(1..=MAX_INPUTS).contains(&arg_types.len()) { + return plan_err!( + "{NAME} expects 1 to {MAX_INPUTS} arguments, got {}", + arg_types.len() + ); + } + match arg_types.iter().find(|t| **t != DataType::Int32) { + Some(other) => plan_err!("{NAME} expects Int32 input, got {other}"), + None => Ok(()), + } +} + +fn interleave_columns(columns: &[&Int32Array], rows: usize) -> BinaryArray { + let width = BYTES_PER_INPUT * columns.len(); + let mut builder = BinaryBuilder::with_capacity(rows, rows * width); + let mut values = [0u32; MAX_INPUTS]; + for row in 0..rows { + for (value, column) in values.iter_mut().zip(columns) { + *value = if column.is_null(row) { + 0 + } else { + column.value(row) as u32 + }; + } + let key = interleave(&values[..columns.len()]); + builder.append_value(&key.to_be_bytes()[size_of::() - width..]); + } + builder.finish() +} + +/// Interleaves `values` MSB-first, one bit from each value in turn, into the low +/// `32 * values.len()` bits of the key. +fn interleave(values: &[u32]) -> u64 { + let mut key = 0u64; + for pos in (0..BITS_PER_INPUT).rev() { + for &value in values { + key = (key << 1) | u64::from((value >> pos) & 1); + } + } + key +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::ArrayRef; + use arrow::datatypes::Field; + use datafusion::config::ConfigOptions; + + /// Interleaved keys Delta's own `InterleaveBits.interleaveBits` produced for these inputs, + /// with the fast path both on and off (the two agreed on every row). + const DELTA_GOLDENS: &[(&[i32], &str)] = &[ + (&[i32::MIN], "80000000"), + (&[i32::MAX], "7FFFFFFF"), + (&[0], "00000000"), + (&[-1], "FFFFFFFF"), + (&[1], "00000001"), + (&[1431655765], "55555555"), + (&[-1431655766], "AAAAAAAA"), + (&[1044735], "000FF0FF"), + (&[-1044736], "FFF00F00"), + (&[16909060], "01020304"), + (&[i32::MIN, i32::MIN], "C000000000000000"), + (&[i32::MIN, i32::MAX], "9555555555555555"), + (&[i32::MIN, 0], "8000000000000000"), + (&[i32::MIN, -1], "D555555555555555"), + (&[i32::MIN, 1], "8000000000000001"), + (&[i32::MIN, 1431655765], "9111111111111111"), + (&[i32::MIN, -1431655766], "C444444444444444"), + (&[i32::MIN, 1044735], "8000005555005555"), + (&[i32::MIN, -1044736], "D555550000550000"), + (&[i32::MIN, 16909060], "8001000400050010"), + (&[i32::MAX, i32::MIN], "6AAAAAAAAAAAAAAA"), + (&[i32::MAX, i32::MAX], "3FFFFFFFFFFFFFFF"), + (&[i32::MAX, 0], "2AAAAAAAAAAAAAAA"), + (&[i32::MAX, -1], "7FFFFFFFFFFFFFFF"), + (&[i32::MAX, 1], "2AAAAAAAAAAAAAAB"), + (&[i32::MAX, 1431655765], "3BBBBBBBBBBBBBBB"), + (&[i32::MAX, -1431655766], "6EEEEEEEEEEEEEEE"), + (&[i32::MAX, 1044735], "2AAAAAFFFFAAFFFF"), + (&[i32::MAX, -1044736], "7FFFFFAAAAFFAAAA"), + (&[i32::MAX, 16909060], "2AABAAAEAAAFAABA"), + (&[0, i32::MIN], "4000000000000000"), + (&[0, i32::MAX], "1555555555555555"), + (&[0, 0], "0000000000000000"), + (&[0, -1], "5555555555555555"), + (&[0, 1], "0000000000000001"), + (&[0, 1431655765], "1111111111111111"), + (&[0, -1431655766], "4444444444444444"), + (&[0, 1044735], "0000005555005555"), + (&[0, -1044736], "5555550000550000"), + (&[0, 16909060], "0001000400050010"), + (&[-1, i32::MIN], "EAAAAAAAAAAAAAAA"), + (&[-1, i32::MAX], "BFFFFFFFFFFFFFFF"), + (&[-1, 0], "AAAAAAAAAAAAAAAA"), + (&[-1, -1], "FFFFFFFFFFFFFFFF"), + (&[-1, 1], "AAAAAAAAAAAAAAAB"), + (&[-1, 1431655765], "BBBBBBBBBBBBBBBB"), + (&[-1, -1431655766], "EEEEEEEEEEEEEEEE"), + (&[-1, 1044735], "AAAAAAFFFFAAFFFF"), + (&[-1, -1044736], "FFFFFFAAAAFFAAAA"), + (&[-1, 16909060], "AAABAAAEAAAFAABA"), + (&[1, i32::MIN], "4000000000000002"), + (&[1, i32::MAX], "1555555555555557"), + (&[1, 0], "0000000000000002"), + (&[1, -1], "5555555555555557"), + (&[1, 1], "0000000000000003"), + (&[1, 1431655765], "1111111111111113"), + (&[1, -1431655766], "4444444444444446"), + (&[1, 1044735], "0000005555005557"), + (&[1, -1044736], "5555550000550002"), + (&[1, 16909060], "0001000400050012"), + (&[1431655765, i32::MIN], "6222222222222222"), + (&[1431655765, i32::MAX], "3777777777777777"), + (&[1431655765, 0], "2222222222222222"), + (&[1431655765, -1], "7777777777777777"), + (&[1431655765, 1], "2222222222222223"), + (&[1431655765, 1431655765], "3333333333333333"), + (&[1431655765, -1431655766], "6666666666666666"), + (&[1431655765, 1044735], "2222227777227777"), + (&[1431655765, -1044736], "7777772222772222"), + (&[1431655765, 16909060], "2223222622272232"), + (&[-1431655766, i32::MIN], "C888888888888888"), + (&[-1431655766, i32::MAX], "9DDDDDDDDDDDDDDD"), + (&[-1431655766, 0], "8888888888888888"), + (&[-1431655766, -1], "DDDDDDDDDDDDDDDD"), + (&[-1431655766, 1], "8888888888888889"), + (&[-1431655766, 1431655765], "9999999999999999"), + (&[-1431655766, -1431655766], "CCCCCCCCCCCCCCCC"), + (&[-1431655766, 1044735], "888888DDDD88DDDD"), + (&[-1431655766, -1044736], "DDDDDD8888DD8888"), + (&[-1431655766, 16909060], "8889888C888D8898"), + (&[1044735, i32::MIN], "400000AAAA00AAAA"), + (&[1044735, i32::MAX], "155555FFFF55FFFF"), + (&[1044735, 0], "000000AAAA00AAAA"), + (&[1044735, -1], "555555FFFF55FFFF"), + (&[1044735, 1], "000000AAAA00AAAB"), + (&[1044735, 1431655765], "111111BBBB11BBBB"), + (&[1044735, -1431655766], "444444EEEE44EEEE"), + (&[1044735, 1044735], "000000FFFF00FFFF"), + (&[1044735, -1044736], "555555AAAA55AAAA"), + (&[1044735, 16909060], "000100AEAA05AABA"), + (&[-1044736, i32::MIN], "EAAAAA0000AA0000"), + (&[-1044736, i32::MAX], "BFFFFF5555FF5555"), + (&[-1044736, 0], "AAAAAA0000AA0000"), + (&[-1044736, -1], "FFFFFF5555FF5555"), + (&[-1044736, 1], "AAAAAA0000AA0001"), + (&[-1044736, 1431655765], "BBBBBB1111BB1111"), + (&[-1044736, -1431655766], "EEEEEE4444EE4444"), + (&[-1044736, 1044735], "AAAAAA5555AA5555"), + (&[-1044736, -1044736], "FFFFFF0000FF0000"), + (&[-1044736, 16909060], "AAABAA0400AF0010"), + (&[16909060, i32::MIN], "40020008000A0020"), + (&[16909060, i32::MAX], "1557555D555F5575"), + (&[16909060, 0], "00020008000A0020"), + (&[16909060, -1], "5557555D555F5575"), + (&[16909060, 1], "00020008000A0021"), + (&[16909060, 1431655765], "11131119111B1131"), + (&[16909060, -1431655766], "4446444C444E4464"), + (&[16909060, 1044735], "0002005D550A5575"), + (&[16909060, -1044736], "55575508005F0020"), + (&[16909060, 16909060], "0003000C000F0030"), + (&[-194441980, -713859073], "FB313D8711F75575"), + (&[1465097161, 551116650], "262A734B93DFB4C6"), + (&[-2044496387, -1402509360], "C4781C1F3E37FBA2"), + (&[-42021569, 936478219], "AFB77BA9E0B60AEF"), + (&[-371001683, -1494708188], "EC96FC48EB288CB2"), + (&[-11981971, -1746055248], "EBBF74D318DF6DA2"), + (&[1396948856, -1908031548], "625E301BCFFF7A90"), + (&[-964581762, -717158876], "F1399003883C2EB8"), + (&[-1802157460, 629602743], "8631C23675496DB5"), + (&[47397773, -796979038], "5108B75F0B8AC4A6"), + (&[2119901077, 859130386], "2FAD279B12EE8326"), + (&[-681693820, -2140680892], "E22A36BD58F39030"), + (&[2129441368, -1087931688], "6FFDACB59F4C73C0"), + (&[458445216, -393048690], "56CA630E6253C854"), + (&[1762056806, -117967399], "7DC2553DF7997969"), + (&[1232985102, -1191240484], "65C27FF7A3F051F8"), + (&[-650702380, -866216185], "F2D21B7E4317A235"), + (&[-1061812808, 807767500], "A5008E394073DAD0"), + (&[-293689653, 887702561], "ADB87EE99838A48B"), + (&[1552189062, -1186747436], "67E19025C46DD138"), + ]; + + fn invoke(args: Vec, number_rows: usize) -> Result { + DeltaInterleaveBits::new().invoke_with_args(ScalarFunctionArgs { + args, + arg_fields: vec![], + number_rows, + return_field: Arc::new(Field::new("result", DataType::Binary, false)), + config_options: Arc::new(ConfigOptions::default()), + }) + } + + fn int_array(values: Vec>) -> ColumnarValue { + ColumnarValue::Array(Arc::new(Int32Array::from(values))) + } + + /// Invokes the function on array arguments and returns each row as upper-case hex. + fn hex_rows(columns: Vec>>) -> Vec { + let rows = columns[0].len(); + let args = columns.into_iter().map(int_array).collect(); + let ColumnarValue::Array(result) = invoke(args, rows).unwrap() else { + panic!("expected an array result"); + }; + assert_eq!(result.null_count(), 0); + to_hex_rows(&result) + } + + fn to_hex_rows(array: &ArrayRef) -> Vec { + array + .as_binary::() + .iter() + .map(|bytes| to_hex(bytes.expect("result is never null"))) + .collect() + } + + fn to_hex(bytes: &[u8]) -> String { + bytes.iter().map(|b| format!("{b:02X}")).collect() + } + + #[test] + fn oracle_two_inputs() { + let rows = hex_rows(vec![ + vec![Some(0x000ff0ff)], + vec![Some(0xfff00f00_u32 as i32)], + ]); + assert_eq!(rows, ["555555AAAA55AAAA"]); + } + + #[test] + fn null_counts_as_zero() { + let rows = hex_rows(vec![vec![Some(-1), None], vec![None, Some(-1)]]); + assert_eq!(rows, ["AAAAAAAAAAAAAAAA", "5555555555555555"]); + } + + #[test] + fn one_input_is_big_endian() { + let rows = hex_rows(vec![vec![Some(i32::MIN), Some(-1), Some(0x01020304)]]); + assert_eq!(rows, ["80000000", "FFFFFFFF", "01020304"]); + } + + #[test] + fn scalar_broadcasts() { + let column = vec![Some(1), Some(2), None]; + let minus_one = || ColumnarValue::Scalar(ScalarValue::Int32(Some(-1))); + let cases = [ + ( + vec![int_array(column.clone()), minus_one()], + vec![column.clone(), vec![Some(-1); 3]], + ), + ( + vec![minus_one(), int_array(column.clone())], + vec![vec![Some(-1); 3], column.clone()], + ), + ]; + for (args, expanded) in cases { + let ColumnarValue::Array(result) = invoke(args, column.len()).unwrap() else { + panic!("expected an array result"); + }; + assert_eq!(to_hex_rows(&result), hex_rows(expanded)); + } + } + + #[test] + fn all_scalar_args() { + let cases: [Vec>; 5] = [ + vec![Some(0x000ff0ff), Some(0xfff00f00_u32 as i32)], + vec![None, Some(-1)], + vec![None, None], + vec![Some(0x01020304)], + vec![None], + ]; + for values in cases { + let args = values + .iter() + .map(|v| ColumnarValue::Scalar(ScalarValue::Int32(*v))) + .collect(); + let ColumnarValue::Scalar(ScalarValue::Binary(Some(bytes))) = invoke(args, 1).unwrap() + else { + panic!("expected a non-null binary scalar for {values:?}"); + }; + let expected = hex_rows(values.iter().map(|v| vec![*v]).collect()); + assert_eq!(vec![to_hex(&bytes)], expected, "{values:?}"); + } + } + + #[test] + fn delta_goldens() { + for (inputs, expected) in DELTA_GOLDENS { + let columns = inputs.iter().map(|v| vec![Some(*v)]).collect(); + assert_eq!(hex_rows(columns), [*expected], "{inputs:?}"); + } + } + + #[test] + fn rejects_bad_arity() { + let udf = DeltaInterleaveBits::new(); + for arity in [0, 3] { + assert!(udf.return_type(&vec![DataType::Int32; arity]).is_err()); + let args = vec![ColumnarValue::Scalar(ScalarValue::Int32(Some(1))); arity]; + assert!(invoke(args, 1).is_err()); + } + } + + #[test] + fn rejects_non_int32_input() { + let udf = DeltaInterleaveBits::new(); + for arg_types in [ + vec![DataType::Int64], + vec![DataType::Int32, DataType::Int64], + ] { + assert!(udf.return_type(&arg_types).is_err(), "{arg_types:?}"); + } + let args = vec![ColumnarValue::Scalar(ScalarValue::Int64(Some(1)))]; + assert!(invoke(args, 1).is_err()); + } +} diff --git a/native/spark-expr/src/delta_funcs/mod.rs b/native/spark-expr/src/delta_funcs/mod.rs new file mode 100644 index 00000000000..f31b4aad5b5 --- /dev/null +++ b/native/spark-expr/src/delta_funcs/mod.rs @@ -0,0 +1,22 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Native implementations of Delta Lake expressions. + +mod interleave_bits; + +pub use interleave_bits::DeltaInterleaveBits; diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index ea8316bcce5..a557a38c664 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -48,8 +48,10 @@ pub mod hash_funcs; mod string_funcs; mod datetime_funcs; +mod delta_funcs; mod iceberg_funcs; pub use agg_funcs::*; +pub use delta_funcs::DeltaInterleaveBits; pub use iceberg_funcs::{SparkIcebergBucket, SparkIcebergTemporalTransform, SparkIcebergTruncate}; pub use cast::{spark_cast, Cast, SparkCastOptions};