From e44c5236edc908a7a448961259df4f63b0beef15 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 17:37:40 -0600 Subject: [PATCH 1/5] feat: native uuid() implementation compatible with Spark --- .../expression-audits/misc_funcs.md | 7 + docs/source/user-guide/latest/expressions.md | 2 +- .../core/src/execution/expressions/random.rs | 18 +- .../execution/planner/expression_registry.rs | 4 + native/proto/src/proto/expr.proto | 9 + .../src/nondetermenistic_funcs/mod.rs | 2 + .../src/nondetermenistic_funcs/shuffle.rs | 8 + .../src/nondetermenistic_funcs/uuid.rs | 190 ++++++++++++++++++ .../apache/comet/serde/QueryPlanSerde.scala | 3 +- .../apache/comet/serde/nondetermenistic.scala | 28 ++- .../sql-tests/expressions/misc/uuid.sql | 43 ++++ .../expressions/misc/uuid_with_seed.sql | 57 ++++++ 12 files changed, 367 insertions(+), 4 deletions(-) create mode 100644 native/spark-expr/src/nondetermenistic_funcs/uuid.rs create mode 100644 spark/src/test/resources/sql-tests/expressions/misc/uuid.sql create mode 100644 spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql diff --git a/docs/source/contributor-guide/expression-audits/misc_funcs.md b/docs/source/contributor-guide/expression-audits/misc_funcs.md index 6fbefe1d247..bab152b4db5 100644 --- a/docs/source/contributor-guide/expression-audits/misc_funcs.md +++ b/docs/source/contributor-guide/expression-audits/misc_funcs.md @@ -82,4 +82,11 @@ - Spark 4.0.1 (audited 2026-05-27): identical to 3.4.3 except the resulting literal carries the default string collation. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. +## uuid + +- Spark 3.4.3 (audited 2026-07-24): `Uuid(randomSeed: Option[Long]) extends LeafExpression with Nondeterministic with ExpressionWithRandomSeed`. The analyzer's `ResolveRandomSeed` fills `randomSeed` with a random `Long`, so it is always defined before Comet sees the plan. Per partition it seeds `RandomUUIDGenerator(randomSeed + partitionIndex)`, a Commons Math3 `MersenneTwister`, and per row draws two `nextLong()`s, masks in the RFC 4122 version 4 and variant bits, and formats via `java.util.UUID.toString`. Only the no-argument `uuid()` form exists (no seed constructor). Comet emits a `Uuid` proto with the resolved seed and reproduces the generator bit for bit via `SparkMersenneTwister`. +- Spark 3.5.8 (audited 2026-07-24): identical to 3.4.3. +- Spark 4.0.1 (audited 2026-07-24): adds `def this(seed: Expression)`, exposing the `uuid(seed)` SQL form (the seed must be an integer or long literal, validated at analysis time). `RandomUUIDGenerator` and the per-row algorithm are unchanged, so results are identical to 3.4.3 for a given seed. +- Spark 4.1.1 (audited 2026-07-24): identical to 4.0.1, plus `withShiftedSeed`. No runtime change. + [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index efbd33ba1c0..4073742d0d6 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -498,7 +498,7 @@ expression-level). The `outer` variants are wired but marked `Incompatible`; the | `try_variant_get` | 🔜 | tracking [#4098](https://github.com/apache/datafusion-comet/issues/4098) | | `typeof` | ✅ | Foldable; resolved to a literal before Comet sees the plan | | `user` | ✅ | Resolved to a literal by the Spark analyzer before reaching Comet | -| `uuid` | 🔜 | Nondeterministic random UUID | +| `uuid` | ✅ | | | `variant_get` | 🔜 | tracking [#4098](https://github.com/apache/datafusion-comet/issues/4098) | --- diff --git a/native/core/src/execution/expressions/random.rs b/native/core/src/execution/expressions/random.rs index 6a11bdbfb99..be12d64ed99 100644 --- a/native/core/src/execution/expressions/random.rs +++ b/native/core/src/execution/expressions/random.rs @@ -22,7 +22,7 @@ use crate::extract_expr; use arrow::datatypes::SchemaRef; use datafusion::physical_expr::PhysicalExpr; use datafusion_comet_proto::spark_expression::Expr; -use datafusion_comet_spark_expr::{RandExpr, RandnExpr, ShuffleExpr}; +use datafusion_comet_spark_expr::{RandExpr, RandnExpr, ShuffleExpr, UuidExpr}; use std::sync::Arc; pub struct RandBuilder; @@ -57,6 +57,22 @@ impl ExpressionBuilder for ShuffleBuilder { } } +pub struct UuidBuilder; + +impl ExpressionBuilder for UuidBuilder { + fn build( + &self, + spark_expr: &Expr, + _input_schema: SchemaRef, + planner: &PhysicalPlanner, + ) -> Result, ExecutionError> { + let expr = extract_expr!(spark_expr, Uuid); + // Spark seeds a fresh generator per partition with `randomSeed + partitionIndex`. + let seed = expr.seed.wrapping_add(planner.partition().into()); + Ok(Arc::new(UuidExpr::new(seed))) + } +} + pub struct RandnBuilder; impl ExpressionBuilder for RandnBuilder { diff --git a/native/core/src/execution/planner/expression_registry.rs b/native/core/src/execution/planner/expression_registry.rs index 7fe7a477dd6..1b50b13c707 100644 --- a/native/core/src/execution/planner/expression_registry.rs +++ b/native/core/src/execution/planner/expression_registry.rs @@ -101,6 +101,7 @@ pub enum ExpressionType { Rand, Randn, Shuffle, + Uuid, SparkPartitionId, MonotonicallyIncreasingId, ArraysZip, @@ -376,6 +377,7 @@ impl ExpressionRegistry { Some(ExprStruct::Rand(_)) => Ok(ExpressionType::Rand), Some(ExprStruct::Randn(_)) => Ok(ExpressionType::Randn), Some(ExprStruct::Shuffle(_)) => Ok(ExpressionType::Shuffle), + Some(ExprStruct::Uuid(_)) => Ok(ExpressionType::Uuid), Some(ExprStruct::SparkPartitionId(_)) => Ok(ExpressionType::SparkPartitionId), Some(ExprStruct::MonotonicallyIncreasingId(_)) => { Ok(ExpressionType::MonotonicallyIncreasingId) @@ -409,6 +411,8 @@ impl ExpressionRegistry { .insert(ExpressionType::Randn, Box::new(RandnBuilder)); self.builders .insert(ExpressionType::Shuffle, Box::new(ShuffleBuilder)); + self.builders + .insert(ExpressionType::Uuid, Box::new(UuidBuilder)); } /// Register partition expression builders diff --git a/native/proto/src/proto/expr.proto b/native/proto/src/proto/expr.proto index cbb1043f21e..783dd254920 100644 --- a/native/proto/src/proto/expr.proto +++ b/native/proto/src/proto/expr.proto @@ -92,6 +92,7 @@ message Expr { JvmScalarUdf jvm_scalar_udf = 70; PreciseTimestampConversion precise_timestamp_conversion = 71; Shuffle shuffle = 72; + Uuid uuid = 73; } reserved 20; @@ -563,6 +564,14 @@ message Shuffle { int64 seed = 2; } +// Spark's Uuid returns a random RFC 4122 version 4 UUID string. It is +// non-deterministic: the resolved random seed is combined with the partition +// index and seeds a Commons Math3 MersenneTwister, matching +// org.apache.spark.sql.catalyst.util.RandomUUIDGenerator. +message Uuid { + int64 seed = 1; +} + // Spark's ArraysZip takes children: Seq[Expression] and names: Seq[Expression] // https://github.com/apache/spark/blob/branch-4.1/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala#L296 message ArraysZip { diff --git a/native/spark-expr/src/nondetermenistic_funcs/mod.rs b/native/spark-expr/src/nondetermenistic_funcs/mod.rs index 3675bb04b9c..e8ba6f122ac 100644 --- a/native/spark-expr/src/nondetermenistic_funcs/mod.rs +++ b/native/spark-expr/src/nondetermenistic_funcs/mod.rs @@ -20,7 +20,9 @@ pub mod monotonically_increasing_id; pub mod rand; pub mod randn; pub mod shuffle; +pub mod uuid; pub use rand::RandExpr; pub use randn::RandnExpr; pub use shuffle::ShuffleExpr; +pub use uuid::UuidExpr; diff --git a/native/spark-expr/src/nondetermenistic_funcs/shuffle.rs b/native/spark-expr/src/nondetermenistic_funcs/shuffle.rs index 490a499125a..28357903f6b 100644 --- a/native/spark-expr/src/nondetermenistic_funcs/shuffle.rs +++ b/native/spark-expr/src/nondetermenistic_funcs/shuffle.rs @@ -160,6 +160,14 @@ impl SparkMersenneTwister { ((y as u32) >> (32 - bits)) as i32 } + /// Port of `BitsStreamGenerator.nextLong()`: two 32-bit draws combined into a + /// signed 64-bit value. Used by `RandomUUIDGenerator`. + pub(crate) fn next_long(&mut self) -> i64 { + let high = (self.next(32) as i64) << 32; + let low = (self.next(32) as i64) & 0xffffffffi64; + high | low + } + /// Port of `BitsStreamGenerator.nextInt(int n)`. The caller always passes a /// strictly positive `n`, matching Spark's `random.nextInt(i + 1)`. fn next_int(&mut self, n: i32) -> i32 { diff --git a/native/spark-expr/src/nondetermenistic_funcs/uuid.rs b/native/spark-expr/src/nondetermenistic_funcs/uuid.rs new file mode 100644 index 00000000000..fd33c12e392 --- /dev/null +++ b/native/spark-expr/src/nondetermenistic_funcs/uuid.rs @@ -0,0 +1,190 @@ +// 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 crate::nondetermenistic_funcs::shuffle::SparkMersenneTwister; +use arrow::array::{RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Schema}; +use datafusion::common::Result; +use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_expr::PhysicalExpr; +use std::fmt::{Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::{Arc, Mutex}; + +/// Draw one RFC 4122 version 4 UUID string from the generator, matching +/// `org.apache.spark.sql.catalyst.util.RandomUUIDGenerator.getNextUUID`: two +/// `nextLong()` draws with the version (4) and variant (10) bits masked in. +fn next_uuid_string(rng: &mut SparkMersenneTwister) -> String { + let most = (rng.next_long() as u64 & 0xFFFF_FFFF_FFFF_0FFF) | 0x0000_0000_0000_4000; + let least = (rng.next_long() as u64 | 0x8000_0000_0000_0000) & 0xBFFF_FFFF_FFFF_FFFF; + // Matches `java.util.UUID.toString()`: lowercase, zero-padded, hyphenated. + format!( + "{:08x}-{:04x}-{:04x}-{:04x}-{:012x}", + (most >> 32) & 0xFFFF_FFFF, + (most >> 16) & 0xFFFF, + most & 0xFFFF, + (least >> 48) & 0xFFFF, + least & 0xFFFF_FFFF_FFFF, + ) +} + +/// Physical expression for Spark's `uuid()`. Like `ShuffleExpr`, the generator +/// state is kept in a `Mutex` so that it advances continuously across every +/// batch in a partition, matching Spark's stateful per-partition evaluation. +/// Spark seeds a fresh `RandomUUIDGenerator` (a Commons Math3 `MersenneTwister`) +/// per partition with `randomSeed + partitionIndex`. +#[derive(Debug)] +pub struct UuidExpr { + /// Random seed already combined with the partition index by the planner. + seed: i64, + state_holder: Arc>>, +} + +impl UuidExpr { + pub fn new(seed: i64) -> Self { + Self { + seed, + state_holder: Arc::new(Mutex::new(None)), + } + } +} + +impl Display for UuidExpr { + fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "Uuid({})", self.seed) + } +} + +impl PartialEq for UuidExpr { + fn eq(&self, other: &Self) -> bool { + self.seed.eq(&other.seed) + } +} + +impl Eq for UuidExpr {} + +impl Hash for UuidExpr { + fn hash(&self, state: &mut H) { + self.seed.hash(state); + } +} + +impl PhysicalExpr for UuidExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Utf8) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let num_rows = batch.num_rows(); + + let mut state = self.state_holder.lock().unwrap(); + let rng = state.get_or_insert_with(|| SparkMersenneTwister::new(self.seed)); + + let result: StringArray = (0..num_rows).map(|_| Some(next_uuid_string(rng))).collect(); + Ok(ColumnarValue::Array(Arc::new(result))) + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(Arc::new(UuidExpr::new(self.seed))) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(self, f) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, RecordBatchOptions}; + + fn empty_batch(num_rows: usize) -> RecordBatch { + RecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &RecordBatchOptions::new().with_row_count(Some(num_rows)), + ) + .unwrap() + } + + fn eval_uuids(seed: i64, num_rows: usize) -> Vec { + let batch = empty_batch(num_rows); + let expr = UuidExpr::new(seed); + let result = expr.evaluate(&batch).unwrap().into_array(num_rows).unwrap(); + let arr = result.as_any().downcast_ref::().unwrap(); + (0..arr.len()).map(|i| arr.value(i).to_string()).collect() + } + + #[test] + fn test_uuid_format_and_version() { + for uuid in eval_uuids(42, 20) { + // Canonical 8-4-4-4-12 lowercase hex. + assert_eq!(uuid.len(), 36); + let parts: Vec<&str> = uuid.split('-').collect(); + assert_eq!( + parts.iter().map(|p| p.len()).collect::>(), + vec![8, 4, 4, 4, 12] + ); + assert!(uuid + .chars() + .all(|c| c == '-' || c.is_ascii_hexdigit() && !c.is_ascii_uppercase())); + // Version 4: first nibble of the third group. + assert_eq!(parts[2].as_bytes()[0], b'4'); + // Variant 10xx: first nibble of the fourth group is 8, 9, a, or b. + assert!(matches!(parts[3].as_bytes()[0], b'8' | b'9' | b'a' | b'b')); + } + } + + #[test] + fn test_uuid_deterministic_for_seed() { + // Same seed -> identical sequence. + assert_eq!(eval_uuids(42, 5), eval_uuids(42, 5)); + // Different seed -> different sequence. + assert_ne!(eval_uuids(42, 5), eval_uuids(0, 5)); + } + + #[test] + fn test_uuid_state_advances_across_batches() { + // A single expression evaluated over two batches yields the same UUIDs as + // one batch of the combined size (state persists across batches). + let expr = UuidExpr::new(7); + let mut streamed = Vec::new(); + for n in [3usize, 4usize] { + let batch = empty_batch(n); + let arr = expr.evaluate(&batch).unwrap().into_array(n).unwrap(); + let arr = arr.as_any().downcast_ref::().unwrap(); + streamed.extend((0..arr.len()).map(|i| arr.value(i).to_string())); + } + assert_eq!(streamed, eval_uuids(7, 7)); + // All distinct. + let mut sorted = streamed.clone(); + sorted.sort(); + sorted.dedup(); + assert_eq!(sorted.len(), streamed.len()); + } +} diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 4892d6f966e..80cdeef178b 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -371,7 +371,8 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { classOf[SortOrder] -> CometSortOrder, classOf[StaticInvoke] -> CometStaticInvoke, classOf[TryEval] -> CometTryEval, - classOf[UnscaledValue] -> CometUnscaledValue) + classOf[UnscaledValue] -> CometUnscaledValue, + classOf[Uuid] -> CometUuid) base ++ sparkVersionSpecificMiscExpressions } diff --git a/spark/src/main/scala/org/apache/comet/serde/nondetermenistic.scala b/spark/src/main/scala/org/apache/comet/serde/nondetermenistic.scala index 8269896c2ec..a88ce92c9de 100644 --- a/spark/src/main/scala/org/apache/comet/serde/nondetermenistic.scala +++ b/spark/src/main/scala/org/apache/comet/serde/nondetermenistic.scala @@ -19,7 +19,9 @@ package org.apache.comet.serde -import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Literal, MonotonicallyIncreasingID, Rand, Randn, SparkPartitionID} +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Literal, MonotonicallyIncreasingID, Rand, Randn, SparkPartitionID, Uuid} + +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason object CometSparkPartitionId extends CometExpressionSerde[SparkPartitionID] { override def convert( @@ -68,6 +70,30 @@ sealed abstract class CometRandCommonSerde[T <: Expression] extends CometExpress } } +object CometUuid extends CometExpressionSerde[Uuid] { + + // Comet reproduces Spark's UUIDs exactly: the resolved seed is combined with the partition index + // and seeds the same Commons Math3 MersenneTwister that drives + // org.apache.spark.sql.catalyst.util.RandomUUIDGenerator, so results match Spark bit for bit. + override def convert( + expr: Uuid, + inputs: Seq[Attribute], + binding: Boolean): Option[ExprOuterClass.Expr] = { + // In a resolved plan `randomSeed` is always defined (resolution requires it). Guard anyway. + expr.randomSeed match { + case Some(seed) => + Some( + ExprOuterClass.Expr + .newBuilder() + .setUuid(ExprOuterClass.Uuid.newBuilder().setSeed(seed)) + .build()) + case None => + withFallbackReason(expr, "uuid requires a resolved random seed") + None + } + } +} + object CometRand extends CometRandCommonSerde[Rand] { override protected def seedExprOf(expr: Rand): Expression = expr.child diff --git a/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql b/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql new file mode 100644 index 00000000000..afb5bc7d664 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql @@ -0,0 +1,43 @@ +-- 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. + +-- uuid() runs natively. The raw random value cannot be compared against Spark because Spark assigns +-- a fresh random seed on each planning pass, so these queries assert deterministic properties that +-- hold identically on both engines. The seeded uuid(seed) form (Spark 4.0+) asserts bit-for-bit +-- equality with Spark and lives in uuid_with_seed.sql. + +statement +CREATE TABLE test_uuid(id int) USING parquet + +statement +INSERT INTO test_uuid VALUES (1), (2), (3), (4), (5) + +-- canonical form is 36 characters +query +SELECT length(uuid()) FROM test_uuid + +-- matches the RFC 4122 version 4 layout (version nibble 4, variant nibble 8/9/a/b), lowercase hex +query +SELECT uuid() RLIKE '^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$' FROM test_uuid + +-- version nibble is always '4' +query +SELECT substring(uuid(), 15, 1) FROM test_uuid + +-- variant nibble is one of 8, 9, a, b +query +SELECT substring(uuid(), 20, 1) RLIKE '^[89ab]$' FROM test_uuid diff --git a/spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql b/spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql new file mode 100644 index 00000000000..6d569be4e72 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql @@ -0,0 +1,57 @@ +-- 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. + +-- The one argument uuid(seed) form only exists in Spark 4.0+. With a fixed seed the output is +-- deterministic, so Comet must reproduce Spark's output exactly. Comet drives the same Commons +-- Math3 MersenneTwister as org.apache.spark.sql.catalyst.util.RandomUUIDGenerator, combining the +-- seed with the partition index, so these queries assert bit-for-bit equality with Spark in the +-- default query mode. + +-- MinSparkVersion: 4.0 + +statement +CREATE TABLE test_uuid_seed(id int) USING parquet + +statement +INSERT INTO test_uuid_seed VALUES (1), (2), (3), (4), (5) + +-- fixed seed, single row +query +SELECT uuid(0) + +-- fixed seed, multiple rows: the generator advances per row within the partition +query +SELECT uuid(42) FROM test_uuid_seed + +-- zero seed over multiple rows +query +SELECT uuid(0) FROM test_uuid_seed + +-- negative seed +query +SELECT uuid(-12345) FROM test_uuid_seed + +-- Long.MinValue and Long.MaxValue seeds +query +SELECT uuid(-9223372036854775808) FROM test_uuid_seed + +query +SELECT uuid(9223372036854775807) FROM test_uuid_seed + +-- seeded uuid combined with other expressions stays native and deterministic +query +SELECT upper(uuid(7)), length(uuid(7)) FROM test_uuid_seed From 0ba81e90da23b1274095e08d0e21ab84c141c385 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 24 Jul 2026 17:54:46 -0600 Subject: [PATCH 2/5] refactor: share MersenneTwister, use uuid crate for formatting Extract SparkMersenneTwister into internal/mersenne.rs so shuffle and uuid both depend on it rather than uuid reaching into the shuffle module. Format UUIDs via the uuid crate encoded into a pre-sized StringBuilder, removing per-row String allocation and the hand-rolled hex formatting. --- native/Cargo.lock | 1 + native/spark-expr/Cargo.toml | 1 + .../internal/mersenne.rs | 172 +++++++++++++++ .../nondetermenistic_funcs/internal/mod.rs | 1 + .../src/nondetermenistic_funcs/shuffle.rs | 203 +++--------------- .../src/nondetermenistic_funcs/uuid.rs | 76 ++++--- .../sql-tests/expressions/misc/uuid.sql | 11 +- 7 files changed, 237 insertions(+), 228 deletions(-) create mode 100644 native/spark-expr/src/nondetermenistic_funcs/internal/mersenne.rs diff --git a/native/Cargo.lock b/native/Cargo.lock index 8194fdd70ab..652a9235121 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2074,6 +2074,7 @@ dependencies = [ "serde_json", "tokio", "twox-hash", + "uuid", ] [[package]] diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index c05ae897931..d790df4984d 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -43,6 +43,7 @@ futures = { workspace = true } twox-hash = "2.1.2" rand = { workspace = true } base64 = "0.23.0" +uuid = "1.23.3" [dev-dependencies] arrow = {workspace = true} diff --git a/native/spark-expr/src/nondetermenistic_funcs/internal/mersenne.rs b/native/spark-expr/src/nondetermenistic_funcs/internal/mersenne.rs new file mode 100644 index 00000000000..783fd9e3600 --- /dev/null +++ b/native/spark-expr/src/nondetermenistic_funcs/internal/mersenne.rs @@ -0,0 +1,172 @@ +// 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. + +//! A bit-for-bit port of Apache Commons Math3's `MersenneTwister`, the PRNG Spark +//! seeds per partition with `randomSeed + partitionIndex`. It backs both `shuffle` +//! (via `RandomIndicesGenerator`) and `uuid` (via `RandomUUIDGenerator`). +//! +//! See: +//! - `org/apache/commons/math3/random/MersenneTwister.java` +//! - `org/apache/commons/math3/random/BitsStreamGenerator.java` + +#[derive(Debug, Clone)] +pub(crate) struct SparkMersenneTwister { + /// Bytes pool. + mt: [i32; Self::N], + /// Current index in the bytes pool. + mti: usize, +} + +impl SparkMersenneTwister { + /// Size of the bytes pool. + const N: usize = 624; + /// Period second parameter. + const M: usize = 397; + /// X * MATRIX_A for X = {0, 1}. + const MAG01: [i32; 2] = [0x0, 0x9908b0dfu32 as i32]; + + pub(crate) fn new(seed: i64) -> Self { + // `mti` is set by seeding before it is ever read; 0 is just a placeholder. + let mut twister = SparkMersenneTwister { + mt: [0i32; Self::N], + mti: 0, + }; + twister.set_seed_long(seed); + twister + } + + fn set_seed_int(&mut self, seed: i32) { + // We use a long masked by 0xffffffff as a poor man's unsigned int. + let mut long_mt = seed as i64; + self.mt[0] = long_mt as i32; + let mut mti = 1usize; + while mti < Self::N { + long_mt = (1812433253i64.wrapping_mul(long_mt ^ (long_mt >> 30)) + mti as i64) + & 0xffffffffi64; + self.mt[mti] = long_mt as i32; + mti += 1; + } + self.mti = mti; + } + + fn set_seed_int_array(&mut self, seed: &[i32]) { + self.set_seed_int(19650218); + let mut i = 1usize; + let mut j = 0usize; + + for _ in 0..Self::N.max(seed.len()) { + let mt_i = self.mt[i] as i64 & 0xffffffffi64; + let mt_im1 = self.mt[i - 1] as i64 & 0xffffffffi64; + let l = (mt_i ^ ((mt_im1 ^ (mt_im1 >> 30)).wrapping_mul(1664525))) + .wrapping_add(seed[j] as i64) + .wrapping_add(j as i64); + self.mt[i] = (l & 0xffffffffi64) as i32; + i += 1; + j += 1; + if i >= Self::N { + self.mt[0] = self.mt[Self::N - 1]; + i = 1; + } + if j >= seed.len() { + j = 0; + } + } + + for _ in 0..(Self::N - 1) { + let mt_i = self.mt[i] as i64 & 0xffffffffi64; + let mt_im1 = self.mt[i - 1] as i64 & 0xffffffffi64; + let l = (mt_i ^ ((mt_im1 ^ (mt_im1 >> 30)).wrapping_mul(1566083941))) + .wrapping_sub(i as i64); + self.mt[i] = (l & 0xffffffffi64) as i32; + i += 1; + if i >= Self::N { + self.mt[0] = self.mt[Self::N - 1]; + i = 1; + } + } + + // MSB is 1, assuring a non-zero initial array. + self.mt[0] = 0x80000000u32 as i32; + } + + fn set_seed_long(&mut self, seed: i64) { + self.set_seed_int_array(&[(seed >> 32) as i32, seed as i32]); + } + + fn next(&mut self, bits: u32) -> i32 { + let mut y: i32; + if self.mti >= Self::N { + // Generate N words at one time. + let mut mt_next = self.mt[0]; + for k in 0..(Self::N - Self::M) { + let mt_curr = mt_next; + mt_next = self.mt[k + 1]; + y = (mt_curr & (0x80000000u32 as i32)) | (mt_next & 0x7fffffff); + self.mt[k] = self.mt[k + Self::M] + ^ (((y as u32) >> 1) as i32) + ^ Self::MAG01[(y & 0x1) as usize]; + } + for k in (Self::N - Self::M)..(Self::N - 1) { + let mt_curr = mt_next; + mt_next = self.mt[k + 1]; + y = (mt_curr & (0x80000000u32 as i32)) | (mt_next & 0x7fffffff); + self.mt[k] = self.mt[k + Self::M - Self::N] + ^ (((y as u32) >> 1) as i32) + ^ Self::MAG01[(y & 0x1) as usize]; + } + y = (mt_next & (0x80000000u32 as i32)) | (self.mt[0] & 0x7fffffff); + self.mt[Self::N - 1] = + self.mt[Self::M - 1] ^ (((y as u32) >> 1) as i32) ^ Self::MAG01[(y & 0x1) as usize]; + self.mti = 0; + } + + y = self.mt[self.mti]; + self.mti += 1; + + // Tempering. + y ^= ((y as u32) >> 11) as i32; + y ^= (y << 7) & (0x9d2c5680u32 as i32); + y ^= (y << 15) & (0xefc60000u32 as i32); + y ^= ((y as u32) >> 18) as i32; + + ((y as u32) >> (32 - bits)) as i32 + } + + /// Port of `BitsStreamGenerator.nextInt(int n)`. The caller always passes a + /// strictly positive `n`, matching Spark's `random.nextInt(i + 1)`. + pub(crate) fn next_int(&mut self, n: i32) -> i32 { + if (n & n.wrapping_neg()) == n { + // n is a power of two. + return ((n as i64 * self.next(31) as i64) >> 31) as i32; + } + loop { + let bits = self.next(31); + let val = bits % n; + if bits.wrapping_sub(val).wrapping_add(n.wrapping_sub(1)) >= 0 { + return val; + } + } + } + + /// Port of `BitsStreamGenerator.nextLong()`: two 32-bit draws combined into a + /// signed 64-bit value. Used by `RandomUUIDGenerator`. + pub(crate) fn next_long(&mut self) -> i64 { + let high = (self.next(32) as i64) << 32; + let low = (self.next(32) as i64) & 0xffffffffi64; + high | low + } +} diff --git a/native/spark-expr/src/nondetermenistic_funcs/internal/mod.rs b/native/spark-expr/src/nondetermenistic_funcs/internal/mod.rs index c7437f0667b..b708b89628f 100644 --- a/native/spark-expr/src/nondetermenistic_funcs/internal/mod.rs +++ b/native/spark-expr/src/nondetermenistic_funcs/internal/mod.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +pub(crate) mod mersenne; mod rand_utils; pub use rand_utils::evaluate_batch_for_rand; diff --git a/native/spark-expr/src/nondetermenistic_funcs/shuffle.rs b/native/spark-expr/src/nondetermenistic_funcs/shuffle.rs index 28357903f6b..081e9305251 100644 --- a/native/spark-expr/src/nondetermenistic_funcs/shuffle.rs +++ b/native/spark-expr/src/nondetermenistic_funcs/shuffle.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::nondetermenistic_funcs::internal::mersenne::SparkMersenneTwister; use arrow::array::{ Array, ArrayRef, FixedSizeListArray, GenericListArray, OffsetSizeTrait, RecordBatch, UInt64Array, @@ -29,184 +30,21 @@ use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::{Arc, Mutex}; -/// A bit-for-bit port of Apache Commons Math3's `MersenneTwister`, the PRNG Spark -/// uses to shuffle arrays. Spark seeds a fresh generator per partition with -/// `randomSeed + partitionIndex` and drives the "inside-out" Fisher-Yates -/// algorithm from `org.apache.spark.sql.catalyst.util.RandomIndicesGenerator`. -/// -/// See: -/// - `org/apache/commons/math3/random/MersenneTwister.java` -/// - `org/apache/commons/math3/random/BitsStreamGenerator.java` (`nextInt(int)`) -#[derive(Debug, Clone)] -pub(crate) struct SparkMersenneTwister { - /// Bytes pool. - mt: [i32; Self::N], - /// Current index in the bytes pool. - mti: usize, -} - -impl SparkMersenneTwister { - /// Size of the bytes pool. - const N: usize = 624; - /// Period second parameter. - const M: usize = 397; - /// X * MATRIX_A for X = {0, 1}. - const MAG01: [i32; 2] = [0x0, 0x9908b0dfu32 as i32]; - - pub(crate) fn new(seed: i64) -> Self { - // `mti` is set by seeding before it is ever read; 0 is just a placeholder. - let mut twister = SparkMersenneTwister { - mt: [0i32; Self::N], - mti: 0, - }; - twister.set_seed_long(seed); - twister - } - - fn set_seed_int(&mut self, seed: i32) { - // We use a long masked by 0xffffffff as a poor man's unsigned int. - let mut long_mt = seed as i64; - self.mt[0] = long_mt as i32; - let mut mti = 1usize; - while mti < Self::N { - long_mt = (1812433253i64.wrapping_mul(long_mt ^ (long_mt >> 30)) + mti as i64) - & 0xffffffffi64; - self.mt[mti] = long_mt as i32; - mti += 1; - } - self.mti = mti; - } - - fn set_seed_int_array(&mut self, seed: &[i32]) { - self.set_seed_int(19650218); - let mut i = 1usize; - let mut j = 0usize; - - for _ in 0..Self::N.max(seed.len()) { - let mt_i = self.mt[i] as i64 & 0xffffffffi64; - let mt_im1 = self.mt[i - 1] as i64 & 0xffffffffi64; - let l = (mt_i ^ ((mt_im1 ^ (mt_im1 >> 30)).wrapping_mul(1664525))) - .wrapping_add(seed[j] as i64) - .wrapping_add(j as i64); - self.mt[i] = (l & 0xffffffffi64) as i32; - i += 1; - j += 1; - if i >= Self::N { - self.mt[0] = self.mt[Self::N - 1]; - i = 1; - } - if j >= seed.len() { - j = 0; - } +/// Port of `RandomIndicesGenerator.getNextIndices`. Fills `out` with, for each +/// output position, the source position it should draw from. `out` is reused +/// across rows to avoid a per-row allocation. Advances the PRNG state (a Commons +/// Math3 `MersenneTwister`), which is shared across every row in the partition. +fn next_indices_into(rng: &mut SparkMersenneTwister, length: usize, out: &mut Vec) { + out.clear(); + out.resize(length, 0); + let mut i = 0usize; + while i < length { + let j = rng.next_int((i + 1) as i32) as usize; + if j != i { + out[i] = out[j]; } - - for _ in 0..(Self::N - 1) { - let mt_i = self.mt[i] as i64 & 0xffffffffi64; - let mt_im1 = self.mt[i - 1] as i64 & 0xffffffffi64; - let l = (mt_i ^ ((mt_im1 ^ (mt_im1 >> 30)).wrapping_mul(1566083941))) - .wrapping_sub(i as i64); - self.mt[i] = (l & 0xffffffffi64) as i32; - i += 1; - if i >= Self::N { - self.mt[0] = self.mt[Self::N - 1]; - i = 1; - } - } - - // MSB is 1, assuring a non-zero initial array. - self.mt[0] = 0x80000000u32 as i32; - } - - fn set_seed_long(&mut self, seed: i64) { - self.set_seed_int_array(&[(seed >> 32) as i32, seed as i32]); - } - - fn next(&mut self, bits: u32) -> i32 { - let mut y: i32; - if self.mti >= Self::N { - // Generate N words at one time. - let mut mt_next = self.mt[0]; - for k in 0..(Self::N - Self::M) { - let mt_curr = mt_next; - mt_next = self.mt[k + 1]; - y = (mt_curr & (0x80000000u32 as i32)) | (mt_next & 0x7fffffff); - self.mt[k] = self.mt[k + Self::M] - ^ (((y as u32) >> 1) as i32) - ^ Self::MAG01[(y & 0x1) as usize]; - } - for k in (Self::N - Self::M)..(Self::N - 1) { - let mt_curr = mt_next; - mt_next = self.mt[k + 1]; - y = (mt_curr & (0x80000000u32 as i32)) | (mt_next & 0x7fffffff); - self.mt[k] = self.mt[k + Self::M - Self::N] - ^ (((y as u32) >> 1) as i32) - ^ Self::MAG01[(y & 0x1) as usize]; - } - y = (mt_next & (0x80000000u32 as i32)) | (self.mt[0] & 0x7fffffff); - self.mt[Self::N - 1] = - self.mt[Self::M - 1] ^ (((y as u32) >> 1) as i32) ^ Self::MAG01[(y & 0x1) as usize]; - self.mti = 0; - } - - y = self.mt[self.mti]; - self.mti += 1; - - // Tempering. - y ^= ((y as u32) >> 11) as i32; - y ^= (y << 7) & (0x9d2c5680u32 as i32); - y ^= (y << 15) & (0xefc60000u32 as i32); - y ^= ((y as u32) >> 18) as i32; - - ((y as u32) >> (32 - bits)) as i32 - } - - /// Port of `BitsStreamGenerator.nextLong()`: two 32-bit draws combined into a - /// signed 64-bit value. Used by `RandomUUIDGenerator`. - pub(crate) fn next_long(&mut self) -> i64 { - let high = (self.next(32) as i64) << 32; - let low = (self.next(32) as i64) & 0xffffffffi64; - high | low - } - - /// Port of `BitsStreamGenerator.nextInt(int n)`. The caller always passes a - /// strictly positive `n`, matching Spark's `random.nextInt(i + 1)`. - fn next_int(&mut self, n: i32) -> i32 { - if (n & n.wrapping_neg()) == n { - // n is a power of two. - return ((n as i64 * self.next(31) as i64) >> 31) as i32; - } - loop { - let bits = self.next(31); - let val = bits % n; - if bits.wrapping_sub(val).wrapping_add(n.wrapping_sub(1)) >= 0 { - return val; - } - } - } - - /// Port of `RandomIndicesGenerator.getNextIndices`. Fills `out` with, for each - /// output position, the source position it should draw from. `out` is reused - /// across rows to avoid a per-row allocation. Advances the PRNG state, which is - /// shared across every row in the partition. - fn next_indices_into(&mut self, length: usize, out: &mut Vec) { - out.clear(); - out.resize(length, 0); - let mut i = 0usize; - while i < length { - let j = self.next_int((i + 1) as i32) as usize; - if j != i { - out[i] = out[j]; - } - out[j] = i; - i += 1; - } - } - - #[cfg(test)] - fn next_indices(&mut self, length: usize) -> Vec { - let mut out = Vec::new(); - self.next_indices_into(length, &mut out); - out + out[j] = i; + i += 1; } } @@ -318,7 +156,7 @@ fn gather_shuffled_indices( if is_null(row) { gathered.extend((start..start + length).map(|idx| idx as u64)); } else { - rng.next_indices_into(length, &mut scratch); + next_indices_into(rng, length, &mut scratch); gathered.extend(scratch.iter().map(|&local| (start + local) as u64)); } } @@ -407,7 +245,14 @@ mod tests { /// driving Spark's `RandomIndicesGenerator.getNextIndices`. fn indices(seed: i64, lengths: &[usize]) -> Vec> { let mut rng = SparkMersenneTwister::new(seed); - lengths.iter().map(|&len| rng.next_indices(len)).collect() + lengths + .iter() + .map(|&len| { + let mut out = Vec::new(); + next_indices_into(&mut rng, len, &mut out); + out + }) + .collect() } #[test] diff --git a/native/spark-expr/src/nondetermenistic_funcs/uuid.rs b/native/spark-expr/src/nondetermenistic_funcs/uuid.rs index fd33c12e392..18a99c084d7 100644 --- a/native/spark-expr/src/nondetermenistic_funcs/uuid.rs +++ b/native/spark-expr/src/nondetermenistic_funcs/uuid.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use crate::nondetermenistic_funcs::shuffle::SparkMersenneTwister; -use arrow::array::{RecordBatch, StringArray}; +use crate::nondetermenistic_funcs::internal::mersenne::SparkMersenneTwister; +use arrow::array::{RecordBatch, StringBuilder}; use arrow::datatypes::{DataType, Schema}; use datafusion::common::Result; use datafusion::logical_expr::ColumnarValue; @@ -24,22 +24,16 @@ use datafusion::physical_expr::PhysicalExpr; use std::fmt::{Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::{Arc, Mutex}; +use uuid::Uuid; -/// Draw one RFC 4122 version 4 UUID string from the generator, matching +/// Draw one RFC 4122 version 4 UUID from the generator, matching /// `org.apache.spark.sql.catalyst.util.RandomUUIDGenerator.getNextUUID`: two /// `nextLong()` draws with the version (4) and variant (10) bits masked in. -fn next_uuid_string(rng: &mut SparkMersenneTwister) -> String { +/// `Uuid`'s canonical lowercase hyphenated form matches `java.util.UUID.toString()`. +fn next_uuid(rng: &mut SparkMersenneTwister) -> Uuid { let most = (rng.next_long() as u64 & 0xFFFF_FFFF_FFFF_0FFF) | 0x0000_0000_0000_4000; let least = (rng.next_long() as u64 | 0x8000_0000_0000_0000) & 0xBFFF_FFFF_FFFF_FFFF; - // Matches `java.util.UUID.toString()`: lowercase, zero-padded, hyphenated. - format!( - "{:08x}-{:04x}-{:04x}-{:04x}-{:012x}", - (most >> 32) & 0xFFFF_FFFF, - (most >> 16) & 0xFFFF, - most & 0xFFFF, - (least >> 48) & 0xFFFF, - least & 0xFFFF_FFFF_FFFF, - ) + Uuid::from_u64_pair(most, least) } /// Physical expression for Spark's `uuid()`. Like `ShuffleExpr`, the generator @@ -98,8 +92,15 @@ impl PhysicalExpr for UuidExpr { let mut state = self.state_holder.lock().unwrap(); let rng = state.get_or_insert_with(|| SparkMersenneTwister::new(self.seed)); - let result: StringArray = (0..num_rows).map(|_| Some(next_uuid_string(rng))).collect(); - Ok(ColumnarValue::Array(Arc::new(result))) + // Each canonical UUID is exactly 36 bytes, so pre-size both builder buffers and encode + // into a reused stack buffer to avoid a per-row heap allocation. + const LEN: usize = uuid::fmt::Hyphenated::LENGTH; + let mut builder = StringBuilder::with_capacity(num_rows, num_rows * LEN); + let mut buf = [0u8; LEN]; + for _ in 0..num_rows { + builder.append_value(next_uuid(rng).hyphenated().encode_lower(&mut buf)); + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) } fn children(&self) -> Vec<&Arc> { @@ -121,7 +122,7 @@ impl PhysicalExpr for UuidExpr { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, RecordBatchOptions}; + use arrow::array::{Array, RecordBatchOptions, StringArray}; fn empty_batch(num_rows: usize) -> RecordBatch { RecordBatch::try_new_with_options( @@ -132,31 +133,29 @@ mod tests { .unwrap() } - fn eval_uuids(seed: i64, num_rows: usize) -> Vec { - let batch = empty_batch(num_rows); - let expr = UuidExpr::new(seed); - let result = expr.evaluate(&batch).unwrap().into_array(num_rows).unwrap(); - let arr = result.as_any().downcast_ref::().unwrap(); + fn collect_uuids(expr: &UuidExpr, batch: &RecordBatch) -> Vec { + let arr = expr + .evaluate(batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + let arr = arr.as_any().downcast_ref::().unwrap(); (0..arr.len()).map(|i| arr.value(i).to_string()).collect() } + fn eval_uuids(seed: i64, num_rows: usize) -> Vec { + collect_uuids(&UuidExpr::new(seed), &empty_batch(num_rows)) + } + #[test] - fn test_uuid_format_and_version() { + fn test_uuid_version_and_variant_bits() { + // The RNG and the version/variant masking are ours; the canonical string layout is + // guaranteed by the `uuid` crate. Assert the RFC 4122 v4 bits our masking sets, at their + // fixed positions in `xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx`. for uuid in eval_uuids(42, 20) { - // Canonical 8-4-4-4-12 lowercase hex. - assert_eq!(uuid.len(), 36); - let parts: Vec<&str> = uuid.split('-').collect(); - assert_eq!( - parts.iter().map(|p| p.len()).collect::>(), - vec![8, 4, 4, 4, 12] - ); - assert!(uuid - .chars() - .all(|c| c == '-' || c.is_ascii_hexdigit() && !c.is_ascii_uppercase())); - // Version 4: first nibble of the third group. - assert_eq!(parts[2].as_bytes()[0], b'4'); - // Variant 10xx: first nibble of the fourth group is 8, 9, a, or b. - assert!(matches!(parts[3].as_bytes()[0], b'8' | b'9' | b'a' | b'b')); + let bytes = uuid.as_bytes(); + assert_eq!(bytes[14], b'4'); // version nibble + assert!(matches!(bytes[19], b'8' | b'9' | b'a' | b'b')); // variant nibble } } @@ -175,10 +174,7 @@ mod tests { let expr = UuidExpr::new(7); let mut streamed = Vec::new(); for n in [3usize, 4usize] { - let batch = empty_batch(n); - let arr = expr.evaluate(&batch).unwrap().into_array(n).unwrap(); - let arr = arr.as_any().downcast_ref::().unwrap(); - streamed.extend((0..arr.len()).map(|i| arr.value(i).to_string())); + streamed.extend(collect_uuids(&expr, &empty_batch(n))); } assert_eq!(streamed, eval_uuids(7, 7)); // All distinct. diff --git a/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql b/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql index afb5bc7d664..612afff46e9 100644 --- a/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql +++ b/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql @@ -30,14 +30,7 @@ INSERT INTO test_uuid VALUES (1), (2), (3), (4), (5) query SELECT length(uuid()) FROM test_uuid --- matches the RFC 4122 version 4 layout (version nibble 4, variant nibble 8/9/a/b), lowercase hex +-- matches the RFC 4122 version 4 layout (version nibble 4, variant nibble 8/9/a/b), lowercase hex. +-- This regex subsumes the length, alphabet, version, and variant checks. query SELECT uuid() RLIKE '^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$' FROM test_uuid - --- version nibble is always '4' -query -SELECT substring(uuid(), 15, 1) FROM test_uuid - --- variant nibble is one of 8, 9, a, b -query -SELECT substring(uuid(), 20, 1) RLIKE '^[89ab]$' FROM test_uuid From 7ef6420c6a0036e96014b1b43dc5ba982866ff7c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 27 Jul 2026 08:14:26 -0600 Subject: [PATCH 3/5] test: address review feedback on uuid tests Fills the five coverage gaps flagged on #5034: - Multi-partition bit-for-bit: uuid(42) FROM ... DISTRIBUTE BY id exercises partitionIndex != 0, which a single-partition test cannot catch (both engines would silently agree even if partitionIndex were ignored). - Aliased projections: uuid(0) = uuid(0) (seeded) and uuid() = uuid() (unseeded, spark_answer_only) pin the stateful-alias semantics that freshCopyIfContainsStatefulExpression relies on. - Empty batch: a Rust test that evaluate() on a zero-row batch does not advance RNG state, guarding against a stray next_uuid() outside the row loop. - Golden fixture: five seeds x five UUIDs captured from Commons Math3's MersenneTwister (the exact RNG behind Spark's RandomUUIDGenerator) lock all 128 bits and cover negative / MIN / MAX seeds; catches next_long regressions that shuffle's next_int tests would miss. - No-scan path: length(uuid(0)) pins the OneRowRelation shape with a deterministic assertion alongside the existing SELECT uuid(0). --- .../src/nondetermenistic_funcs/uuid.rs | 95 +++++++++++++++++++ .../sql-tests/expressions/misc/uuid.sql | 7 ++ .../expressions/misc/uuid_with_seed.sql | 24 ++++- 3 files changed, 125 insertions(+), 1 deletion(-) diff --git a/native/spark-expr/src/nondetermenistic_funcs/uuid.rs b/native/spark-expr/src/nondetermenistic_funcs/uuid.rs index 18a99c084d7..bfb7040e036 100644 --- a/native/spark-expr/src/nondetermenistic_funcs/uuid.rs +++ b/native/spark-expr/src/nondetermenistic_funcs/uuid.rs @@ -183,4 +183,99 @@ mod tests { sorted.dedup(); assert_eq!(sorted.len(), streamed.len()); } + + #[test] + fn test_uuid_empty_batch_does_not_advance_state() { + // A zero-row batch (e.g. `LIMIT 0` over a `uuid` projection, or a fully filtered + // partition) must not consume any RNG draws. If the code ever grew a `next_uuid()` + // call outside the `0..num_rows` loop, the state would drift and downstream + // bit-for-bit tests would fail intermittently. + let expr = UuidExpr::new(7); + let _ = collect_uuids(&expr, &empty_batch(0)); + let after_empty = collect_uuids(&expr, &empty_batch(3)); + assert_eq!(after_empty, eval_uuids(7, 3)); + } + + /// Golden fixtures captured from Commons Math3's `MersenneTwister` (the exact RNG that + /// backs Spark's `org.apache.spark.sql.catalyst.util.RandomUUIDGenerator`). Regenerate + /// with a tiny Java program: + /// + /// ```java + /// import org.apache.commons.math3.random.MersenneTwister; + /// import java.util.UUID; + /// static UUID next(MersenneTwister r) { + /// long m = (r.nextLong() & 0xFFFFFFFFFFFF0FFFL) | 0x0000000000004000L; + /// long l = (r.nextLong() | 0x8000000000000000L) & 0xBFFFFFFFFFFFFFFFL; + /// return new UUID(m, l); + /// } + /// ``` + /// + /// `next_long()` is used *only* by `uuid`, so the shuffle tests (which exercise + /// `next_int`) do not guard against a `next_long` regression such as swapping the + /// high/low words or masking with the wrong constant. This test locks all 128 bits + /// per row, catches sign-extension bugs in `set_seed_long` via the negative seeds, + /// and runs entirely in Rust so it fires without a JVM roundtrip. + #[test] + fn test_uuid_matches_commons_math3_random_uuid_generator() { + let cases: &[(i64, &[&str])] = &[ + ( + 0, + &[ + "269567e9-5d09-4af5-b20f-16851fc4a81a", + "2a52c3b8-890f-4aae-9607-3331ab0d4f01", + "46beb52b-622b-4226-bb12-4f40c6cdba04", + "24fdcf77-ec9f-4ac6-a096-b36d3b9fe378", + "cd614b40-85b9-4227-be28-64d85fcfbb24", + ], + ), + ( + -1, + &[ + "05965e7e-3faf-4328-968d-6a409e667b13", + "36438051-74c0-4a52-9d78-e5c31853117e", + "eb6900f4-e69b-4bc6-9efc-c0ebc49b2c2e", + "bb6bb8b3-b20f-49e8-be16-41ab0dea1a19", + "bcb43f50-314a-4fdc-badf-4da2884aa53c", + ], + ), + ( + i64::MIN, + &[ + "444b2d95-8a54-40b5-950e-9ef88de5a4fd", + "e3ff5c8a-242c-467f-b1d3-ca6b6acb483d", + "4d07fa9a-6df6-43c2-9523-a5fc295411cd", + "4f9549c4-92ed-4e93-9fa4-412e5e6740ed", + "88776d39-8d3b-4942-9eb2-25821ac806a7", + ], + ), + ( + i64::MAX, + &[ + "596c0db4-d77a-4abb-9c91-5e7f323f02b0", + "98e2f79d-4b33-4ca4-b094-4ff3c2ee7d7b", + "39bae786-bdb1-4092-88eb-3aa52b52c0e8", + "d00c376b-87e7-4c8d-85d8-f7053996d6e6", + "ae65bca1-3143-4dc4-9a54-56d786f99699", + ], + ), + ( + 42, + &[ + "6f155395-c8b9-436b-a39c-d247226bc2b2", + "92da2253-2186-4525-a440-17ad3083a275", + "24ba1728-e648-487f-8c20-446d4411d15c", + "1188645d-be7b-4681-9936-8b24d7c51835", + "b4b1bfea-867b-4c0a-9f8f-bfe01cc571c0", + ], + ), + ]; + for (seed, expected) in cases { + let actual = eval_uuids(*seed, expected.len()); + assert_eq!( + actual, + expected.iter().map(|s| s.to_string()).collect::>(), + "mismatch for seed {seed}" + ); + } + } } diff --git a/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql b/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql index 612afff46e9..7bb3264f1e2 100644 --- a/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql +++ b/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql @@ -34,3 +34,10 @@ SELECT length(uuid()) FROM test_uuid -- This regex subsumes the length, alphabet, version, and variant checks. query SELECT uuid() RLIKE '^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$' FROM test_uuid + +-- Two unseeded uuid() nodes get distinct fresh random seeds during analysis (Uuid is stateful, +-- freshCopyIfContainsStatefulExpression rewrites aliases). Every row must therefore compare +-- unequal -- if Comet ever hoisted a single instance, this would flip to true. +-- spark_answer_only: value is random, so we only cross-check that Spark and Comet agree. +query spark_answer_only +SELECT uuid() = uuid() FROM test_uuid diff --git a/spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql b/spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql index 6d569be4e72..de9562ea604 100644 --- a/spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql +++ b/spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql @@ -29,14 +29,36 @@ CREATE TABLE test_uuid_seed(id int) USING parquet statement INSERT INTO test_uuid_seed VALUES (1), (2), (3), (4), (5) --- fixed seed, single row +-- Multi-partition table used by the DISTRIBUTE BY query below to exercise partitionIndex != 0. +statement +CREATE TABLE test_uuid_parts(id int) USING parquet + +statement +INSERT INTO test_uuid_parts SELECT id FROM range(0, 32) + +-- fixed seed, single row -- also exercises the no-scan (OneRowRelation) planning path query SELECT uuid(0) +-- pin the no-scan path with a deterministic assertion on top of the raw value above +query +SELECT length(uuid(0)) + -- fixed seed, multiple rows: the generator advances per row within the partition query SELECT uuid(42) FROM test_uuid_seed +-- Forces a hash exchange so uuid runs post-shuffle across several partitions. Locks the +-- `seed + partitionIndex` offset used by RandomUUIDGenerator; a single-partition test would +-- silently agree with Spark even if Comet ignored partitionIndex. +query +SELECT uuid(42) FROM test_uuid_parts DISTRIBUTE BY id + +-- Aliased seeded projections: both nodes carry the same explicit seed, so freshCopyIfContainsStatefulExpression +-- must not accidentally hoist a shared instance -- each row must satisfy uuid(0) = uuid(0). +query +SELECT uuid(0) = uuid(0) FROM test_uuid_seed + -- zero seed over multiple rows query SELECT uuid(0) FROM test_uuid_seed From 2f6a5839e33e4796ef96ca152523b54dc8ae0742 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 31 Jul 2026 14:11:33 -0600 Subject: [PATCH 4/5] test: make uuid partition and seed coverage non-vacuous The DISTRIBUTE BY query never exercised partitionIndex != 0. `... FROM t DISTRIBUTE BY id` parses as RepartitionByExpression on top of the Project, so uuid ran on the scan's partitions, and five tiny files pack into a single FilePartition. Moved the projection above the exchange: SELECT uuid(42) FROM (SELECT id FROM test_uuid_parts DISTRIBUTE BY id) Add CometUuidExpressionSuite, which builds `Uuid(Some(seed))` directly through the existing version-shimmed `getColumnFromExpression`. This gives 3.4 and 3.5 a real cross-engine assertion for the first place: the SQL `uuid(seed)` form is 4.0+, so uuid_with_seed.sql is skipped there and the Rust golden constants were the only guard, with no way to confirm they came from Commons Math3 rather than from the Rust code they are meant to test. One test covers five seeds including negative/MIN/MAX; the other repartitions and asserts getNumPartitions > 1 before comparing, so the partition-offset coverage is checked rather than assumed. Verified by mutation: dropping `expr.seed.wrapping_add(planner.partition())` in UuidBuilder::build fails both new Scala tests. Restored, all pass on spark-3.5 and spark-4.1. Also: - Drop `spark_answer_only` from `uuid() = uuid()` in uuid.sql. The values are random but the comparison is not -- distinct fresh seeds mean false on every row -- so default query mode keeps the value check and adds the nativeness assertion that the guard depends on. - Record the real reason typeof(uuid(0)) cannot be used: Spark's TypeOf.doGenCode interpolates `child.dataType.catalogString` unquoted into generated Java, emitting `UTF8String.fromString(string)`. That code is byte-identical on 3.5, 4.0 and master, so it is not a 4.1 quirk as previously recorded. It is normally hidden because TypeOf.foldable is true and ConstantFolding removes the node; this suite excludes ConstantFolding, which exposes it. Captured as a harness constraint: no SQL fixture here can use typeof on a non-foldable input. - Note on the Rust golden test that CometUuidExpressionSuite is now its backstop, and that regenerating the constants from the Rust side would make it circular. No partition-count assertion is added to the SQL fixture: that suite compares Comet against Spark, so if the plan collapsed to one partition both engines would collapse identically and any such check would still pass. It is only assertable from Scala, which is where it now lives. --- .../src/nondetermenistic_funcs/uuid.rs | 9 +++ .../sql-tests/expressions/misc/uuid.sql | 8 ++- .../expressions/misc/uuid_with_seed.sql | 27 ++++++-- .../comet/CometUuidExpressionSuite.scala | 65 +++++++++++++++++++ 4 files changed, 101 insertions(+), 8 deletions(-) create mode 100644 spark/src/test/scala/org/apache/comet/CometUuidExpressionSuite.scala diff --git a/native/spark-expr/src/nondetermenistic_funcs/uuid.rs b/native/spark-expr/src/nondetermenistic_funcs/uuid.rs index bfb7040e036..c2f9944b0a7 100644 --- a/native/spark-expr/src/nondetermenistic_funcs/uuid.rs +++ b/native/spark-expr/src/nondetermenistic_funcs/uuid.rs @@ -215,6 +215,15 @@ mod tests { /// high/low words or masking with the wrong constant. This test locks all 128 bits /// per row, catches sign-extension bugs in `set_seed_long` via the negative seeds, /// and runs entirely in Rust so it fires without a JVM roundtrip. + /// + /// These constants are not the only line of defense, and must not become it. They were + /// hand-captured, so their provenance cannot be checked from the repo, and regenerating them + /// from this Rust code instead of from Commons Math3 would make the test circular and silently + /// stop it testing anything. The backstop is `CometUuidExpressionSuite`, which builds + /// `Uuid(Some(seed))` directly and compares Comet against Spark bit for bit on *every* + /// supported profile -- including 3.4 and 3.5, where the SQL `uuid(seed)` form does not exist + /// and `uuid_with_seed.sql` is skipped. If these values ever disagree with that suite, trust + /// the suite. #[test] fn test_uuid_matches_commons_math3_random_uuid_generator() { let cases: &[(i64, &[&str])] = &[ diff --git a/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql b/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql index 7bb3264f1e2..3effbdb50c9 100644 --- a/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql +++ b/spark/src/test/resources/sql-tests/expressions/misc/uuid.sql @@ -37,7 +37,9 @@ SELECT uuid() RLIKE '^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0- -- Two unseeded uuid() nodes get distinct fresh random seeds during analysis (Uuid is stateful, -- freshCopyIfContainsStatefulExpression rewrites aliases). Every row must therefore compare --- unequal -- if Comet ever hoisted a single instance, this would flip to true. --- spark_answer_only: value is random, so we only cross-check that Spark and Comet agree. -query spark_answer_only +-- unequal -- if Comet ever hoisted a single instance, this would flip to true. The values are +-- random but the comparison is not, so this runs in default query mode: that keeps the value +-- check and also asserts the projection actually ran natively, which is the precondition for +-- this guard to mean anything. +query SELECT uuid() = uuid() FROM test_uuid diff --git a/spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql b/spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql index de9562ea604..260e729f43a 100644 --- a/spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql +++ b/spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql @@ -40,7 +40,15 @@ INSERT INTO test_uuid_parts SELECT id FROM range(0, 32) query SELECT uuid(0) --- pin the no-scan path with a deterministic assertion on top of the raw value above +-- pin the no-scan path with a deterministic assertion on top of the raw value above. +-- typeof(uuid(0)) would read better here but cannot be used: Spark's TypeOf.doGenCode +-- interpolates the catalog string unquoted into generated Java -- +-- `UTF8String.fromString(${child.dataType.catalogString})` -- which emits +-- `UTF8String.fromString(string)` and fails to compile. That code is byte-identical on 3.5, 4.0 +-- and master, so it is not version specific; it is normally invisible because TypeOf.foldable is +-- true and ConstantFolding removes the node before codegen. This suite excludes ConstantFolding +-- (see CometSqlFileTestSuite), which is what exposes it. So no SQL fixture here can use typeof on +-- a non-foldable input. query SELECT length(uuid(0)) @@ -48,11 +56,20 @@ SELECT length(uuid(0)) query SELECT uuid(42) FROM test_uuid_seed --- Forces a hash exchange so uuid runs post-shuffle across several partitions. Locks the --- `seed + partitionIndex` offset used by RandomUUIDGenerator; a single-partition test would --- silently agree with Spark even if Comet ignored partitionIndex. +-- Runs the projection *above* a hash exchange so uuid is evaluated on post-shuffle partitions and +-- exercises the `seed + partitionIndex` offset in RandomUUIDGenerator. The projection must sit on +-- top of the exchange: `... FROM t DISTRIBUTE BY id` parses as RepartitionByExpression on top of +-- the Project, so uuid would run on the scan's partitions instead, and those pack into a single +-- FilePartition here -- partitionIndex would be 0 and the test would agree with Spark even if the +-- offset were dropped entirely. query -SELECT uuid(42) FROM test_uuid_parts DISTRIBUTE BY id +SELECT uuid(42) FROM (SELECT id FROM test_uuid_parts DISTRIBUTE BY id) + +-- The value of the query above rests on there being more than one post-shuffle partition, which +-- cannot be asserted here: this suite compares Comet against Spark, so if the plan ever collapsed +-- to one partition both engines would collapse identically and any check would still pass. The +-- partition count is asserted directly in CometUuidExpressionSuite ("across multiple partitions"), +-- which checks getNumPartitions > 1 before comparing, and covers 3.4/3.5 as well. -- Aliased seeded projections: both nodes carry the same explicit seed, so freshCopyIfContainsStatefulExpression -- must not accidentally hoist a shared instance -- each row must satisfy uuid(0) = uuid(0). diff --git a/spark/src/test/scala/org/apache/comet/CometUuidExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometUuidExpressionSuite.scala new file mode 100644 index 00000000000..56029cba282 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometUuidExpressionSuite.scala @@ -0,0 +1,65 @@ +/* + * 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. + */ + +package org.apache.comet + +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.catalyst.expressions.Uuid +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.functions.col + +/** + * Bit-for-bit comparisons of seeded `uuid` against Spark. + * + * The SQL `uuid(seed)` form only exists in Spark 4.0+, so `uuid_with_seed.sql` carries a + * `MinSparkVersion: 4.0` marker and is skipped on 3.4 and 3.5. Unseeded `uuid()` cannot be + * compared across engines. That left the Rust golden constants in + * `nondetermenistic_funcs/uuid.rs` as the only guard on those profiles, and their provenance is + * not checkable from the repo -- if they were ever regenerated from the Rust side the test would + * become circular. + * + * `Uuid(Some(seed))` is constructible from Scala on every supported version even where the SQL + * form is not, so building the expression directly gives all profiles a real cross-engine + * assertion and removes that dependence. + */ +class CometUuidExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + test("seeded uuid matches Spark bit for bit") { + withParquetTable((0 until 20).map(i => (i, i.toString)), "tbl") { + Seq(0L, -1L, 42L, Long.MinValue, Long.MaxValue).foreach { seed => + val df = spark.table("tbl").select(getColumnFromExpression(Uuid(Some(seed)))) + checkSparkAnswerAndOperator(df) + } + } + } + + test("seeded uuid matches Spark bit for bit across multiple partitions") { + // Exercises the `seed + partitionIndex` offset that RandomUUIDGenerator applies. A + // single-partition run agrees with Spark even if the offset were dropped entirely, so + // repartition first and assert the partition count rather than assuming it. + withParquetTable((0 until 64).map(i => (i, i.toString)), "tbl") { + val repartitioned = spark.table("tbl").repartition(4, col("_1")) + assert(repartitioned.rdd.getNumPartitions > 1) + Seq(0L, 42L).foreach { seed => + checkSparkAnswerAndOperator( + repartitioned.select(getColumnFromExpression(Uuid(Some(seed))))) + } + } + } +} From b0ca636660e78cc7d763c2cf40a181e4912acdae Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 31 Jul 2026 14:51:51 -0600 Subject: [PATCH 5/5] ci: register CometUuidExpressionSuite in PR build workflows The check-suites.py preflight requires every *Suite.scala to be listed in both pr_build_linux.yml and pr_build_macos.yml. Add the new uuid suite to the expressions bucket. --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index d5476710955..ff69a2cc68d 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -401,6 +401,7 @@ jobs: org.apache.comet.CometCodegenFuzzSuite org.apache.comet.CometStringDecodeSuite org.apache.comet.CometWidthBucketSuite + org.apache.comet.CometUuidExpressionSuite fail-fast: false name: ${{ matrix.profile.name }} [${{ matrix.suite.name }}] runs-on: ubuntu-24.04 diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 1183eb0a82b..f3f3f082695 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -217,6 +217,7 @@ jobs: org.apache.comet.CometCodegenFuzzSuite org.apache.comet.CometStringDecodeSuite org.apache.comet.CometWidthBucketSuite + org.apache.comet.CometUuidExpressionSuite fail-fast: false name: ${{ matrix.os }}/${{ matrix.profile.name }} [${{ matrix.suite.name }}]