[AURON #1680] initCap semantics are aligned with Spark #1681
+164
−27
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,119 @@ | ||
| // 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::sync::Arc; | ||
|
|
||
| use arrow::array::{ArrayRef, StringArray}; | ||
| use datafusion::{ | ||
| common::{Result, ScalarValue, cast::as_string_array}, | ||
| logical_expr::ColumnarValue, | ||
| }; | ||
| use datafusion_ext_commons::df_execution_err; | ||
|
|
||
| pub fn string_initcap(args: &[ColumnarValue]) -> Result<ColumnarValue> { | ||
| match &args[0] { | ||
| ColumnarValue::Array(array) => { | ||
| let input_array = as_string_array(array)?; | ||
| let output_array = | ||
| StringArray::from_iter(input_array.into_iter().map(|s| s.map(|s| initcap(s)))); | ||
| Ok(ColumnarValue::Array(Arc::new(output_array) as ArrayRef)) | ||
| } | ||
| ColumnarValue::Scalar(ScalarValue::Utf8(Some(str))) => { | ||
| Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(initcap(str))))) | ||
| } | ||
| _ => df_execution_err!("Unsupported args {args:?} for `string_initcap`"), | ||
| } | ||
| } | ||
|
|
||
| fn initcap(input: &str) -> String { | ||
| let mut out = String::with_capacity(input.len()); | ||
| let mut prev_is_space = true; // i == 0 or chars[i-1] == ' ' | ||
|
|
||
| if input.is_ascii() { | ||
| // ASCII | ||
| for ch in input.chars() { | ||
| if prev_is_space && ch.is_ascii_alphanumeric() { | ||
| out.push(ch.to_ascii_uppercase()); | ||
| } else { | ||
| out.push(ch.to_ascii_lowercase()); | ||
| }; | ||
| prev_is_space = ch == ' '; | ||
| } | ||
| } else { | ||
| // Non-ASCII | ||
| for ch in input.chars() { | ||
| if prev_is_space && ch.is_alphabetic() { | ||
| out.extend(ch.to_uppercase()); | ||
| } else { | ||
| out.extend(ch.to_lowercase()); | ||
| } | ||
| prev_is_space = ch == ' '; | ||
| } | ||
| } | ||
| out | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod test { | ||
| use std::sync::Arc; | ||
|
|
||
| use arrow::array::{ArrayRef, StringArray}; | ||
| use datafusion::{ | ||
| common::{Result, ScalarValue}, | ||
| physical_plan::ColumnarValue, | ||
| }; | ||
|
|
||
| use crate::spark_initcap::string_initcap; | ||
|
|
||
| #[test] | ||
| fn test_initcap_array() -> Result<()> { | ||
| let input_data = vec![ | ||
| None, | ||
| Some(""), | ||
| Some("hI THOmAS"), | ||
| Some("James-Smith"), | ||
| Some("michael rose"), | ||
| Some("a1b2 c3D4"), | ||
| Some(" ---abc--- ABC --ABC-- a-b A B eB Ac c d"), | ||
| Some(" 世 界 世界 "), | ||
| ]; | ||
| let input_columnar_value = ColumnarValue::Array(Arc::new(StringArray::from(input_data))); | ||
|
|
||
| let result = string_initcap(&vec![input_columnar_value])?.into_array(6)?; | ||
|
|
||
| let expected_data = vec![ | ||
| None, | ||
| Some(""), | ||
| Some("Hi Thomas"), | ||
| Some("James-smith"), | ||
| Some("Michael Rose"), | ||
| Some("A1b2 C3d4"), | ||
| Some(" ---abc--- Abc --abc-- A-b A B Eb Ac C D"), | ||
| Some(" 世 界 世界 "), | ||
| ]; | ||
| let expected: ArrayRef = Arc::new(StringArray::from(expected_data)); | ||
| assert_eq!(&result, &expected); | ||
| Ok(()) | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_initcap_scalar() -> Result<()> { | ||
| let input_columnar_value = ColumnarValue::Scalar(ScalarValue::from("abC c3D4")); | ||
| let result = string_initcap(&vec![input_columnar_value])?.into_array(1)?; | ||
| let expected: ArrayRef = Arc::new(StringArray::from(vec![Some("Abc C3d4")])); | ||
| assert_eq!(&result, &expected); | ||
| Ok(()) | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -312,37 +312,53 @@ class AuronQuerySuite extends AuronQueryTest with BaseAuronSQLSuite with AuronSQ | |
| } | ||
|
|
||
| test("initcap basic") { | ||
| withTable("initcap_basic_tbl") { | ||
| sql(s"CREATE TABLE initcap_basic_tbl(id INT, txt STRING) USING parquet") | ||
| sql(s""" | ||
| |INSERT INTO initcap_basic_tbl VALUES | ||
| | (1, 'spark sql'), | ||
| | (2, 'SPARK'), | ||
| | (3, 'sPaRk'), | ||
| | (4, ''), | ||
| | (5, NULL) | ||
| """.stripMargin) | ||
| checkSparkAnswerAndOperator("select id, initcap(txt) from initcap_basic_tbl") | ||
| } | ||
| } | ||
|
|
||
| test("initcap: word boundaries and punctuation") { | ||
| withTable("initcap_bound_tbl") { | ||
| sql(s"CREATE TABLE initcap_bound_tbl(id INT, txt STRING) USING parquet") | ||
| sql(s""" | ||
| |INSERT INTO initcap_bound_tbl VALUES | ||
| | (1, 'hello world'), | ||
| | (2, 'hello_world'), | ||
| | (3, 'über-alles'), | ||
| | (4, 'foo.bar/baz'), | ||
| | (5, 'v2Ray is COOL'), | ||
| | (6, 'rock''n''roll'), | ||
| | (7, 'hi\tthere'), | ||
| | (8, 'hi\nthere') | ||
| """.stripMargin) | ||
| checkSparkAnswerAndOperator("select id, initcap(txt) from initcap_bound_tbl") | ||
| } | ||
| } | ||
|
|
||
| test("initcap: mixed cases and edge cases") { | ||
| withTable("initcap_mixed_tbl") { | ||
| sql(s"CREATE TABLE initcap_mixed_tbl(id INT, txt STRING) USING parquet") | ||
| sql(s""" | ||
| |INSERT INTO initcap_mixed_tbl VALUES | ||
| | (1, 'a1b2 c3D4'), | ||
| | (2, '---abc--- ABC --ABC-- 世界 世 界 '), | ||
|
Contributor
There was a problem hiding this comment. Sorry about that! I think it would be better to remove the Chinese test cases here and use English instead. This way, other team members can more easily understand and verify the code during review. Thanks for understanding!
Contributor
There was a problem hiding this comment. it is ok to use chinese (or any other non-ascii characters) because we have to test with real unicode strings.
Contributor
There was a problem hiding this comment. Thanks to @richox for the feedback. From the perspective of Unicode testing, I believe this is acceptable. |
||
| | (3, ' multiple spaces '), | ||
| | (4, 'AbCdE aBcDe'), | ||
| | (5, ' A B A b '), | ||
| | (6, 'aBćDe ab世De AbĆdE aB世De ÄBĆΔE'), | ||
| | (7, 'i\u0307onic FIDELİO'), | ||
| | (8, 'a🙃B🙃c 😄 😆') | ||
| """.stripMargin) | ||
| checkSparkAnswerAndOperator("select id, initcap(txt) from initcap_mixed_tbl") | ||
| } | ||
| } | ||
|
|
||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It would better not to use parquet table to do testing, the test would fail on second run because it would create a directory on local disk, on second run, the directory name is already taken.
Use temp view or something similar would be better.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Great suggestion. We'll fix this in a follow-up PR.