This is an automated email from the ASF dual-hosted git repository.
ShreyeshArangath pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git
The following commit(s) were added to refs/heads/master by this push:
new 003f0d95 [AURON #1891] Implement randn() function (#1938)
003f0d95 is described below
commit 003f0d955c07cd42168bf2112a7e684b0a1b68e7
Author: Rob Reeves <[email protected]>
AuthorDate: Thu Jul 23 14:27:57 2026 -0700
[AURON #1891] Implement randn() function (#1938)
<!--
- Start the PR title with the related issue ID, e.g. '[AURON #XXXX]
Short summary...'.
-->
# Which issue does this PR close?
Closes #1891
# Rationale for this change
This improves function coverage in Auron by creating a native randn
implementation.
# What changes are included in this PR?
Adds a native randn implementation.
# Are there any user-facing changes?
Yes, it adds the randn function.
# How was this patch tested?
Added unit tests and manually tested in spark-shell.
```scala
import org.apache.spark.sql.functions.randn
val df = spark.range(5)
val outputPath = "/tmp/spark_range_output.parquet"
df.write.mode("overwrite").parquet(outputPath)
val readDf = spark.read.parquet(outputPath)
val resultDf = readDf.withColumn("random_normal", randn(18))
resultDf.collect
```
Output:
```
26/01/30 15:41:22 WARN NativeHelper: memory total: 1408.0 MiB, onheap:
1024.0 MiB, offheap: 384.0 MiB
26/01/30 15:41:24 WARN AuronCallNativeWrapper: Start executing native plan
26/01/30 15:41:24 WARN AuronCallNativeWrapper: Start executing native plan
26/01/30 15:41:24 WARN AuronCallNativeWrapper: Start executing native plan
26/01/30 15:41:24 WARN AuronCallNativeWrapper: Start executing native plan
26/01/30 15:41:24 WARN AuronCallNativeWrapper: Start executing native plan
26/01/30 15:41:24 WARN AuronCallNativeWrapper: Start executing native plan
------ initializing auron native environment ------
initializing logging with level: info
2026-01-30 15:41:24.368 (+0.000s) [INFO] [auron::exec:73] (stage: 0,
partition: 0, tid: 0) - initializing JNI bridge
2026-01-30 15:41:24.369 (+0.001s) [INFO] [auron_jni_bridge::jni_bridge:473]
(stage: 0, partition: 0, tid: 0) - Initializing JavaClasses...
2026-01-30 15:41:24.375 (+0.007s) [INFO] [auron_jni_bridge::jni_bridge:529]
(stage: 0, partition: 0, tid: 0) - Initializing JavaClasses finished
2026-01-30 15:41:24.375 (+0.007s) [INFO] [auron::exec:77] (stage: 0,
partition: 0, tid: 0) - initializing datafusion session
2026-01-30 15:41:24.375 (+0.007s) [INFO] [auron_memmgr:48] (stage: 0,
partition: 0, tid: 0) - mem manager initialized with total memory: 230.4 MiB
2026-01-30 15:41:24.385 (+0.017s) [INFO] [auron::rt:146] (stage: 2,
partition: 1, tid: 12) - start executing plan:
ProjectExec [#3@0 AS #3, Randn(seed=18, partition=1) AS #5],
schema=[#3:Int64;N, #5:Float64]
RenameColumnsExec: ["#3"], schema=[#3:Int64;N]
ParquetExec: limit=None, file_group=[FileGroup { files: [], statistics:
None }, FileGroup { files: [PartitionedFile { object_meta: ObjectMeta {
location: Path { raw:
"ZmlsZTovLy90bXAvc3BhcmtfcmFuZ2Vfb3V0cHV0LnBhcnF1ZXQvcGFydC0wMDAwMS04ZTkwNmRiYS0zZDg3LTRkZWMtYjM0NC1hYjdiZWUyODEwZWQtYzAwMC5zbmFwcHkucGFycXVldA"
}, last_modified: 1970-01-01T00:00:00Z, size: 472, e_tag: None, version: None
}, partition_values: [], range: Some(FileRange { start: 0, end: 472 }),
statistics: None, extens [...]
2026-01-30 15:41:24.385 (+0.017s) [INFO] [auron::rt:146] (stage: 2,
partition: 5, tid: 16) - start executing plan:
ProjectExec [#3@0 AS #3, Randn(seed=18, partition=5) AS #5],
schema=[#3:Int64;N, #5:Float64]
RenameColumnsExec: ["#3"], schema=[#3:Int64;N]
ParquetExec: limit=None, file_group=[FileGroup { files: [], statistics:
None }, FileGroup { files: [], statistics: None }, FileGroup { files: [],
statistics: None }, FileGroup { files: [], statistics: None }, FileGroup {
files: [], statistics: None }, FileGroup { files: [PartitionedFile {
object_meta: ObjectMeta { location: Path { raw:
"ZmlsZTovLy90bXAvc3BhcmtfcmFuZ2Vfb3V0cHV0LnBhcnF1ZXQvcGFydC0wMDAwMC04ZTkwNmRiYS0zZDg3LTRkZWMtYjM0NC1hYjdiZWUyODEwZWQtYzAwMC5zbmFwcHkucGFycXVldA"
}, [...]
2026-01-30 15:41:24.385 (+0.017s) [INFO] [auron::rt:146] (stage: 2,
partition: 2, tid: 13) - start executing plan:
ProjectExec [#3@0 AS #3, Randn(seed=18, partition=2) AS #5],
schema=[#3:Int64;N, #5:Float64]
RenameColumnsExec: ["#3"], schema=[#3:Int64;N]
ParquetExec: limit=None, file_group=[FileGroup { files: [], statistics:
None }, FileGroup { files: [], statistics: None }, FileGroup { files:
[PartitionedFile { object_meta: ObjectMeta { location: Path { raw:
"ZmlsZTovLy90bXAvc3BhcmtfcmFuZ2Vfb3V0cHV0LnBhcnF1ZXQvcGFydC0wMDAwMy04ZTkwNmRiYS0zZDg3LTRkZWMtYjM0NC1hYjdiZWUyODEwZWQtYzAwMC5zbmFwcHkucGFycXVldA"
}, last_modified: 1970-01-01T00:00:00Z, size: 472, e_tag: None, version: None
}, partition_values: [], range: Some(FileRange { star [...]
2026-01-30 15:41:24.385 (+0.017s) [INFO] [auron::rt:146] (stage: 2,
partition: 4, tid: 15) - start executing plan:
ProjectExec [#3@0 AS #3, Randn(seed=18, partition=4) AS #5],
schema=[#3:Int64;N, #5:Float64]
RenameColumnsExec: ["#3"], schema=[#3:Int64;N]
ParquetExec: limit=None, file_group=[FileGroup { files: [], statistics:
None }, FileGroup { files: [], statistics: None }, FileGroup { files: [],
statistics: None }, FileGroup { files: [], statistics: None }, FileGroup {
files: [PartitionedFile { object_meta: ObjectMeta { location: Path { raw:
"ZmlsZTovLy90bXAvc3BhcmtfcmFuZ2Vfb3V0cHV0LnBhcnF1ZXQvcGFydC0wMDAwNS04ZTkwNmRiYS0zZDg3LTRkZWMtYjM0NC1hYjdiZWUyODEwZWQtYzAwMC5zbmFwcHkucGFycXVldA"
}, last_modified: 1970-01-01T00:00:00Z, size: [...]
2026-01-30 15:41:24.385 (+0.017s) [INFO] [auron::rt:146] (stage: 2,
partition: 3, tid: 14) - start executing plan:
ProjectExec [#3@0 AS #3, Randn(seed=18, partition=3) AS #5],
schema=[#3:Int64;N, #5:Float64]
RenameColumnsExec: ["#3"], schema=[#3:Int64;N]
ParquetExec: limit=None, file_group=[FileGroup { files: [], statistics:
None }, FileGroup { files: [], statistics: None }, FileGroup { files: [],
statistics: None }, FileGroup { files: [PartitionedFile { object_meta:
ObjectMeta { location: Path { raw:
"ZmlsZTovLy90bXAvc3BhcmtfcmFuZ2Vfb3V0cHV0LnBhcnF1ZXQvcGFydC0wMDAwOS04ZTkwNmRiYS0zZDg3LTRkZWMtYjM0NC1hYjdiZWUyODEwZWQtYzAwMC5zbmFwcHkucGFycXVldA"
}, last_modified: 1970-01-01T00:00:00Z, size: 472, e_tag: None, version: None
}, partiti [...]
2026-01-30 15:41:24.385 (+0.017s) [INFO] [auron::rt:146] (stage: 2,
partition: 0, tid: 11) - start executing plan:
ProjectExec [#3@0 AS #3, Randn(seed=18, partition=0) AS #5],
schema=[#3:Int64;N, #5:Float64]
RenameColumnsExec: ["#3"], schema=[#3:Int64;N]
ParquetExec: limit=None, file_group=[FileGroup { files:
[PartitionedFile { object_meta: ObjectMeta { location: Path { raw:
"ZmlsZTovLy90bXAvc3BhcmtfcmFuZ2Vfb3V0cHV0LnBhcnF1ZXQvcGFydC0wMDAwNy04ZTkwNmRiYS0zZDg3LTRkZWMtYjM0NC1hYjdiZWUyODEwZWQtYzAwMC5zbmFwcHkucGFycXVldA"
}, last_modified: 1970-01-01T00:00:00Z, size: 472, e_tag: None, version: None
}, partition_values: [], range: Some(FileRange { start: 0, end: 472 }),
statistics: None, extensions: None, metadata_size_hint: None }], st [...]
2026-01-30 15:41:24.394 (+0.026s) [INFO]
[datafusion_datasource_parquet::opener:421] (stage: 2, partition: 4, tid: 15) -
executing parquet scan with adaptive batch size: 10000
2026-01-30 15:41:24.394 (+0.026s) [INFO]
[datafusion_datasource_parquet::opener:421] (stage: 2, partition: 0, tid: 11) -
executing parquet scan with adaptive batch size: 10000
2026-01-30 15:41:24.394 (+0.026s) [INFO]
[datafusion_datasource_parquet::opener:421] (stage: 2, partition: 3, tid: 14) -
executing parquet scan with adaptive batch size: 10000
2026-01-30 15:41:24.394 (+0.026s) [INFO]
[datafusion_datasource_parquet::opener:421] (stage: 2, partition: 1, tid: 12) -
executing parquet scan with adaptive batch size: 10000
2026-01-30 15:41:24.394 (+0.026s) [INFO]
[datafusion_datasource_parquet::opener:421] (stage: 2, partition: 2, tid: 13) -
executing parquet scan with adaptive batch size: 10000
2026-01-30 15:41:24.394 (+0.026s) [INFO]
[datafusion_datasource_parquet::opener:421] (stage: 2, partition: 5, tid: 16) -
executing parquet scan with adaptive batch size: 1
2026-01-30 15:41:24.488 (+0.120s) [INFO] [auron::rt:183] (stage: 2,
partition: 5, tid: 16) - task finished
2026-01-30 15:41:24.488 (+0.120s) [INFO] [auron::rt:183] (stage: 2,
partition: 0, tid: 11) - task finished
2026-01-30 15:41:24.488 (+0.120s) [INFO] [auron::rt:183] (stage: 2,
partition: 4, tid: 15) - task finished
2026-01-30 15:41:24.488 (+0.120s) [INFO] [auron::rt:183] (stage: 2,
partition: 3, tid: 14) - task finished
2026-01-30 15:41:24.488 (+0.120s) [INFO] [auron::rt:266] (stage: 0,
partition: 0, tid: 0) - (partition=5) native execution finalizing
2026-01-30 15:41:24.488 (+0.120s) [INFO] [auron::rt:183] (stage: 2,
partition: 2, tid: 13) - task finished
2026-01-30 15:41:24.488 (+0.120s) [INFO] [auron::rt:183] (stage: 2,
partition: 1, tid: 12) - task finished
2026-01-30 15:41:24.488 (+0.120s) [INFO] [auron::rt:274] (stage: 0,
partition: 0, tid: 0) - (partition=5) native execution finalized
2026-01-30 15:41:24.511 (+0.143s) [INFO] [auron::rt:266] (stage: 0,
partition: 0, tid: 0) - (partition=3) native execution finalizing
2026-01-30 15:41:24.511 (+0.143s) [INFO] [auron::rt:266] (stage: 0,
partition: 0, tid: 0) - (partition=4) native execution finalizing
2026-01-30 15:41:24.511 (+0.143s) [INFO] [auron::rt:266] (stage: 0,
partition: 0, tid: 0) - (partition=0) native execution finalizing
2026-01-30 15:41:24.511 (+0.143s) [INFO] [auron::rt:266] (stage: 0,
partition: 0, tid: 0) - (partition=2) native execution finalizing
2026-01-30 15:41:24.511 (+0.143s) [INFO] [auron::rt:266] (stage: 0,
partition: 0, tid: 0) - (partition=1) native execution finalizing
2026-01-30 15:41:24.512 (+0.144s) [INFO] [auron::rt:274] (stage: 0,
partition: 0, tid: 0) - (partition=0) native execution finalized
2026-01-30 15:41:24.512 (+0.144s) [INFO] [auron::rt:274] (stage: 0,
partition: 0, tid: 0) - (partition=4) native execution finalized
2026-01-30 15:41:24.512 (+0.144s) [INFO] [auron::rt:274] (stage: 0,
partition: 0, tid: 0) - (partition=1) native execution finalized
2026-01-30 15:41:24.512 (+0.144s) [INFO] [auron::rt:274] (stage: 0,
partition: 0, tid: 0) - (partition=2) native execution finalized
2026-01-30 15:41:24.512 (+0.144s) [INFO] [auron::rt:274] (stage: 0,
partition: 0, tid: 0) - (partition=3) native execution finalized
import org.apache.spark.sql.functions.randn
df: org.apache.spark.sql.Dataset[Long] = [id: bigint]
outputPath: String = /tmp/spark_range_output.parquet
readDf: org.apache.spark.sql.DataFrame = [id: bigint]
resultDf: org.apache.spark.sql.DataFrame = [id: bigint, random_normal:
double]
res0: Array[org.apache.spark.sql.Row] = Array([3,1.4607292672705405],
[0,-0.3268302897860617], [1,-0.09087682847007866], [4,-1.2271197538792842],
[2,-0.546398027932835])
```
---------
Co-authored-by: Claude Opus 4.5 <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>
---
Cargo.lock | 12 +
Cargo.toml | 1 +
native-engine/auron-planner/proto/auron.proto | 7 +
native-engine/auron-planner/src/planner.rs | 3 +-
native-engine/datafusion-ext-exprs/Cargo.toml | 2 +
native-engine/datafusion-ext-exprs/src/lib.rs | 1 +
.../datafusion-ext-exprs/src/spark_randn.rs | 305 +++++++++++++++++++++
.../org/apache/auron/AuronFunctionSuite.scala | 62 +++++
.../org/apache/spark/sql/AuronQueryTest.scala | 23 +-
.../apache/spark/sql/auron/NativeConverters.scala | 15 +
10 files changed, 421 insertions(+), 10 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index aa57e0d1..6d956cb3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1167,6 +1167,8 @@ dependencies = [
"log",
"once_cell",
"parking_lot",
+ "rand",
+ "rand_distr",
]
[[package]]
@@ -3351,6 +3353,16 @@ dependencies = [
"getrandom 0.3.3",
]
+[[package]]
+name = "rand_distr"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463"
+dependencies = [
+ "num-traits",
+ "rand",
+]
+
[[package]]
name = "rdkafka"
version = "0.36.2"
diff --git a/Cargo.toml b/Cargo.toml
index 208054c3..e5b80b91 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -185,6 +185,7 @@ prost = "0.14.4"
prost-types = "0.14.3"
prost-reflect = "0.16.4"
rand = "0.9.3"
+rand_distr = "0.5"
smallvec = "2.0.0-alpha.11"
sonic-rs = "0.5.8"
tempfile = "3"
diff --git a/native-engine/auron-planner/proto/auron.proto
b/native-engine/auron-planner/proto/auron.proto
index a905c8a3..74b0bebb 100644
--- a/native-engine/auron-planner/proto/auron.proto
+++ b/native-engine/auron-planner/proto/auron.proto
@@ -121,6 +121,9 @@ message PhysicalExprNode {
// MonotonicIncreasingID
MonotonicIncreasingIdExprNode monotonic_increasing_id_expr = 20102;
+ // SparkRandn
+ SparkRandnExprNode spark_randn_expr = 20103;
+
// BloomFilterMightContain
BloomFilterMightContainExprNode bloom_filter_might_contain_expr = 20200;
}
@@ -384,6 +387,10 @@ message SparkPartitionIdExprNode {
message MonotonicIncreasingIdExprNode {
}
+message SparkRandnExprNode {
+ int64 seed = 1;
+}
+
message BloomFilterMightContainExprNode {
string uuid = 1;
PhysicalExprNode bloom_filter_expr = 2;
diff --git a/native-engine/auron-planner/src/planner.rs
b/native-engine/auron-planner/src/planner.rs
index b1ee1584..b31e5363 100644
--- a/native-engine/auron-planner/src/planner.rs
+++ b/native-engine/auron-planner/src/planner.rs
@@ -54,7 +54,7 @@ use datafusion_ext_exprs::{
get_indexed_field::GetIndexedFieldExpr, get_map_value::GetMapValueExpr,
named_struct::NamedStructExpr, row_num::RowNumExpr,
spark_monotonically_increasing_id::SparkMonotonicallyIncreasingIdExpr,
- spark_partition_id::SparkPartitionIdExpr,
+ spark_partition_id::SparkPartitionIdExpr, spark_randn::SparkRandnExpr,
spark_scalar_subquery_wrapper::SparkScalarSubqueryWrapperExpr,
spark_udf_wrapper::SparkUDFWrapperExpr,
string_contains::StringContainsExpr,
string_ends_with::StringEndsWithExpr,
string_starts_with::StringStartsWithExpr,
@@ -1095,6 +1095,7 @@ impl PhysicalPlanner {
ExprType::MonotonicIncreasingIdExpr(_) => {
Arc::new(SparkMonotonicallyIncreasingIdExpr::new(self.partition_id))
}
+ ExprType::SparkRandnExpr(e) =>
Arc::new(SparkRandnExpr::new(e.seed, self.partition_id)),
ExprType::BloomFilterMightContainExpr(e) =>
Arc::new(BloomFilterMightContainExpr::new(
e.uuid.clone(),
self.try_parse_physical_expr_box_required(&e.bloom_filter_expr, input_schema)?,
diff --git a/native-engine/datafusion-ext-exprs/Cargo.toml
b/native-engine/datafusion-ext-exprs/Cargo.toml
index a38b71e4..96e77bf0 100644
--- a/native-engine/datafusion-ext-exprs/Cargo.toml
+++ b/native-engine/datafusion-ext-exprs/Cargo.toml
@@ -35,3 +35,5 @@ jni = { workspace = true }
log = { workspace = true }
once_cell = { workspace = true }
parking_lot = { workspace = true }
+rand = { workspace = true }
+rand_distr = { workspace = true }
diff --git a/native-engine/datafusion-ext-exprs/src/lib.rs
b/native-engine/datafusion-ext-exprs/src/lib.rs
index 6400f7d2..8c781ae6 100644
--- a/native-engine/datafusion-ext-exprs/src/lib.rs
+++ b/native-engine/datafusion-ext-exprs/src/lib.rs
@@ -25,6 +25,7 @@ pub mod named_struct;
pub mod row_num;
pub mod spark_monotonically_increasing_id;
pub mod spark_partition_id;
+pub mod spark_randn;
pub mod spark_scalar_subquery_wrapper;
pub mod spark_udf_wrapper;
pub mod string_contains;
diff --git a/native-engine/datafusion-ext-exprs/src/spark_randn.rs
b/native-engine/datafusion-ext-exprs/src/spark_randn.rs
new file mode 100644
index 00000000..9a947a33
--- /dev/null
+++ b/native-engine/datafusion-ext-exprs/src/spark_randn.rs
@@ -0,0 +1,305 @@
+// 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 std::{
+ any::Any,
+ fmt::{Debug, Display, Formatter},
+ hash::{Hash, Hasher},
+ sync::Arc,
+};
+
+use arrow::{
+ array::{Float64Array, RecordBatch},
+ datatypes::{DataType, Schema},
+};
+use datafusion::{
+ common::Result,
+ logical_expr::ColumnarValue,
+ physical_expr::{PhysicalExpr, PhysicalExprRef},
+};
+use parking_lot::Mutex;
+use rand::{SeedableRng, rngs::StdRng};
+use rand_distr::{Distribution, StandardNormal};
+
+use crate::down_cast_any_ref;
+
+/// Returns random values with independent and identically distributed (i.i.d.)
+/// samples drawn from the standard normal distribution.
+///
+/// Spark-compatible semantics:
+/// - RNG is seeded with `seed + partition_id`
+/// - RNG state advances for each row (stateful across batches)
+///
+/// Note: the underlying RNG/gaussian implementation is not intended to
+/// reproduce Spark's exact output sequence for a given seed/partition.
+pub struct SparkRandnExpr {
+ seed: i64,
+ partition_id: usize,
+ rng: Mutex<StdRng>,
+}
+
+impl SparkRandnExpr {
+ pub fn new(seed: i64, partition_id: usize) -> Self {
+ let effective_seed = (seed as u64).wrapping_add(partition_id as u64);
+ Self {
+ seed,
+ partition_id,
+ rng: Mutex::new(StdRng::seed_from_u64(effective_seed)),
+ }
+ }
+}
+
+impl Display for SparkRandnExpr {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(
+ f,
+ "Randn(seed={}, partition={})",
+ self.seed, self.partition_id
+ )
+ }
+}
+
+impl Debug for SparkRandnExpr {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(
+ f,
+ "Randn(seed={}, partition={})",
+ self.seed, self.partition_id
+ )
+ }
+}
+
+impl PartialEq for SparkRandnExpr {
+ fn eq(&self, other: &Self) -> bool {
+ self.seed == other.seed && self.partition_id == other.partition_id
+ }
+}
+
+impl Eq for SparkRandnExpr {}
+
+impl Hash for SparkRandnExpr {
+ fn hash<H: Hasher>(&self, state: &mut H) {
+ self.seed.hash(state);
+ self.partition_id.hash(state);
+ }
+}
+
+impl PhysicalExpr for SparkRandnExpr {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
+ Ok(DataType::Float64)
+ }
+
+ 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 rng = self.rng.lock();
+ let values =
+ Float64Array::from_iter_values(StandardNormal.sample_iter(&mut
*rng).take(num_rows));
+ Ok(ColumnarValue::Array(Arc::new(values)))
+ }
+
+ fn children(&self) -> Vec<&PhysicalExprRef> {
+ vec![]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<PhysicalExprRef>,
+ ) -> Result<PhysicalExprRef> {
+ Ok(Arc::new(Self::new(self.seed, self.partition_id)))
+ }
+
+ fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "randn({})", self.seed)
+ }
+}
+
+impl PartialEq<dyn Any> for SparkRandnExpr {
+ fn eq(&self, other: &dyn Any) -> bool {
+ down_cast_any_ref(other)
+ .downcast_ref::<Self>()
+ .map(|other| self.seed == other.seed && self.partition_id ==
other.partition_id)
+ .unwrap_or(false)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::{array::RecordBatch, datatypes::Schema};
+ use datafusion::common::{Result, cast::as_float64_array};
+
+ use super::*;
+
+ fn create_empty_batch(num_rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::empty());
+ RecordBatch::try_new_with_options(
+ schema,
+ vec![],
+
&arrow::array::RecordBatchOptions::new().with_row_count(Some(num_rows)),
+ )
+ .expect("Failed to create empty batch")
+ }
+
+ #[test]
+ fn test_randn_generates_different_values_per_row() -> Result<()> {
+ let expr = SparkRandnExpr::new(42, 0);
+ let batch = create_empty_batch(5);
+
+ let result = expr.evaluate(&batch)?;
+ let array = result.into_array(5)?;
+ let float_arr = as_float64_array(&array)?;
+
+ // Values should not be constant across rows, which verifies a value is
+ // generated per row rather than a single value being broadcast.
+ // (Individual samples are allowed to repeat, so we don't require all
+ // values to be distinct.)
+ let values: Vec<f64> = (0..5).map(|i| float_arr.value(i)).collect();
+ assert!(
+ values.iter().any(|&v| v != values[0]),
+ "Expected per-row values, but all rows were identical: {values:?}"
+ );
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_randn_reproducible_with_same_seed() -> Result<()> {
+ let expr1 = SparkRandnExpr::new(42, 0);
+ let expr2 = SparkRandnExpr::new(42, 0);
+ let batch = create_empty_batch(5);
+
+ let result1 = expr1.evaluate(&batch)?;
+ let result2 = expr2.evaluate(&batch)?;
+
+ let arr1_binding = result1.into_array(5)?;
+ let arr2_binding = result2.into_array(5)?;
+ let arr1 = as_float64_array(&arr1_binding)?;
+ let arr2 = as_float64_array(&arr2_binding)?;
+
+ for i in 0..5 {
+ assert_eq!(
+ arr1.value(i),
+ arr2.value(i),
+ "Same seed should produce same values"
+ );
+ }
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_randn_different_seeds_produce_different_values() -> Result<()> {
+ let expr1 = SparkRandnExpr::new(42, 0);
+ let expr2 = SparkRandnExpr::new(123, 0);
+ let batch = create_empty_batch(5);
+
+ let result1 = expr1.evaluate(&batch)?;
+ let result2 = expr2.evaluate(&batch)?;
+
+ let arr1_binding = result1.into_array(5)?;
+ let arr2_binding = result2.into_array(5)?;
+ let arr1 = as_float64_array(&arr1_binding)?;
+ let arr2 = as_float64_array(&arr2_binding)?;
+
+ // At least one value should be different
+ let any_different = (0..5).any(|i| arr1.value(i) != arr2.value(i));
+ assert!(
+ any_different,
+ "Different seeds should produce different values"
+ );
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_randn_different_partitions_produce_different_values() ->
Result<()> {
+ let expr1 = SparkRandnExpr::new(42, 0);
+ let expr2 = SparkRandnExpr::new(42, 1);
+ let batch = create_empty_batch(5);
+
+ let result1 = expr1.evaluate(&batch)?;
+ let result2 = expr2.evaluate(&batch)?;
+
+ let arr1_binding = result1.into_array(5)?;
+ let arr2_binding = result2.into_array(5)?;
+ let arr1 = as_float64_array(&arr1_binding)?;
+ let arr2 = as_float64_array(&arr2_binding)?;
+
+ // At least one value should be different
+ let any_different = (0..5).any(|i| arr1.value(i) != arr2.value(i));
+ assert!(
+ any_different,
+ "Different partitions should produce different values"
+ );
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_randn_stateful_across_batches() -> Result<()> {
+ let expr = SparkRandnExpr::new(42, 0);
+ let batch1 = create_empty_batch(3);
+ let batch2 = create_empty_batch(3);
+
+ // Evaluate two batches sequentially
+ let result1 = expr.evaluate(&batch1)?;
+ let result2 = expr.evaluate(&batch2)?;
+
+ let arr1_binding = result1.into_array(3)?;
+ let arr2_binding = result2.into_array(3)?;
+ let arr1 = as_float64_array(&arr1_binding)?;
+ let arr2 = as_float64_array(&arr2_binding)?;
+
+ // Collect all values
+ let values1: Vec<f64> = (0..3).map(|i| arr1.value(i)).collect();
+ let values2: Vec<f64> = (0..3).map(|i| arr2.value(i)).collect();
+
+ // Second batch should continue from where first left off (not restart)
+ // So values should be different between batches
+ assert_ne!(values1, values2, "Batches should have different values");
+
+ // Compare with fresh expr that evaluates both batches together
+ let expr_fresh = SparkRandnExpr::new(42, 0);
+ let batch_combined = create_empty_batch(6);
+ let result_combined = expr_fresh.evaluate(&batch_combined)?;
+ let arr_combined_binding = result_combined.into_array(6)?;
+ let arr_combined = as_float64_array(&arr_combined_binding)?;
+
+ // First 3 values should match values1, next 3 should match values2
+ for i in 0..3 {
+ assert_eq!(
+ arr_combined.value(i),
+ values1[i],
+ "First batch values should match"
+ );
+ assert_eq!(
+ arr_combined.value(i + 3),
+ values2[i],
+ "Second batch values should match continuation"
+ );
+ }
+
+ Ok(())
+ }
+}
diff --git
a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
index 0a34dedf..40f44d43 100644
---
a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
+++
b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
@@ -1014,6 +1014,68 @@ class AuronFunctionSuite extends AuronQueryTest with
BaseAuronSQLSuite {
}
}
+ test("randn function with seed") {
+ withTable("t1") {
+ sql("CREATE TABLE t1(id INT) USING parquet")
+ sql("INSERT INTO t1 VALUES(1), (2), (3)")
+
+ // randn is non-deterministic and intentionally does not replicate
Spark's RNG, so its
+ // values cannot be compared against vanilla Spark. Verify it runs
natively, produces a
+ // non-null value per row, and is reproducible for a fixed seed.
+ val query = "SELECT id, randn(42) AS r1, randn(100) AS r2 FROM t1 ORDER
BY id"
+ val df = sql(query)
+ val rows = df.collect()
+ assertPlanIsNative(df)
+
+ assert(rows.length == 3)
+ assert(rows.forall(r => !r.isNullAt(1) && !r.isNullAt(2)))
+
+ // Same seed -> same values across executions.
+ val rows2 = sql(query).collect()
+ assert(rows.map(_.getDouble(1)).sameElements(rows2.map(_.getDouble(1))))
+ assert(rows.map(_.getDouble(2)).sameElements(rows2.map(_.getDouble(2))))
+
+ // Different seeds -> different values.
+ assert(rows.exists(r => r.getDouble(1) != r.getDouble(2)))
+ }
+ }
+
+ test("randn function with foldable seed expression") {
+ withTable("t1") {
+ sql("CREATE TABLE t1(id INT) USING parquet")
+ sql("INSERT INTO t1 VALUES(1), (2), (3)")
+
+ // A foldable (non-literal) seed must still be converted to the native
randn expression.
+ val query = "SELECT id, randn(cast(42 as bigint)) AS r FROM t1 ORDER BY
id"
+ val df = sql(query)
+ val rows = df.collect()
+ assertPlanIsNative(df)
+
+ assert(rows.length == 3)
+ assert(rows.forall(r => !r.isNullAt(1)))
+
+ // Reproducible for a fixed seed.
+ val rows2 = sql(query).collect()
+ assert(rows.map(_.getDouble(1)).sameElements(rows2.map(_.getDouble(1))))
+ }
+ }
+
+ test("randn function without seed") {
+ withTable("t1") {
+ sql("CREATE TABLE t1(id INT) USING parquet")
+ sql("INSERT INTO t1 VALUES(1), (2), (3)")
+
+ // randn() with no seed uses a randomly assigned seed, so values are not
reproducible
+ // across executions. Verify it still runs natively and produces a
non-null value per row.
+ val df = sql("SELECT id, randn() AS r FROM t1 ORDER BY id")
+ val rows = df.collect()
+ assertPlanIsNative(df)
+
+ assert(rows.length == 3)
+ assert(rows.forall(r => !r.isNullAt(1)))
+ }
+ }
+
test("ascii function") {
withTable("t1") {
sql("create table t1(c1 string) using parquet")
diff --git
a/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/AuronQueryTest.scala
b/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/AuronQueryTest.scala
index 678faea8..263f0623 100644
---
a/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/AuronQueryTest.scala
+++
b/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/AuronQueryTest.scala
@@ -63,20 +63,25 @@ abstract class AuronQueryTest
checkAnswer(dfAuron, expected)
if (requireNative) {
- val plan = stripAQEPlan(dfAuron.queryExecution.executedPlan)
- plan
- .collectFirst { case op if !isNativeOrPassThrough(op) => op }
- .foreach { op: SparkPlan =>
- fail(s"""
- |Found non-native operator: ${op.nodeName}
- |plan:
- |${plan}""".stripMargin)
- }
+ assertPlanIsNative(dfAuron)
}
dfAuron
}
+ /** Fail if any operator in the executed plan is not native or a
pass-through. */
+ protected def assertPlanIsNative(df: DataFrame): Unit = {
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ plan
+ .collectFirst { case op if !isNativeOrPassThrough(op) => op }
+ .foreach { op: SparkPlan =>
+ fail(s"""
+ |Found non-native operator: ${op.nodeName}
+ |plan:
+ |${plan}""".stripMargin)
+ }
+ }
+
protected def isNativeOrPassThrough(op: SparkPlan): Boolean = op match {
case _: NativeSupports => true
case e: UnaryExecNode
diff --git
a/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala
b/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala
index 378a8d66..95f42e0f 100644
---
a/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala
+++
b/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala
@@ -878,6 +878,21 @@ object NativeConverters extends Logging {
case e: Hex => buildScalarFunction(pb.ScalarFunction.Hex, e.children,
e.dataType)
case e: IsNaN =>
buildExtScalarFunction("Spark_IsNaN", e.children, e.dataType)
+ case e: Randn =>
+ val seedOpt: Option[Long] = e.child match {
+ case Literal(s: Long, _) => Some(s)
+ case Literal(s: Int, _) => Some(s.toLong)
+ case expr if expr.foldable =>
+ Option(expr.eval(null)).collect { case n: Number => n.longValue() }
+ case _ => None
+ }
+ seedOpt
+ .map { seed =>
+ buildExprNode {
+
_.setSparkRandnExpr(pb.SparkRandnExprNode.newBuilder().setSeed(seed))
+ }
+ }
+ .getOrElse(fallback(e))
case e: Round =>
e.scale match {
case Literal(n: Int, _) =>
