andygrove commented on code in PR #5034: URL: https://github.com/apache/datafusion-comet/pull/5034#discussion_r3693106096
########## spark/src/test/resources/sql-tests/expressions/misc/uuid_with_seed.sql: ########## @@ -0,0 +1,79 @@ +-- 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) + +-- 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)) Review Comment: You are right, and I checked the mechanism rather than just taking the correction: `TypeOf.foldable` is true, so ConstantFolding normally removes the node before codegen ever sees it, and `CometSqlFileTestSuite` sets `spark.sql.optimizer.excludedRules` to `ConstantFolding` (lines 91-92), which is exactly what exposes the latent bug. So the version-specific framing in my earlier comment was wrong on both counts — it is not 4.1 and it is not uuid. Recorded in 2f6a5839e as a comment on the query, framed as the harness-level constraint you asked for: no SQL fixture in this suite can use `typeof` on a non-foldable input, with the unquoted-`catalogString` cause spelled out so the next person does not rediscover it through the Janino error. I will correct the PR description too. On the upstream report: I have not filed it, since that is an outward-facing action on someone else's tracker and I would rather confirm before opening a Spark ticket. Happy to file it if you want, or leave it to you — the repro is `SELECT typeof(<non-foldable>)` with ConstantFolding excluded, on any of 3.5, 4.0 or master. ########## native/spark-expr/src/nondetermenistic_funcs/uuid.rs: ########## @@ -0,0 +1,281 @@ +// 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::internal::mersenne::SparkMersenneTwister; +use arrow::array::{RecordBatch, StringBuilder}; +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}; +use uuid::Uuid; + +/// 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. +/// `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; + Uuid::from_u64_pair(most, least) +} + +/// 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<Mutex<Option<SparkMersenneTwister>>>, +} + +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<H: Hasher>(&self, state: &mut H) { + self.seed.hash(state); + } +} + +impl PhysicalExpr for UuidExpr { + fn data_type(&self, _input_schema: &Schema) -> Result<DataType> { + Ok(DataType::Utf8) + } + + fn nullable(&self, _input_schema: &Schema) -> Result<bool> { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { + 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)); + + // 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<dyn PhysicalExpr>> { + vec![] + } + + fn with_new_children( + self: Arc<Self>, + _children: Vec<Arc<dyn PhysicalExpr>>, + ) -> Result<Arc<dyn PhysicalExpr>> { + 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, StringArray}; + + 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 collect_uuids(expr: &UuidExpr, batch: &RecordBatch) -> Vec<String> { + let arr = expr + .evaluate(batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + let arr = arr.as_any().downcast_ref::<StringArray>().unwrap(); + (0..arr.len()).map(|i| arr.value(i).to_string()).collect() + } + + fn eval_uuids(seed: i64, num_rows: usize) -> Vec<String> { + collect_uuids(&UuidExpr::new(seed), &empty_batch(num_rows)) + } + + #[test] + 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) { + 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 + } + } + + #[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] { + streamed.extend(collect_uuids(&expr, &empty_batch(n))); + } + assert_eq!(streamed, eval_uuids(7, 7)); + // All distinct. + let mut sorted = streamed.clone(); + sorted.sort(); + 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() { Review Comment: This was the most useful comment on the PR — the circularity risk is real and I had not protected against it. Fixed in 2f6a5839e. Added `CometUuidExpressionSuite`, which builds `Uuid(Some(seed))` directly via the existing version-shimmed `getColumnFromExpression` in `ShimCometTestBase` (`new Column(expr)` on 3.x, `ExpressionColumnNode` on 4.x), so it compiles and runs on every profile. Two tests: one over seeds 0, -1, 42, `Long.MinValue` and `Long.MaxValue`, and one that repartitions and asserts `getNumPartitions > 1` before comparing. Both use `checkSparkAnswerAndOperator`, so they assert bit-for-bit equality with Spark *and* that the projection stayed native. That gives 3.4 and 3.5 a real cross-engine assertion for the first time, so the golden constants are no longer the only line of defense on those profiles. Confirmed non-vacuous by mutation: dropping the partition offset in `UuidBuilder::build` fails **both** new tests (the single-partition one too — `withParquetTable` over 20 rows under `local[5]` already gives more than one partition). Restored, and all pass on spark-3.5 and spark-4.1. I also left a note on the Rust golden test pointing at the new suite as its backstop and stating explicitly that regenerating these constants from the Rust side would make the test circular — so the trap is documented where someone would fall into it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
