[AURON #1891] Implement randn() function by robreeves · Pull Request #1938 · apache/auron

robreeves · 2026-01-21T05:35:05Z

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.

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, extensions: None, metadata_size_hint: None }], statistics: Some(Statistics { num_rows: Exact(0), total_byte_size: Exact(0), column_statistics: [ColumnStatistics { null_count: Absent, max_value: Absent, min_value: Absent, sum_value: Absent, distinct_count: Absent }] }) }, FileGroup { files: [], statistics: None }, FileGroup { files: [], statistics: None }, FileGroup { files: [], statistics: None }, FileGroup { files: [], statistics: None }], predicate=Some(Literal { value: Boolean(true), field: Field { name: "lit", data_type: Boolean, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }), schema=[id:Int64;N]

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" }, last_modified: 1970-01-01T00:00:00Z, size: 297, e_tag: None, version: None }, partition_values: [], range: Some(FileRange { start: 0, end: 297 }), statistics: None, extensions: None, metadata_size_hint: None }], statistics: Some(Statistics { num_rows: Exact(0), total_byte_size: Exact(0), column_statistics: [ColumnStatistics { null_count: Absent, max_value: Absent, min_value: Absent, sum_value: Absent, distinct_count: Absent }] }) }], predicate=Some(Literal { value: Boolean(true), field: Field { name: "lit", data_type: Boolean, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }), schema=[id:Int64;N]

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 { start: 0, end: 472 }), statistics: None, extensions: None, metadata_size_hint: None }], statistics: Some(Statistics { num_rows: Exact(0), total_byte_size: Exact(0), column_statistics: [ColumnStatistics { null_count: Absent, max_value: Absent, min_value: Absent, sum_value: Absent, distinct_count: Absent }] }) }, FileGroup { files: [], statistics: None }, FileGroup { files: [], statistics: None }, FileGroup { files: [], statistics: None }], predicate=Some(Literal { value: Boolean(true), field: Field { name: "lit", data_type: Boolean, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }), schema=[id:Int64;N]

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: 471, e_tag: None, version: None }, partition_values: [], range: Some(FileRange { start: 0, end: 471 }), statistics: None, extensions: None, metadata_size_hint: None }], statistics: Some(Statistics { num_rows: Exact(0), total_byte_size: Exact(0), column_statistics: [ColumnStatistics { null_count: Absent, max_value: Absent, min_value: Absent, sum_value: Absent, distinct_count: Absent }] }) }, FileGroup { files: [], statistics: None }], predicate=Some(Literal { value: Boolean(true), field: Field { name: "lit", data_type: Boolean, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }), schema=[id:Int64;N]

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 }, partition_values: [], range: Some(FileRange { start: 0, end: 472 }), statistics: None, extensions: None, metadata_size_hint: None }], statistics: Some(Statistics { num_rows: Exact(0), total_byte_size: Exact(0), column_statistics: [ColumnStatistics { null_count: Absent, max_value: Absent, min_value: Absent, sum_value: Absent, distinct_count: Absent }] }) }, FileGroup { files: [], statistics: None }, FileGroup { files: [], statistics: None }], predicate=Some(Literal { value: Boolean(true), field: Field { name: "lit", data_type: Boolean, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }), schema=[id:Int64;N]

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 }], statistics: Some(Statistics { num_rows: Exact(0), total_byte_size: Exact(0), column_statistics: [ColumnStatistics { null_count: Absent, max_value: Absent, min_value: Absent, sum_value: Absent, distinct_count: Absent }] }) }, FileGroup { files: [], statistics: None }, FileGroup { files: [], statistics: None }, FileGroup { files: [], statistics: None }, FileGroup { files: [], statistics: None }, FileGroup { files: [], statistics: None }], predicate=Some(Literal { value: Boolean(true), field: Field { name: "lit", data_type: Boolean, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }), schema=[id:Int64;N]

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])

Copilot

Pull request overview

This PR implements the randn() function to improve Spark function coverage in Auron. The function generates random values from a standard normal distribution with optional seed support.

Changes:

Added Rust implementation of spark_randn function with seed handling
Registered the new function in the Scala converter and Rust function registry
Added rand_distr dependency for normal distribution sampling

Reviewed changes

Copilot reviewed 5 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala	Added case handler for Randn expression to route to native implementation
native-engine/datafusion-ext-functions/src/spark_randn.rs	New implementation of randn function with seed handling and unit tests
native-engine/datafusion-ext-functions/src/lib.rs	Registered Spark_Randn function in the extension function factory
native-engine/datafusion-ext-functions/Cargo.toml	Added rand and rand_distr dependencies
Cargo.toml	Added rand_distr workspace dependency
Cargo.lock	Updated lock file with rand_distr package metadata

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

native-engine/datafusion-ext-functions/src/spark_randn.rs

Resolve conflicts between randn and spark_partition_id features: - Proto: spark_partition_id_expr at 20101, randn_expr at 20102 - Planner: include both expression handlers Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

yew1eb · 2026-01-31T04:35:33Z

@robreeves Nice work! LGTM.
Could add SQL unit tests (per AuronFunctionSuite) to align with Spark SQl's semantics, thanks!

Resolved conflicts by assigning separate IDs to randn and monotonically_increasing_id: - MonotonicIncreasingIdExprNode: ID 20102 - RandnExprNode: ID 20103 Both expressions are now supported in the proto definitions and planner. Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Add test to AuronFunctionSuite to verify randn functionality with seeds. The test validates that Auron's native randn implementation produces the same reproducible results as Spark's baseline when using explicit seeds. Test covers: - randn with seed 42 - randn with seed 100 - Validates against Spark baseline using checkSparkAnswerAndOperator Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

robreeves · 2026-02-04T04:26:48Z

@robreeves Nice work! LGTM. Could add SQL unit tests (per AuronFunctionSuite) to align with Spark SQl's semantics, thanks!

I added a AuronFunctionSuite test. Thanks.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

robreeves · 2026-02-08T15:52:09Z

@richox can you run the PR checks again?

robreeves · 2026-02-12T18:59:38Z

@cxzl25 can you run the PR checks? Thanks

ShreyeshArangath

Changes LGTM, just one comment about the naming of this rust function.

native-engine/datafusion-ext-exprs/src/lib.rs

…ntion Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

…nces Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

robreeves added 4 commits January 18, 2026 10:37

randn

175fa00

reduce arg parsing complexity

5bb5ca5

remove extra method

3d0fa7d

use rand_dist

0f793c4

github-actions bot added spark native build labels Jan 21, 2026

robreeves added 2 commits January 20, 2026 21:40

revert auto format change

0d041f3

revert auto format changes

0b93080

cxzl25 requested a review from Copilot January 21, 2026 05:48

Copilot started reviewing on behalf of cxzl25 January 21, 2026 05:49 View session

Copilot AI reviewed Jan 21, 2026

View reviewed changes

native-engine/datafusion-ext-functions/src/spark_randn.rs Outdated Show resolved Hide resolved

native-engine/datafusion-ext-functions/src/spark_randn.rs Outdated Show resolved Hide resolved

native-engine/datafusion-ext-functions/src/spark_randn.rs Outdated Show resolved Hide resolved

robreeves and others added 8 commits January 21, 2026 21:40

autoformat

3b78e73

handle batches

4b2261b

new approach

81beb0b

forced format changes during build

a8183ae

Merge origin/master into randn

f591295

Resolve conflicts between randn and spark_partition_id features: - Proto: spark_partition_id_expr at 20101, randn_expr at 20102 - Planner: include both expression handlers Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

simplify evaluate

0ea5dc9

revert whitespace

f792b1b

revert whitespace

dd6c98b

robreeves marked this pull request as ready for review January 31, 2026 00:26

robreeves and others added 2 commits February 3, 2026 20:00

robreeves and others added 2 commits February 8, 2026 07:36

Fix rustfmt formatting in randn.rs

388d0a1

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Fix clippy uninlined_format_args in randn.rs

6e1f2d4

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

ShreyeshArangath reviewed Feb 15, 2026

View reviewed changes

native-engine/datafusion-ext-exprs/src/lib.rs Outdated Show resolved Hide resolved

robreeves and others added 2 commits February 15, 2026 21:24

Rename randn to spark_randn to align with Spark-specific naming conve…

048fdfb

…ntion Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Rename RandnExprNode to SparkRandnExprNode in proto and update refere…

8f554f8

…nces Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

ShreyeshArangath approved these changes Feb 16, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[AURON #1891] Implement randn() function#1938

[AURON #1891] Implement randn() function#1938
robreeves wants to merge 20 commits intoapache:masterfrom
robreeves:randn

robreeves commented Jan 21, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yew1eb commented Jan 31, 2026

Uh oh!

robreeves commented Feb 4, 2026

Uh oh!

robreeves commented Feb 8, 2026

Uh oh!

robreeves commented Feb 12, 2026

Uh oh!

ShreyeshArangath left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

robreeves commented Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

How was this patch tested?

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yew1eb commented Jan 31, 2026

Uh oh!

robreeves commented Feb 4, 2026

Uh oh!

robreeves commented Feb 8, 2026

Uh oh!

robreeves commented Feb 12, 2026

Uh oh!

ShreyeshArangath left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

robreeves commented Jan 21, 2026 •

edited

Loading