This is an automated email from the ASF dual-hosted git repository.
richox 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 15391928 [AURON #1518] Implement native function of quarter. (#1519)
15391928 is described below
commit 153919283ef4a3ded80852564c14694544c51f7e
Author: slfan1989 <[email protected]>
AuthorDate: Mon Nov 3 10:53:34 2025 +0800
[AURON #1518] Implement native function of quarter. (#1519)
* [AURON #1518] Implement native function of quarter.
Signed-off-by: slfan1989 <[email protected]>
* [AURON #1518] Fix CheckStyle Issue.
Signed-off-by: slfan1989 <[email protected]>
---------
Signed-off-by: slfan1989 <[email protected]>
---
native-engine/datafusion-ext-functions/src/lib.rs | 1 +
.../datafusion-ext-functions/src/spark_dates.rs | 71 ++++++++++++++++++++++
.../apache/spark/sql/auron/AuronQuerySuite.scala | 22 +++++++
.../apache/spark/sql/auron/NativeConverters.scala | 1 +
4 files changed, 95 insertions(+)
diff --git a/native-engine/datafusion-ext-functions/src/lib.rs
b/native-engine/datafusion-ext-functions/src/lib.rs
index f2311f37..8c2f7560 100644
--- a/native-engine/datafusion-ext-functions/src/lib.rs
+++ b/native-engine/datafusion-ext-functions/src/lib.rs
@@ -60,6 +60,7 @@ pub fn create_spark_ext_function(name: &str) ->
Result<ScalarFunctionImplementat
"Year" => Arc::new(spark_dates::spark_year),
"Month" => Arc::new(spark_dates::spark_month),
"Day" => Arc::new(spark_dates::spark_day),
+ "Quarter" => Arc::new(spark_dates::spark_quarter),
"BrickhouseArrayUnion" =>
Arc::new(brickhouse::array_union::array_union),
"Round" => Arc::new(spark_round::spark_round),
"NormalizeNanAndZero" => {
diff --git a/native-engine/datafusion-ext-functions/src/spark_dates.rs
b/native-engine/datafusion-ext-functions/src/spark_dates.rs
index 6bb99555..6221d9ab 100644
--- a/native-engine/datafusion-ext-functions/src/spark_dates.rs
+++ b/native-engine/datafusion-ext-functions/src/spark_dates.rs
@@ -13,7 +13,10 @@
// See the License for the specific language governing permissions and
// limitations under the License.
+use std::sync::Arc;
+
use arrow::{
+ array::{ArrayRef, Int32Array},
compute::{DatePart, date_part},
datatypes::DataType,
};
@@ -35,6 +38,33 @@ pub fn spark_day(args: &[ColumnarValue]) ->
Result<ColumnarValue> {
Ok(ColumnarValue::Array(date_part(&input, DatePart::Day)?))
}
+/// `spark_quarter(date/timestamp/compatible-string)`
+///
+/// Simulates Spark's `quarter()` function.
+/// Converts the input to `Date32`, extracts the month (1–12),
+/// and computes the quarter as `((month - 1) / 3) + 1`.
+/// Null values are propagated transparently.
+pub fn spark_quarter(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ // Cast input to Date32 for compatibility with date_part()
+ let input = cast(&args[0].clone().into_array(1)?, &DataType::Date32)?;
+
+ // Extract month (1–12) using Arrow's date_part
+ let month_arr: ArrayRef = date_part(&input, DatePart::Month)?;
+ let month_arr = month_arr
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .expect("date_part(Month) must return Int32Array");
+
+ // Compute quarter: ((month - 1) / 3) + 1, preserving NULLs
+ let quarter = Int32Array::from_iter(
+ month_arr
+ .iter()
+ .map(|opt_m| opt_m.map(|m| ((m - 1) / 3 + 1) as i32)),
+ );
+
+ Ok(ColumnarValue::Array(Arc::new(quarter)))
+}
+
#[cfg(test)]
mod tests {
use std::sync::Arc;
@@ -100,4 +130,45 @@ mod tests {
&expected_ret
);
}
+
+ #[test]
+ fn test_spark_quarter_basic() {
+ // Date32 days relative to 1970-01-01:
+ // 0 -> 1970-01-01 (Q1)
+ // 40 -> ~1970-02-10 (Q1)
+ // 100 -> ~1970-04-11 (Q2)
+ // 200 -> ~1970-07-20 (Q3)
+ // 300 -> ~1970-10-28 (Q4)
+ let input = Arc::new(Date32Array::from(vec![
+ Some(0),
+ Some(40),
+ Some(100),
+ Some(200),
+ Some(300),
+ None,
+ ]));
+ let args = vec![ColumnarValue::Array(input)];
+ let expected: ArrayRef = Arc::new(Int32Array::from(vec![
+ Some(1),
+ Some(1),
+ Some(2),
+ Some(3),
+ Some(4),
+ None,
+ ]));
+
+ let out = spark_quarter(&args).unwrap().into_array(1).unwrap();
+ assert_eq!(&out, &expected);
+ }
+
+ #[test]
+ fn test_spark_quarter_null_only() {
+ // Ensure NULL propagation
+ let input = Arc::new(Date32Array::from(vec![None, None]));
+ let args = vec![ColumnarValue::Array(input)];
+ let expected: ArrayRef = Arc::new(Int32Array::from(vec![None, None]));
+
+ let out = spark_quarter(&args).unwrap().into_array(1).unwrap();
+ assert_eq!(&out, &expected);
+ }
}
diff --git
a/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/auron/AuronQuerySuite.scala
b/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/auron/AuronQuerySuite.scala
index 4b1fec6e..d36fd096 100644
---
a/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/auron/AuronQuerySuite.scala
+++
b/spark-extension-shims-spark/src/test/scala/org/apache/spark/sql/auron/AuronQuerySuite.scala
@@ -287,6 +287,28 @@ class AuronQuerySuite
}
}
+ test("test filter with quarter function") {
+ withTable("t1") {
+ sql("""
+ |create table t1 using parquet as
+ |select '2024-02-10' as event_time
+ |union all select '2024-04-11'
+ |union all select '2024-07-20'
+ |union all select '2024-12-18'
+ |""".stripMargin)
+
+ checkAnswer(
+ sql("""
+ |select q, count(*)
+ |from (select event_time, quarter(event_time) as q from t1) t
+ |where q <= 3
+ |group by q
+ |order by q
+ |""".stripMargin),
+ Seq(Row(1, 1), Row(2, 1), Row(3, 1)))
+ }
+ }
+
test("lpad/rpad basic") {
Seq(
("select lpad('abc', 5, '*')", Row("**abc")),
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 d8418b17..8f0e90e1 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
@@ -879,6 +879,7 @@ object NativeConverters extends Logging {
case Year(child) => buildExtScalarFunction("Year", child :: Nil,
IntegerType)
case Month(child) => buildExtScalarFunction("Month", child :: Nil,
IntegerType)
case DayOfMonth(child) => buildExtScalarFunction("Day", child :: Nil,
IntegerType)
+ case Quarter(child) => buildExtScalarFunction("Quarter", child :: Nil,
IntegerType)
case e: Levenshtein =>
buildScalarFunction(pb.ScalarFunction.Levenshtein, e.children,
e.dataType)
