Skip to content

[SPARK-61177][CONNECT] Quote the field names of DataType::to_ddl, and canonicalize the collation name of DataType::from_ddl, in the Rust client #374

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
HyukjinKwon wants to merge 1 commit into from
+329 −16
Closed

[SPARK-61177][CONNECT] Quote the field names of DataType::to_ddl, and canonicalize the collation name of DataType::from_ddl, in the Rust client #374

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 105 additions & 15 deletions crates/spark-connect/src/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -920,20 +920,19 @@ impl DataType {
}

/// DDL string for a StructType, mirroring `StructType.toDDL()`:
/// comma-separated `name type[ NOT NULL][ COMMENT '...']` per field.
/// comma-separated `name type[ NOT NULL][ COMMENT '...']` per field. The names are quoted
/// and the types written as `DataType.sql`, so the server parses the DDL back.
pub fn to_ddl(&self) -> Result<String> {
match self {
DataType::Struct { fields } => Ok(fields
.iter()
.map(|f| {
let mut s = format!("{} {}", f.name, f.data_type.simple_string());
if !f.nullable {
s.push_str(" NOT NULL");
}
if let Some(comment) = f.metadata.get("comment").and_then(|c| c.as_str()) {
s.push_str(&format!(" COMMENT '{}'", comment.replace('\'', "\\'")));
}
s
format!(
"{} {}{}",
quote_if_needed(&f.name),
f.data_type.sql_type(),
field_ddl_suffix(f)
)
})
.collect::<Vec<_>>()
.join(",")),
Expand All @@ -944,6 +943,39 @@ impl DataType {
}
}

/// The SQL form of a type, mirroring `DataType.sql`: `INT`, `STRING COLLATE UTF8_LCASE`,
/// `ARRAY<...>`, `MAP<K, V>` and `STRUCT<name: TYPE, ...>`, the form `toDDL` writes.
fn sql_type(&self) -> String {
match self {
DataType::Struct { fields } => {
let fields: Vec<String> = fields
.iter()
.map(|f| {
format!(
"{}: {}{}",
quote_if_needed(&f.name),
f.data_type.sql_type(),
field_ddl_suffix(f)
)
})
.collect();
format!("STRUCT<{}>", fields.join(", "))
}
DataType::Array { element_type, .. } => format!("ARRAY<{}>", element_type.sql_type()),
DataType::Map {
key_type,
value_type,
..
} => format!("MAP<{}, {}>", key_type.sql_type(), value_type.sql_type()),
DataType::String { collation }
if !(collation.is_empty() || collation == "UTF8_BINARY") =>
{
format!("STRING COLLATE {}", collation.to_uppercase())
}
_ => self.simple_string().to_uppercase(),
}
}

/// Tree-string for a StructType, mirroring `StructType.treeString()`.
pub fn tree_string(&self) -> Result<String> {
self.tree_string_with_depth(i32::MAX)
Expand Down Expand Up @@ -1176,6 +1208,32 @@ fn interval_string(start_field: i32, end_field: i32, fields: &[(&str, i32)]) ->
}
}

/// Quotes a field name with backticks, as Spark's `QuotingUtils.quoteIfNeeded` does: a name
/// that is not `[a-zA-Z_][a-zA-Z0-9_]*` is quoted, and each backtick in it is doubled.
fn quote_if_needed(name: &str) -> String {
let mut chars = name.chars();
let plain = matches!(chars.next(), Some(c) if c == '_' || c.is_ascii_alphabetic())
&& chars.all(|c| c == '_' || c.is_ascii_alphanumeric());
if plain {
name.to_string()
} else {
format!("`{}`", name.replace('`', "``"))
}
}

/// The ` NOT NULL` and ` COMMENT '...'` suffix of a field, as `StructField.toDDL` writes it. The
/// single quotes of the comment are escaped as Spark's `escapeSingleQuotedString` does.
fn field_ddl_suffix(field: &StructField) -> String {
let mut s = String::new();
if !field.nullable {
s.push_str(" NOT NULL");
}
if let Some(comment) = field.metadata.get("comment").and_then(|c| c.as_str()) {
s.push_str(&format!(" COMMENT '{}'", comment.replace('\'', "\\'")));
}
s
}

/// The key of the metadata of a struct field of the collations of its strings, as Spark's.
const COLLATIONS_METADATA_KEY: &str = "__COLLATIONS";

Expand Down Expand Up @@ -1907,6 +1965,35 @@ fn split_top_level_comma(s: &str) -> Vec<&str> {
result
}

/// The canonical form of a collation name, as Spark's server gives it. A `UTF8_*` or `UNICODE*`
/// name is in upper case. An ICU locale name, `language[_Script][_COUNTRY][_SPECIFIERS]`, has a
/// lower-case language, a title-case script (a four-letter part) and upper-case other parts, such
/// as `sr_Cyrl_SRB_CI_AI`. Collation names are case-insensitive, so any case is accepted.
fn canonical_collation_name(name: &str) -> String {
let provider = name.split('_').next().unwrap_or_default();
if provider.eq_ignore_ascii_case("UTF8") || provider.eq_ignore_ascii_case("UNICODE") {
return name.to_uppercase();
}
name.split('_')
.enumerate()
.map(|(i, part)| {
if i == 0 {
part.to_lowercase()
} else if part.len() == 4 && part.chars().all(|c| c.is_ascii_alphabetic()) {
let mut chars = part.chars();
let first = chars.next().map(|c| c.to_ascii_uppercase());
first
.into_iter()
.chain(chars.map(|c| c.to_ascii_lowercase()))
.collect()
} else {
part.to_uppercase()
}
})
.collect::<Vec<_>>()
.join("_")
}

/// Parse a single type (primitive or complex)
fn parse_single_type(input: &str) -> Result<DataType> {
let trimmed = input.trim().to_lowercase();
Expand Down Expand Up @@ -1959,12 +2046,15 @@ fn parse_single_type(input: &str) -> Result<DataType> {
} else if trimmed.starts_with("interval") {
parse_interval_string(&trimmed)
} else if trimmed.starts_with("string collate") {
let collation = trimmed
.trim_start_matches("string")
.trim_start_matches("collate")
.trim()
.to_string();
Ok(DataType::String { collation })
// The name is taken from the input, not from `trimmed`, which is lowercased.
let collation = input
.split_whitespace()
.skip(2)
.collect::<Vec<_>>()
.join(" ");
Ok(DataType::String {
collation: canonical_collation_name(&collation),
})
} else if trimmed.starts_with("geometry(") {
parse_geometry(&trimmed)
} else if trimmed.starts_with("geography(") {
Expand Down
225 changes: 224 additions & 1 deletion crates/spark-connect/tests/types_behavior.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,7 +277,7 @@ fn struct_to_ddl_tree_and_nullable() {
},
],
};
assert_eq!(st.to_ddl().unwrap(), "a int,b string NOT NULL");
assert_eq!(st.to_ddl().unwrap(), "a INT,b STRING NOT NULL");
let tree = st.tree_string().unwrap();
assert!(tree.starts_with("root\n"));
assert!(tree.contains("a: integer (nullable = true)"));
Expand All @@ -293,6 +293,229 @@ fn struct_to_ddl_tree_and_nullable() {
assert!(DataType::Integer.tree_string().is_err());
}

/// The expected strings are the output of `StructType.toDDL()` of PySpark 4.3.0 on a server.
#[test]
fn struct_to_ddl_quotes_names_and_writes_sql_types() {
use spark_connect::types::StructField;
use std::collections::BTreeMap;

fn field(name: &str, data_type: DataType, nullable: bool) -> StructField {
StructField {
name: name.to_string(),
data_type,
nullable,
metadata: BTreeMap::new(),
}
}
fn commented(name: &str, data_type: DataType, nullable: bool, comment: &str) -> StructField {
let mut f = field(name, data_type, nullable);
f.metadata
.insert("comment".to_string(), serde_json::json!(comment));
f
}
fn ddl(fields: Vec<StructField>) -> String {
DataType::Struct { fields }.to_ddl().unwrap()
}
let int = || DataType::Integer;
let string = || DataType::String {
collation: "UTF8_BINARY".to_string(),
};

// Names are quoted unless they are `[a-zA-Z_][a-zA-Z0-9_]*`; a backtick is doubled.
assert_eq!(
ddl(vec![
field("my col", int(), true),
field("a`b", int(), true),
field("a.b", int(), true),
field("1abc", int(), true),
field("123", int(), true),
field("_x", int(), true),
field("ünï", int(), true),
]),
"`my col` INT,`a``b` INT,`a.b` INT,`1abc` INT,`123` INT,_x INT,`ünï` INT"
);
assert_eq!(
ddl(vec![commented("x", int(), false, "it's")]),
"x INT NOT NULL COMMENT 'it\\'s'"
);
// Nested struct fields are written as `name: TYPE`, with the same quoting and suffixes.
assert_eq!(
ddl(vec![field(
"s",
DataType::Struct {
fields: vec![commented("my col", int(), false, "c")],
},
true,
)]),
"s STRUCT<`my col`: INT NOT NULL COMMENT 'c'>"
);
assert_eq!(
ddl(vec![field(
"arr",
DataType::Array {
element_type: Box::new(DataType::Struct {
fields: vec![field("a b", int(), true)],
}),
contains_null: true,
},
true,
)]),
"arr ARRAY<STRUCT<`a b`: INT>>"
);
assert_eq!(
ddl(vec![field(
"m",
DataType::Map {
key_type: Box::new(string()),
value_type: Box::new(DataType::Struct {
fields: vec![field("k`v", DataType::Long, false)],
}),
value_contains_null: true,
},
true,
)]),
"m MAP<STRING, STRUCT<`k``v`: BIGINT NOT NULL>>"
);
assert_eq!(
ddl(vec![field(
"a",
DataType::Array {
element_type: Box::new(DataType::Struct {
fields: vec![commented("q", string(), true, "x'y")],
}),
contains_null: true,
},
true,
)]),
"a ARRAY<STRUCT<q: STRING COMMENT 'x\\'y'>>"
);
// Atomic types are written as `DataType.sql`.
assert_eq!(
ddl(vec![
field(
"d",
DataType::Decimal {
precision: 10,
scale: 2
},
true
),
field("ntz", DataType::TimestampNtz, true),
field(
"ym",
DataType::YearMonthInterval {
start_field: 0,
end_field: 1
},
true
),
field(
"dt",
DataType::DayTimeInterval {
start_field: 0,
end_field: 3
},
true
),
field("ch", DataType::Char { length: 3 }, true),
field("vc", DataType::Varchar { length: 4 }, true),
field("tm", DataType::Time { precision: 3 }, true),
field("bin", DataType::Binary, true),
field("b", DataType::Boolean, true),
field("v", DataType::Variant, true),
]),
"d DECIMAL(10,2),ntz TIMESTAMP_NTZ,ym INTERVAL YEAR TO MONTH,dt INTERVAL DAY TO SECOND,\
ch CHAR(3),vc VARCHAR(4),tm TIME(3),bin BINARY,b BOOLEAN,v VARIANT"
);
// A non-default collation is written as `STRING COLLATE <name>`.
assert_eq!(
ddl(vec![field(
"c",
DataType::String {
collation: "UTF8_LCASE".to_string(),
},
true,
)]),
"c STRING COLLATE UTF8_LCASE"
);
assert_eq!(
ddl(vec![field(
"s",
DataType::Struct {
fields: vec![field(
"c d",
DataType::String {
collation: "UNICODE_CI".to_string(),
},
true,
)],
},
true,
)]),
"s STRUCT<`c d`: STRING COLLATE UNICODE_CI>"
);
}

#[test]
fn from_ddl_canonicalizes_the_collation_name() {
// The canonical names of the server of PySpark 4.3.0, from `createDataFrame([], "c STRING
// COLLATE <name>")`. `en_us` is rejected by the server (COLLATION_INVALID_NAME), so `from_ddl`
// only canonicalizes its case, `en_US`.
let cases = [
("utf8_lcase", "UTF8_LCASE"),
("Unicode_Ci_Ai", "UNICODE_CI_AI"),
("unicode_rtrim", "UNICODE_RTRIM"),
("utf8_lcase_rtrim", "UTF8_LCASE_RTRIM"),
("sr_ci_ai", "sr_CI_AI"),
("SR_CI_AI", "sr_CI_AI"),
("sr_CI_AI", "sr_CI_AI"),
("EN", "en"),
("de_ci", "de_CI"),
("zh_hant_hkg", "zh_Hant_HKG"),
("SR_CYRL_SRB_CI_AI", "sr_Cyrl_SRB_CI_AI"),
("en_us", "en_US"),
];
for (written, canonical) in cases {
assert_eq!(
DataType::from_ddl(&format!("STRING COLLATE {written}")).unwrap(),
DataType::String {
collation: canonical.to_string(),
},
"STRING COLLATE {written}"
);
}
match DataType::from_ddl("a STRING COLLATE UTF8_LCASE").unwrap() {
DataType::Struct { fields } => assert_eq!(
fields[0].data_type,
DataType::String {
collation: "UTF8_LCASE".to_string(),
}
),
other => panic!("expected struct, got {other:?}"),
}
assert_eq!(
DataType::from_ddl("array<string collate UNICODE_CI>").unwrap(),
DataType::Array {
element_type: Box::new(DataType::String {
collation: "UNICODE_CI".to_string(),
}),
contains_null: true,
}
);
// The DDL written by `to_ddl` parses back to the same collation.
let st = DataType::Struct {
fields: vec![spark_connect::types::StructField {
name: "a".to_string(),
data_type: DataType::String {
collation: "UTF8_LCASE".to_string(),
},
nullable: true,
metadata: Default::default(),
}],
};
assert_eq!(DataType::from_ddl(&st.to_ddl().unwrap()).unwrap(), st);
}

#[test]
fn tree_string_with_depth_truncates_nested_structs() {
use spark_connect::types::StructField;
Expand Down
Loading