feat: Add `datafusion-spark` crate #15168

shehabgamin · 2025-03-11T21:24:15Z

Which issue does this PR close?

Rationale for this change

See discussion in #5600

TL;DR Many projects want Spark-compatible expressions for use with DataFusion. There are some in Comet and there are some in the Sail project.

What changes are included in this PR?

Adding Spark crate.

Are these changes tested?

Are there any user-facing changes?

shehabgamin · 2025-03-13T02:26:13Z

datafusion/sqllogictest/test_files/spark/README.md

+  specific language governing permissions and limitations
+  under the License.
+-->
+


Testing instructions here

shehabgamin · 2025-03-13T02:26:31Z

datafusion/spark/README.md

+# datafusion-spark: Spark-compatible Expressions
+
+This crate provides Apache Spark-compatible expressions for use with DataFusion.
+


Testing instructions here

shehabgamin · 2025-03-13T02:26:54Z

datafusion/sqllogictest/test_files/spark/math/expm1.slt

+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+


Example test here

The expm1 probably works identically across all Spark versions and is not affected by different configuration settings, but many expressions are affected by settings such as ANSI mode and different date/time formats and timezones.

How do you think we should handle these different cases with this test approach?

A related question - There will be some functions that we can support with 100% compatibility and some that we cannot. It would be good to think about how we express that.

Sorry for all the questions, but I am really excited about this 😄 ... what would be involved in being able to run these same sqllogic test files in Spark (either in CI or manually locally) to confirm same/similar results

These are great questions and really good things to discuss. I'm about to go into a meeting but I have a bunch of thoughts that I'll share afterwards in a couple of hours.

The expm1 probably works identically across all Spark versions and is not affected by different configuration settings, but many expressions are affected by settings such as ANSI mode and different date/time formats and timezones.

How do you think we should handle these different cases with this test approach?

In the Sail code base, auxiliary information is passed into new() and stored within the struct. For example:
https://github.com/lakehq/sail/blob/be54362ce8bd79bfb3e55214c5a1b80c2c3d2492/crates/sail-plan/src/extension/function/datetime/timestamp_now.rs#L9-L32

#[derive(Debug)] pub struct TimestampNow { signature: Signature, timezone: Arc<str>, time_unit: TimeUnit, } impl TimestampNow { pub fn new(timezone: Arc<str>, time_unit: TimeUnit) -> Self { Self { signature: Signature::nullary(Volatility::Stable), timezone, time_unit, } } pub fn timezone(&self) -> &str { &self.timezone } pub fn time_unit(&self) -> &TimeUnit { &self.time_unit } }

And then in our PhysicalExtensionCodec we can do the following:
https://github.com/lakehq/sail/blob/be54362ce8bd79bfb3e55214c5a1b80c2c3d2492/crates/sail-execution/src/codec.rs#L946-L953

if let Some(func) = node.inner().as_any().downcast_ref::<TimestampNow>() { let timezone = func.timezone().to_string(); let time_unit: gen_datafusion_common::TimeUnit = func.time_unit().into(); let time_unit = time_unit.as_str_name().to_string(); UdfKind::TimestampNow(gen::TimestampNowUdf { timezone, time_unit, })

If we decide to not use sqllogictest (per #15168 (comment)) then we will have no problem testing UDFs with auxiliary information. There are already tests in DataFusion core for this type of pattern as well:

datafusion/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

Lines 1265 to 1304 in 8061485

async fn test_parameterized_scalar_udf() -> Result<()> {

let batch = RecordBatch::try_from_iter([(

"text",

Arc::new(StringArray::from(vec!["foo", "bar", "foobar", "barfoo"])) as ArrayRef,

)])?;

let ctx = SessionContext::new();

ctx.register_batch("t", batch)?;

let t = ctx.table("t").await?;

let foo_udf = ScalarUDF::from(MyRegexUdf::new("fo{2}"));

let bar_udf = ScalarUDF::from(MyRegexUdf::new("[Bb]ar"));

let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)

.filter(

foo_udf

.call(vec![col("text")])

.and(bar_udf.call(vec![col("text")])),

)?

.filter(col("text").is_not_null())?

.build()?;

assert_eq!(

format!("{plan}"),

"Filter: t.text IS NOT NULL\n Filter: regex_udf(t.text) AND regex_udf(t.text)\n TableScan: t projection=[text]"

);

let actual = DataFrame::new(ctx.state(), plan).collect().await?;

let expected = [

"+--------+",

"| text |",

"+--------+",

"| foobar |",

"| barfoo |",

"+--------+",

];

assert_batches_eq!(expected, &actual);

ctx.deregister_table("t")?;

Ok(())

}

A related question - There will be some functions that we can support with 100% compatibility and some that we cannot. It would be good to think about how we express that.

Throw errors when possible, and provide documentation. Depending on how shaky the compatibility for the function is, we may want to avoid implementing it altogether.

... what would be involved in being able to run these same sqllogic test files in Spark (either in CI or manually locally) to confirm same/similar results

One option:

A Python script to automatically generate Spark SQL function test cases and their results using PySpark.

A README for developers explaining how to run the script and commit the test cases.

An optional CI workflow to verify the correctness of the test cases' ground truth on demand.

Separate topic... Do you have ideas about fuzzy testing and its suitability in DataFusion?

Here is my suggestion.

Write a Python script that generates interesting test data (probably in Parquet format) with edge cases using PySpark

Create files containing SQL queries that operate on these test files

Write a Python script to run those queries via PySpark and write results out to file

Write a Rust script to run those queries using datafusion-spark and write results out to file

Write a script that can compare the Spark and datafusion-spark output and report on any differences

shehabgamin · 2025-03-13T02:27:01Z

datafusion/sqllogictest/test_files/spark/string/ascii.slt

+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+


Example test here

shehabgamin · 2025-03-13T02:33:20Z

datafusion/spark/src/function/string/ascii.rs

+        };
+    }
+
+    #[test]


Example invoke test here. Direct invocation tests should only be used to verify that the function is correctly implemented. Further tests are required in the sqllogictest crate (examples for ascii can be found in this PR).

shehabgamin · 2025-03-13T02:34:11Z

datafusion/spark/src/function/math/expm1.rs

+        };
+    }
+
+    #[test]


Example invoke test here. Direct invocation tests should only be used to verify that the function is correctly implemented. Further tests are required in the sqllogictest crate (examples for expm1 can be found in this PR).

shehabgamin · 2025-03-13T02:48:07Z

datafusion/spark/src/function/string/ascii.rs

+    }
+
+    fn name(&self) -> &str {
+        "spark_ascii"


Prefix with spark_ because sqllogictest evaluates both implementations of ascii.

shehabgamin · 2025-03-13T02:49:35Z

datafusion/spark/src/function/math/expm1.rs

+    }
+
+    fn name(&self) -> &str {
+        "spark_expm1"


Prefix with spark_ because sqllogictest may evaluate more than one implementation of expm1.

andygrove · 2025-03-13T23:58:12Z

datafusion/sqllogictest/test_files/spark/math/expm1.slt

+SELECT spark_expm1(1::INT);
+----
+1.718281828459


I tried running this in Spark 3.5.3 and did not get the same answer.

scala> spark.sql("select expm1(1)").show() +-----------------+ | EXPM1(1)| +-----------------+ |1.718281828459045| +-----------------+

@andygrove I believe sqllogic is truncating the answer.

I initially had this test, which tested for the value 1.7182818284590453 (slightly more precise than your result) but removed it because cargo test (amd64) was giving the value 1.718281828459045 (https://github.com/apache/datafusion/actions/runs/13825914914/job/38680868216) while the rest of the cargo tests on different architectures were passing.

Does sqllogictest have a way to test floating point results within some tolerance?

Doesn't look like it:

Remember, the purpose of sqllogictest is to validate the logic behind the evaluation of SQL statements, not the ability to handle extreme values. So keep content in a reasonable range: small integers, short strings, and floating point numbers that use only the most significant bits of an a 32-bit IEEE float.

https://www.sqlite.org/sqllogictest/doc/trunk/about.wiki

We don't have to stick with sqllogictest.

We can create test helpers similar to:
https://github.com/lakehq/datafusion/blob/d78877a55c5e835a07a7ebf23a7bd515faf7d827/datafusion/optimizer/src/analyzer/type_coercion.rs#L2137-L2208

The above link is from an old PR that didn't end up getting merged in, but the general idea seems useful here.

lol

query R SELECT spark_expm1(1::INT); ---- 1.718281828459 query T SELECT spark_expm1(1::INT)::STRING; ---- 1.7182818284590453

andygrove · 2025-03-14T00:08:46Z

datafusion/sqllogictest/test_files/spark/math/expm1.slt

+1.718281828459
+
+query R
+SELECT spark_expm1(a) FROM (VALUES (0::INT), (1::INT)) AS t(a);


I'd suggest adding tests for a wider range of values and edge cases, such as negative numbers, large positive and negative numbers, NaN, null, and so on.

Good call, will do!

andygrove · 2025-03-14T00:11:40Z

datafusion/sqllogictest/test_files/spark/string/ascii.slt

+50
+
+query I
+SELECT spark_ascii(a) FROM (VALUES ('Spark'), ('PySpark'), ('Pandas API')) AS t(a);


Could you add tests for edge cases?

Some ideas from ChatGPT (the results are from me actually running them in Spark):

scala> spark.sql("select ascii('😀')").show() +---------+ |ascii(😀)| +---------+ | 128512| +---------+ scala> spark.sql("select ascii('\n')").show() +---------+ |ascii(\n)| +---------+ | 10| +---------+ scala> spark.sql("select ascii('\t')").show() +---------+ |ascii(\t)| +---------+ | 9| +---------+

Good call, will do!

andygrove · 2025-03-14T00:25:34Z

datafusion/core/src/execution/session_state_defaults.rs

+        #[cfg(feature = "spark")]
+        functions.append(&mut datafusion_spark::all_default_scalar_functions());
+


I don't have a strong opinion on this, but I suspect many DataFusion users may not be keen on having Spark-specific work affecting the core.

I don't have a strong opinion on this either. I'm open to alternatives.

feat: Add datafusion-spark crate

57018b7

github-actions bot added the functions Changes to functions implementation label Mar 11, 2025

shehabgamin mentioned this pull request Mar 11, 2025

[DISCUSSION] Add separate crate to cover spark builtin functions #5600

Open

shehabgamin added 4 commits March 12, 2025 17:21

spark crate setup

762ef03

clean up 2 example functions

6e1aa06

cleanup crate

eba0a51

Spark crate setup

9ca4a52

shehabgamin marked this pull request as ready for review March 13, 2025 02:15

github-actions bot added core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) labels Mar 13, 2025

fix lint issue

074fc70

shehabgamin commented Mar 13, 2025

View reviewed changes

cargo cleanup

4e28310

shehabgamin commented Mar 13, 2025

View reviewed changes

fix collision in sqllogic

7f448f3

shehabgamin commented Mar 13, 2025

View reviewed changes

remove redundant test

c29a9e4

alamb mentioned this pull request Mar 13, 2025

Weekly Plan (Andrew Lamb) March 10, 2025 #15121

Open

11 tasks

andygrove reviewed Mar 13, 2025

View reviewed changes

andygrove reviewed Mar 14, 2025

View reviewed changes

shehabgamin added 3 commits March 13, 2025 22:32

fix merge conflict

b0cdaf3

test float precision when casting to string

b2e9135

reorder

4fedbf4

undo

9482964

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Add `datafusion-spark` crate #15168

feat: Add `datafusion-spark` crate #15168

shehabgamin commented Mar 11, 2025 •

edited

Loading

shehabgamin Mar 13, 2025

shehabgamin Mar 13, 2025

shehabgamin Mar 13, 2025

andygrove Mar 14, 2025

andygrove Mar 14, 2025

andygrove Mar 14, 2025

shehabgamin Mar 14, 2025

andygrove Mar 14, 2025

shehabgamin Mar 14, 2025

shehabgamin Mar 14, 2025

shehabgamin Mar 14, 2025

andygrove Mar 14, 2025

shehabgamin Mar 13, 2025

shehabgamin Mar 13, 2025 •

edited

Loading

shehabgamin Mar 13, 2025

shehabgamin Mar 13, 2025 •

edited

Loading

shehabgamin Mar 13, 2025

andygrove Mar 13, 2025

shehabgamin Mar 14, 2025

andygrove Mar 14, 2025

shehabgamin Mar 14, 2025

shehabgamin Mar 14, 2025

shehabgamin Mar 14, 2025

andygrove Mar 14, 2025

shehabgamin Mar 14, 2025

andygrove Mar 14, 2025 •

edited

Loading

shehabgamin Mar 14, 2025

andygrove Mar 14, 2025

shehabgamin Mar 14, 2025

		# datafusion-spark: Spark-compatible Expressions

		This crate provides Apache Spark-compatible expressions for use with DataFusion.

	async fn test_parameterized_scalar_udf() -> Result<()> {
	let batch = RecordBatch::try_from_iter([(
	"text",
	Arc::new(StringArray::from(vec!["foo", "bar", "foobar", "barfoo"])) as ArrayRef,
	)])?;

	let ctx = SessionContext::new();
	ctx.register_batch("t", batch)?;
	let t = ctx.table("t").await?;
	let foo_udf = ScalarUDF::from(MyRegexUdf::new("fo{2}"));
	let bar_udf = ScalarUDF::from(MyRegexUdf::new("[Bb]ar"));

	let plan = LogicalPlanBuilder::from(t.into_optimized_plan()?)
	.filter(
	foo_udf
	.call(vec![col("text")])
	.and(bar_udf.call(vec![col("text")])),
	)?
	.filter(col("text").is_not_null())?
	.build()?;

	assert_eq!(
	format!("{plan}"),
	"Filter: t.text IS NOT NULL\n Filter: regex_udf(t.text) AND regex_udf(t.text)\n TableScan: t projection=[text]"
	);

	let actual = DataFrame::new(ctx.state(), plan).collect().await?;
	let expected = [
	"+--------+",
	"\| text \|",
	"+--------+",
	"\| foobar \|",
	"\| barfoo \|",
	"+--------+",
	];
	assert_batches_eq!(expected, &actual);

	ctx.deregister_table("t")?;
	Ok(())
	}

+                      };
+                  }
+                  #[test]

		#[cfg(feature = "spark")]
		functions.append(&mut datafusion_spark::all_default_scalar_functions());

feat: Add datafusion-spark crate #15168

Are you sure you want to change the base?

feat: Add datafusion-spark crate #15168

Conversation

shehabgamin commented Mar 11, 2025 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shehabgamin Mar 13, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shehabgamin Mar 13, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

andygrove Mar 14, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

feat: Add `datafusion-spark` crate #15168

feat: Add `datafusion-spark` crate #15168

shehabgamin commented Mar 11, 2025 •

edited

Loading

shehabgamin Mar 13, 2025 •

edited

Loading

shehabgamin Mar 13, 2025 •

edited

Loading

andygrove Mar 14, 2025 •

edited

Loading