[datafusion-spark] Implement ceil&floor function for spark by irenjj · Pull Request #15958 · apache/datafusion

irenjj · 2025-05-06T08:35:12Z

Which issue does this PR close?

Closes [datafusion-spark] Implement ceil function #15916

Rationale for this change

Add Spark-style ceil&floor function, the differences between DataFusion and Spark's implementations of mathematical functions ceil and floor:

Function Signatures:
- Spark provides a more flexible implementation with ceil(x, scale) accepting two parameters
- DataFusion only implements the single parameter version ceil(x)
Type Support:
- Spark has broader type support, handling INTEGER, DECIMAL, DOUBLE, and FLOAT types
- DataFusion currently only supports Float32 and Float64 types

What changes are included in this PR?

Add new ceil&floor function.

Are these changes tested?

Yes.

Are there any user-facing changes?

No.

irenjj · 2025-05-06T08:56:04Z

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

+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Err(generic_internal_err(
+            "ceil",
+            "`return_type` should not be called, call `return_type_from_args` instead",
+        ))
+    }


Hi @shehabgamin, we encountered an issue here: most of the code was copied from sail. In sail, the return value is constructed using return_type_from_args, which requires access to both args.arg_types and args.scalar_arguments to determine the final value. However, in the latest version of DF, return_type_from_args has been removed. It seems that relying solely on datatype is no longer sufficient to construct the return value within return_type().

@irenjj Would return_field_from_args work here?

datafusion/datafusion/expr/src/udf.rs

Lines 447 to 487 in 7e89862

/// What type will be returned by this function, given the arguments?

///

/// By default, this function calls [`Self::return_type`] with the

/// types of each argument.

///

/// # Notes

///

/// Most UDFs should implement [`Self::return_type`] and not this

/// function as the output type for most functions only depends on the types

/// of their inputs (e.g. `sqrt(f32)` is always `f32`).

///

/// This function can be used for more advanced cases such as:

///

/// 1. specifying nullability

/// 2. return types based on the **values** of the arguments (rather than

/// their **types**.

///

/// # Output Type based on Values

///

/// For example, the following two function calls get the same argument

/// types (something and a `Utf8` string) but return different types based

/// on the value of the second argument:

///

/// * `arrow_cast(x, 'Int16')` --> `Int16`

/// * `arrow_cast(x, 'Float32')` --> `Float32`

///

/// # Requirements

///

/// This function **must** consistently return the same type for the same

/// logical input even if the input is simplified (e.g. it must return the same

/// value for `('foo' | 'bar')` as it does for ('foobar').

fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<Field> {

let data_types = args

.arg_fields

.iter()

.map(|f| f.data_type())

.cloned()

.collect::<Vec<_>>();

let return_type = self.return_type(&data_types)?;

Ok(Field::new(self.name(), return_type, true))

}

shehabgamin · 2025-05-06T21:47:56Z

@irenjj Here are some tests for ceil and floor that you can re-use. Sail adheres to Spark SQL syntax, so you may need to adjust as needed for this pr:

shehabgamin · 2025-05-07T22:57:31Z

@irenjj I'll review this by tomorrow!

I'm not a committer though, so we'll still need a review from someone else as well. cc @alamb @andygrove

andygrove · 2025-05-07T23:18:20Z

Hi @irenjj. Could you explain in the PR description how these functions differ from the standard DataFusion implementation (i.e., what is Spark-specific about them)? That will help with reviews.

andygrove · 2025-05-07T23:23:06Z

@parthchandra @huaxingao @comphead @mbutrovich @kazuyukitanimura, FYI, in case you want to review. In Comet, we are currently using DataFusion's ceil and floor functions and have not specialized for Spark.

andygrove · 2025-05-08T02:03:36Z

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

+        let arg_len = args.args.len();
+        let target_scale = if arg_len == 1 {
+            Ok(&None)
+        } else if arg_len == 2 {


Floor and Ceil expressions in Spark are unary expressions and only take a single argument:

case class Ceil(child: _root_.org.apache.spark.sql.catalyst.expressions.Expression) extends _root_.org.apache.spark.sql.catalyst.expressions.UnaryMathExpression {

When calling ceil in Spark SQL with two parameters, it looks like a cast is added to the input expression. For example, ceil(_1, 2) translates to plan: *(1) Project [ceil(cast(_1#6 as decimal(30,15)), 2) AS ceil(_1, 2)#10.

Perhaps you implemented the two-argument version to work around the fact that DataFusion's logical optimizer doesn't implement the same logic as Spark for adding this cast? I wonder if it is worth looking into adding a Spark-compatible optimizer rule to add this cast in the future.

Actually, I have confused myself now. The plan does still show two arguments, but the physical plan only has a single argument 😕

Ah, there is a separate RoundCeil in Spark for the two-argument case:

case class RoundCeil(child: Expression, scale: Expression)

Got it. In this PR, we only need to implement the single-argument version of ceil, and add the new function roundceil as a two-argument implementation in a future PR.

I think it is fine to have both cases covered in a single expression in DataFusion. Sorry for the confusion.

andygrove · 2025-05-08T02:14:40Z

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

+query error
+SELECT ceil(3.1411, -3);
+----
+DataFusion error: Error during planning: Failed to coerce arguments to satisfy a call to 'ceil' function: coercion from [Float64, Int64] to the signature Uniform(1, [Float64, Float32]) failed


I assume that these errors need to be fixed still

@irenjj Sail doesn't use the same SQL parser as DataFusion, as it follows Spark SQL syntax.

You’ll need to cast values to ensure the tests interpret the correct DataType.

For example, instead of this:

SELECT ceil(3.1411, -3);

Use:

SELECT ceil(3.1411::decimal(5,4), -3::INT);

Copy-pasting the exact queries from the files below won't work directly. However, if you look closely, you’ll see that each query includes the necessary cast for the first argument (the second argument is always an INT):

https://github.com/lakehq/sail/blob/main/python/pysail/tests/spark/function/test_ceil.txt

https://github.com/lakehq/sail/blob/main/python/pysail/tests/spark/function/test_floor.txt

However, you should expect the same result values from those two files.

shehabgamin

Nice progress! Just gotta clean up the test files a bit :)

shehabgamin · 2025-05-09T02:56:46Z

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

+query error
+SELECT ceil(3.1411, -3);
+----
+DataFusion error: Error during planning: Failed to coerce arguments to satisfy a call to 'ceil' function: coercion from [Float64, Int64] to the signature Uniform(1, [Float64, Float32]) failed


@irenjj Sail doesn't use the same SQL parser as DataFusion, as it follows Spark SQL syntax.

You’ll need to cast values to ensure the tests interpret the correct DataType.

For example, instead of this:

SELECT ceil(3.1411, -3);

Use:

SELECT ceil(3.1411::decimal(5,4), -3::INT);

Copy-pasting the exact queries from the files below won't work directly. However, if you look closely, you’ll see that each query includes the necessary cast for the first argument (the second argument is always an INT):

https://github.com/lakehq/sail/blob/main/python/pysail/tests/spark/function/test_ceil.txt

https://github.com/lakehq/sail/blob/main/python/pysail/tests/spark/function/test_floor.txt

However, you should expect the same result values from those two files.

alamb · 2025-06-17T22:55:14Z

@irenjj - I wonder if you would be willing to help pick this PR back up now that we have merged a PR with a bunch of tests from @shehabgamin here:

chore: generate basic spark function tests #16409

We are hoping to generate some example PRs to use when starting to fill out the rest of the functions and I was thinking this one might be a good one

irenjj · 2025-06-17T23:25:58Z

@irenjj - I wonder if you would be willing to help pick this PR back up now that we have merged a PR with a bunch of tests from @shehabgamin here:

chore: generate basic spark function tests #16409

We are hoping to generate some example PRs to use when starting to fill out the rest of the functions and I was thinking this one might be a good one

Sure! I need some time to figure out why new added ceil/floor only works with one argument(it seems datafusion has an original implementation for ceil/floor), i'll do this work next week(little busy this week) -- Sorry for dragging out this PR for so long @shehabgamin @andygrove

alamb · 2025-06-18T10:11:36Z

@irenjj - I wonder if you would be willing to help pick this PR back up now that we have merged a PR with a bunch of tests from @shehabgamin here:

chore: generate basic spark function tests #16409

We are hoping to generate some example PRs to use when starting to fill out the rest of the functions and I was thinking this one might be a good one

Sure! I need some time to figure out why new added ceil/floor only works with one argument(it seems datafusion has an original implementation for ceil/floor), i'll do this work next week(little busy this week) -- Sorry for dragging out this PR for so long @shehabgamin @andygrove

Thanks @irenjj -- no worries about the timing. I know you are busy and are also on the leading edge of these spark functions (so the process of implementation is a bit rough). Thank you for helping us figure it out

github-actions · 2025-08-19T02:09:39Z

Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.

irenjj force-pushed the spart_ceil_floor branch from 58a56c9 to fd203b9 Compare May 6, 2025 08:41

Implement ceil&floor function for spark

d13258e

irenjj force-pushed the spart_ceil_floor branch from fd203b9 to d13258e Compare May 6, 2025 08:41

Merge remote-tracking branch 'upstream/main' into spart_ceil_floor

c209438

irenjj force-pushed the spart_ceil_floor branch from bf4dcbc to c209438 Compare May 6, 2025 08:49

irenjj commented May 6, 2025

View reviewed changes

fix clippy

d31e1e0

irenjj added 2 commits May 7, 2025 10:45

add return_field_from_args

5c3eb4a

add test

4009fd8

github-actions bot added the sqllogictest SQL Logic Tests (.slt) label May 7, 2025

irenjj marked this pull request as ready for review May 7, 2025 06:04

andygrove changed the title ~~Implement ceil&floor function for spark~~ [datafusion-spark] Implement ceil&floor function for spark May 7, 2025

andygrove reviewed May 8, 2025

View reviewed changes

shehabgamin reviewed May 9, 2025

View reviewed changes

alamb mentioned this pull request Jun 11, 2025

[EPIC] Complete datafusion-spark Spark Compatible Functions #15914

Open

19 tasks

github-actions bot added the Stale PR has not had any activity for some time label Aug 19, 2025

github-actions bot closed this Aug 28, 2025

codetyri0n mentioned this pull request Oct 21, 2025

Feat: [datafusion-spark] Implement ceil function. #18174

Closed

	/// What type will be returned by this function, given the arguments?
	///
	/// By default, this function calls [`Self::return_type`] with the
	/// types of each argument.
	///
	/// # Notes
	///
	/// Most UDFs should implement [`Self::return_type`] and not this
	/// function as the output type for most functions only depends on the types
	/// of their inputs (e.g. `sqrt(f32)` is always `f32`).
	///
	/// This function can be used for more advanced cases such as:
	///
	/// 1. specifying nullability
	/// 2. return types based on the values of the arguments (rather than
	/// their types.
	///
	/// # Output Type based on Values
	///
	/// For example, the following two function calls get the same argument
	/// types (something and a `Utf8` string) but return different types based
	/// on the value of the second argument:
	///
	/// * `arrow_cast(x, 'Int16')` --> `Int16`
	/// * `arrow_cast(x, 'Float32')` --> `Float32`
	///
	/// # Requirements
	///
	/// This function must consistently return the same type for the same
	/// logical input even if the input is simplified (e.g. it must return the same
	/// value for `('foo' \| 'bar')` as it does for ('foobar').
	fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<Field> {
	let data_types = args
	.arg_fields
	.iter()
	.map(\|f\| f.data_type())
	.cloned()
	.collect::<Vec<_>>();
	let return_type = self.return_type(&data_types)?;
	Ok(Field::new(self.name(), return_type, true))
	}

Conversation

irenjj commented May 6, 2025 • 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 these changes tested?

Are there any user-facing changes?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shehabgamin commented May 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shehabgamin commented May 7, 2025

Uh oh!

andygrove commented May 7, 2025

Uh oh!

andygrove commented May 7, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shehabgamin May 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shehabgamin left a comment

Choose a reason for hiding this comment

Uh oh!

shehabgamin May 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb commented Jun 17, 2025

Uh oh!

irenjj commented Jun 17, 2025

Uh oh!

alamb commented Jun 18, 2025

Uh oh!

github-actions bot commented Aug 19, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

irenjj commented May 6, 2025 •

edited

Loading

shehabgamin commented May 6, 2025 •

edited

Loading

shehabgamin May 9, 2025 •

edited

Loading

shehabgamin May 9, 2025 •

edited

Loading