Merge apache/main into shuffle-complex-type-perf

Author: andygrove

native/core/src/execution/jni_api.rs

@@ -56,6 +56,7 @@ use datafusion_spark::function::math::hex::SparkHex;
 use datafusion_spark::function::math::width_bucket::SparkWidthBucket;
 use datafusion_spark::function::string::char::CharFunc;
 use datafusion_spark::function::string::concat::SparkConcat;
+use datafusion_spark::function::string::luhn_check::SparkLuhnCheck;
 use datafusion_spark::function::string::space::SparkSpace;
 use futures::poll;
 use futures::stream::StreamExt;
@@ -403,6 +404,7 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) {
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkWidthBucket::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(MapFromEntries::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkCrc32::default()));
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkLuhnCheck::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSpace::default()));
     session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitCount::default()));
 }
@@ -849,6 +851,13 @@ pub extern "system" fn Java_org_apache_comet_Native_sortRowPartitionsNative(
             tracing_enabled != JNI_FALSE,
             || {
                 // SAFETY: JVM unsafe memory allocation is aligned with long.
+                debug_assert!(address != 0, "sortRowPartitionsNative: null address");
+                debug_assert!(size >= 0, "sortRowPartitionsNative: negative size {size}");
+                debug_assert_eq!(
+                    (address as usize) % std::mem::align_of::<i64>(),
+                    0,
+                    "sortRowPartitionsNative: address not aligned to i64"
+                );
                 let array =
                     unsafe { std::slice::from_raw_parts_mut(address as *mut i64, size as usize) };
                 array.rdxsort();

native/core/src/execution/shuffle/spark_unsafe/list.rs

@@ -104,6 +104,7 @@ impl SparkUnsafeArray {
     pub fn new(addr: i64) -> Self {
         // SAFETY: addr points to valid Spark UnsafeArray data from the JVM.
         // The first 8 bytes contain the element count as a little-endian i64.
+        debug_assert!(addr != 0, "SparkUnsafeArray::new: null address");
         let slice: &[u8] = unsafe { std::slice::from_raw_parts(addr as *const u8, 8) };
         let num_elements = i64::from_le_bytes(slice.try_into().unwrap());
 
@@ -138,6 +139,11 @@ impl SparkUnsafeArray {
         // SAFETY: row_addr points to valid Spark UnsafeArray data. The null bitset starts
         // at offset 8 and contains ceil(num_elements/64) * 8 bytes. The caller ensures
         // index < num_elements, so word_offset is within the bitset region.
+        debug_assert!(
+            index < self.num_elements,
+            "is_null_at: index {index} >= num_elements {}",
+            self.num_elements
+        );
         unsafe {
             let mask: i64 = 1i64 << (index & 0x3f);
             let word_offset = (self.row_addr + 8 + (((index >> 6) as i64) << 3)) as *const i64;

native/core/src/execution/shuffle/spark_unsafe/map.rs

@@ -32,6 +32,8 @@ impl SparkUnsafeMap {
     pub(crate) fn new(addr: i64, size: i32) -> Self {
         // SAFETY: addr points to valid Spark UnsafeMap data from the JVM.
         // The first 8 bytes contain the key array size as a little-endian i64.
+        debug_assert!(addr != 0, "SparkUnsafeMap::new: null address");
+        debug_assert!(size >= 0, "SparkUnsafeMap::new: negative size {size}");
         let slice: &[u8] = unsafe { std::slice::from_raw_parts(addr as *const u8, 8) };
         let key_array_size = i64::from_le_bytes(slice.try_into().unwrap());
 

native/core/src/execution/shuffle/spark_unsafe/row.rs

@@ -93,6 +93,10 @@ pub trait SparkUnsafeObject {
         let addr = self.get_element_offset(index, 1);
         // SAFETY: addr points to valid element data within the UnsafeRow/UnsafeArray region.
         // The caller ensures index is within bounds.
+        debug_assert!(
+            !addr.is_null(),
+            "get_boolean: null pointer at index {index}"
+        );
         unsafe { *addr != 0 }
     }
 
@@ -101,6 +105,7 @@ pub trait SparkUnsafeObject {
     fn get_byte(&self, index: usize) -> i8 {
         let addr = self.get_element_offset(index, 1);
         // SAFETY: addr points to valid element data (1 byte) within the row/array region.
+        debug_assert!(!addr.is_null(), "get_byte: null pointer at index {index}");
         let slice: &[u8] = unsafe { std::slice::from_raw_parts(addr, 1) };
         i8::from_le_bytes(slice.try_into().unwrap())
     }
@@ -110,6 +115,7 @@ pub trait SparkUnsafeObject {
     fn get_short(&self, index: usize) -> i16 {
         let addr = self.get_element_offset(index, 2);
         // SAFETY: addr points to valid element data (2 bytes) within the row/array region.
+        debug_assert!(!addr.is_null(), "get_short: null pointer at index {index}");
         let slice: &[u8] = unsafe { std::slice::from_raw_parts(addr, 2) };
         i16::from_le_bytes(slice.try_into().unwrap())
     }
@@ -119,6 +125,7 @@ pub trait SparkUnsafeObject {
     fn get_int(&self, index: usize) -> i32 {
         let addr = self.get_element_offset(index, 4);
         // SAFETY: addr points to valid element data (4 bytes) within the row/array region.
+        debug_assert!(!addr.is_null(), "get_int: null pointer at index {index}");
         let slice: &[u8] = unsafe { std::slice::from_raw_parts(addr, 4) };
         i32::from_le_bytes(slice.try_into().unwrap())
     }
@@ -128,6 +135,7 @@ pub trait SparkUnsafeObject {
     fn get_long(&self, index: usize) -> i64 {
         let addr = self.get_element_offset(index, 8);
         // SAFETY: addr points to valid element data (8 bytes) within the row/array region.
+        debug_assert!(!addr.is_null(), "get_long: null pointer at index {index}");
         let slice: &[u8] = unsafe { std::slice::from_raw_parts(addr, 8) };
         i64::from_le_bytes(slice.try_into().unwrap())
     }
@@ -137,6 +145,7 @@ pub trait SparkUnsafeObject {
     fn get_float(&self, index: usize) -> f32 {
         let addr = self.get_element_offset(index, 4);
         // SAFETY: addr points to valid element data (4 bytes) within the row/array region.
+        debug_assert!(!addr.is_null(), "get_float: null pointer at index {index}");
         let slice: &[u8] = unsafe { std::slice::from_raw_parts(addr, 4) };
         f32::from_le_bytes(slice.try_into().unwrap())
     }
@@ -146,6 +155,7 @@ pub trait SparkUnsafeObject {
     fn get_double(&self, index: usize) -> f64 {
         let addr = self.get_element_offset(index, 8);
         // SAFETY: addr points to valid element data (8 bytes) within the row/array region.
+        debug_assert!(!addr.is_null(), "get_double: null pointer at index {index}");
         let slice: &[u8] = unsafe { std::slice::from_raw_parts(addr, 8) };
         f64::from_le_bytes(slice.try_into().unwrap())
     }
@@ -156,6 +166,11 @@ pub trait SparkUnsafeObject {
         let addr = self.get_row_addr() + offset as i64;
         // SAFETY: addr points to valid UTF-8 string data within the variable-length region.
         // Offset and length are read from the fixed-length portion of the row/array.
+        debug_assert!(addr != 0, "get_string: null address at index {index}");
+        debug_assert!(
+            len >= 0,
+            "get_string: negative length {len} at index {index}"
+        );
         let slice: &[u8] = unsafe { std::slice::from_raw_parts(addr as *const u8, len as usize) };
 
         from_utf8(slice).unwrap()
@@ -167,6 +182,11 @@ pub trait SparkUnsafeObject {
         let addr = self.get_row_addr() + offset as i64;
         // SAFETY: addr points to valid binary data within the variable-length region.
         // Offset and length are read from the fixed-length portion of the row/array.
+        debug_assert!(addr != 0, "get_binary: null address at index {index}");
+        debug_assert!(
+            len >= 0,
+            "get_binary: negative length {len} at index {index}"
+        );
         unsafe { std::slice::from_raw_parts(addr as *const u8, len as usize) }
     }
 
@@ -175,6 +195,7 @@ pub trait SparkUnsafeObject {
     fn get_date(&self, index: usize) -> i32 {
         let addr = self.get_element_offset(index, 4);
         // SAFETY: addr points to valid element data (4 bytes) within the row/array region.
+        debug_assert!(!addr.is_null(), "get_date: null pointer at index {index}");
         let slice: &[u8] = unsafe { std::slice::from_raw_parts(addr, 4) };
         i32::from_le_bytes(slice.try_into().unwrap())
     }
@@ -184,6 +205,10 @@ pub trait SparkUnsafeObject {
     fn get_timestamp(&self, index: usize) -> i64 {
         let addr = self.get_element_offset(index, 8);
         // SAFETY: addr points to valid element data (8 bytes) within the row/array region.
+        debug_assert!(
+            !addr.is_null(),
+            "get_timestamp: null pointer at index {index}"
+        );
         let slice: &[u8] = unsafe { std::slice::from_raw_parts(addr, 8) };
         i64::from_le_bytes(slice.try_into().unwrap())
     }
@@ -296,6 +321,7 @@ impl SparkUnsafeRow {
         // SAFETY: row_addr points to valid Spark UnsafeRow data with at least
         // ceil(num_fields/64) * 8 bytes of null bitset. The caller ensures index < num_fields.
         // word_offset is within the bitset region since (index >> 6) << 3 < bitset size.
+        debug_assert!(self.row_addr != -1, "is_null_at: row not initialized");
         unsafe {
             let mask: i64 = 1i64 << (index & 0x3f);
             let word_offset = (self.row_addr + (((index >> 6) as i64) << 3)) as *const i64;
@@ -310,6 +336,7 @@ impl SparkUnsafeRow {
         // ceil(num_fields/64) * 8 bytes of null bitset. The caller ensures index < num_fields.
         // word_offset is within the bitset region since (index >> 6) << 3 < bitset size.
         // Writing is safe because we have mutable access and the memory is owned by the JVM.
+        debug_assert!(self.row_addr != -1, "set_not_null_at: row not initialized");
         unsafe {
             let mask: i64 = 1i64 << (index & 0x3f);
             let word_offset = (self.row_addr + (((index >> 6) as i64) << 3)) as *mut i64;

native/core/src/execution/utils.rs

@@ -97,6 +97,16 @@ impl SparkArrowConvert for ArrayData {
             }
         } else {
             // SAFETY: `array_ptr` and `schema_ptr` are aligned correctly.
+            debug_assert_eq!(
+                array_ptr.align_offset(array_align),
+                0,
+                "move_to_spark: array_ptr not aligned"
+            );
+            debug_assert_eq!(
+                schema_ptr.align_offset(schema_align),
+                0,
+                "move_to_spark: schema_ptr not aligned"
+            );
             unsafe {
                 std::ptr::write(array_ptr, FFI_ArrowArray::new(self));
                 std::ptr::write(schema_ptr, FFI_ArrowSchema::try_from(self.data_type())?);

native/core/src/jvm_bridge/mod.rs

@@ -263,11 +263,19 @@ impl JVMClasses<'_> {
     }
 
     pub fn get() -> &'static JVMClasses<'static> {
+        debug_assert!(
+            JVM_CLASSES.get().is_some(),
+            "JVMClasses::get: not initialized"
+        );
         unsafe { JVM_CLASSES.get_unchecked() }
     }
 
     /// Gets the JNIEnv for the current thread.
     pub fn get_env() -> CometResult<AttachGuard<'static>> {
+        debug_assert!(
+            JAVA_VM.get().is_some(),
+            "JVMClasses::get_env: JAVA_VM not initialized"
+        );
         unsafe {
             let java_vm = JAVA_VM.get_unchecked();
             java_vm.attach_current_thread().map_err(|e| {

spark/src/main/scala/org/apache/comet/serde/statics.scala

@@ -19,7 +19,7 @@
 
 package org.apache.comet.serde
 
-import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.expressions.{Attribute, ExpressionImplUtils}
 import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke
 import org.apache.spark.sql.catalyst.util.CharVarcharCodegenUtils
 
@@ -34,7 +34,8 @@ object CometStaticInvoke extends CometExpressionSerde[StaticInvoke] {
       : Map[(String, Class[_]), CometExpressionSerde[StaticInvoke]] =
     Map(
       ("readSidePadding", classOf[CharVarcharCodegenUtils]) -> CometScalarFunction(
-        "read_side_padding"))
+        "read_side_padding"),
+      ("isLuhnNumber", classOf[ExpressionImplUtils]) -> CometScalarFunction("luhn_check"))
 
   override def convert(
       expr: StaticInvoke,

spark/src/test/resources/sql-tests/expressions/string/luhn_check.sql

@@ -0,0 +1,33 @@
+-- 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.
+
+-- MinSparkVersion: 3.5
+-- ConfigMatrix: parquet.enable.dictionary=false,true
+
+statement
+CREATE TABLE test_luhn(s string) USING parquet
+
+statement
+INSERT INTO test_luhn VALUES ('79927398710'), ('79927398713'), ('1234567812345670'), ('0'), (''), ('abc'), (NULL)
+
+-- column input
+query
+SELECT luhn_check(s) FROM test_luhn
+
+-- literal arguments
+query
+SELECT luhn_check('79927398713'), luhn_check('79927398710'), luhn_check(''), luhn_check(NULL)