Skip to content
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.

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

Escape quotes and newlines when converting strings to json format in to_json #9612

Merged
merged 9 commits into from
Nov 14, 2023
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
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
2 changes: 0 additions & 2 deletions docs/compatibility.md
Original file line number Diff line number Diff line change
Expand Up @@ -342,8 +342,6 @@ with Spark, and can be enabled by setting `spark.rapids.sql.expression.StructsTo

Known issues are:

- String escaping is not implemented, so strings containing quotes, newlines, and other special characters will
not produce valid JSON
- There is no support for timestamp types
- There can be rounding differences when formatting floating-point numbers as strings. For example, Spark may
produce `-4.1243574E26` but the GPU may produce `-4.124357351E26`.
Expand Down
10 changes: 8 additions & 2 deletions integration_tests/src/main/python/json_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -575,8 +575,14 @@ def test_read_case_col_name(spark_tmp_path, v1_enabled_list, col_name):
pytest.param(double_gen, marks=pytest.mark.xfail(reason='https://github.com/NVIDIA/spark-rapids/issues/9350')),
pytest.param(date_gen, marks=pytest.mark.xfail(reason='https://github.com/NVIDIA/spark-rapids/issues/9515')),
pytest.param(timestamp_gen, marks=pytest.mark.xfail(reason='https://github.com/NVIDIA/spark-rapids/issues/9515')),
StringGen('[A-Za-z0-9]{0,10}', nullable=True),
pytest.param(StringGen(nullable=True), marks=pytest.mark.xfail(reason='https://github.com/NVIDIA/spark-rapids/issues/9514')),
StringGen('[A-Za-z0-9\r\n\'"\\\\]{0,10}', nullable=True) \
.with_special_case('\u1f600') \
.with_special_case('"a"') \
.with_special_case('\\"a\\"') \
.with_special_case('\'a\'') \
.with_special_case('\\\'a\\\''),
pytest.param(StringGen('\u001a', nullable=True), marks=pytest.mark.xfail(
reason='cuDF represents two-digit unicode characters in hex format such as \x1a'))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a cudf issue to track this which we could link from here?

Maybe we can use something like this https://github.com/rapidsai/cudf/blob/f97e74f00b7a6bac37c9603def95a11b06cb013f/cpp/src/io/json/write_json.cu#L95-L108

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think that this is something that users will encounter - this is specific to escaping non-printable ASCII characters, and it would seem unusual to want to represent those in JSON. I wanted to note this issue to explain why we can't do a full StringGen in the tests.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@gerashegalov I added a link to a follow-up issue #9705

], ids=idfn)
@pytest.mark.parametrize('ignore_null_fields', [
True,
Expand Down
30 changes: 21 additions & 9 deletions sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCast.scala
Original file line number Diff line number Diff line change
Expand Up @@ -898,7 +898,10 @@ object GpuCast {

val numRows = input.getRowCount.toInt

/** Create a new column with quotes around the supplied string column */
/**
* Create a new column with quotes around the supplied string column. Caller
* is responsible for closing `column`.
*/
def addQuotes(column: ColumnVector, rowCount: Int): ColumnVector = {
withResource(ArrayBuffer.empty[ColumnVector]) { columns =>
withResource(Scalar.fromString("\"")) { quote =>
Expand All @@ -921,7 +924,7 @@ object GpuCast {
// keys must have quotes around them in JSON mode
val strKey: ColumnVector = withResource(kvStructColumn.getChildColumnView(0)) { keyColumn =>
withResource(castToString(keyColumn, from.keyType, options)) { key =>
addQuotes(key.incRefCount(), keyColumn.getRowCount.toInt)
addQuotes(key, keyColumn.getRowCount.toInt)
}
}
// string values must have quotes around them in JSON mode, and null values need
Expand All @@ -930,7 +933,7 @@ object GpuCast {
withResource(kvStructColumn.getChildColumnView(1)) { valueColumn =>
val valueStr = if (valueColumn.getType == DType.STRING) {
withResource(castToString(valueColumn, from.valueType, options)) { valueStr =>
addQuotes(valueStr.incRefCount(), valueColumn.getRowCount.toInt)
addQuotes(valueStr, valueColumn.getRowCount.toInt)
}
} else {
castToString(valueColumn, from.valueType, options)
Expand Down Expand Up @@ -1107,13 +1110,15 @@ object GpuCast {
attrColumns += colon.incRefCount()
}
// write the value
val attrValue = castToString(cv, inputSchema(fieldIndex).dataType, options)
if (needsQuoting) {
attrColumns += quote.incRefCount()
attrColumns += escapeJsonString(attrValue)
withResource(castToString(cv, inputSchema(fieldIndex).dataType, options)) {
attrValue =>
attrColumns += escapeJsonString(attrValue)
}
attrColumns += quote.incRefCount()
} else {
attrColumns += attrValue
attrColumns += castToString(cv, inputSchema(fieldIndex).dataType, options)
}
// now concatenate
val jsonAttr = withResource(Scalar.fromString("")) { emptyString =>
Expand Down Expand Up @@ -1168,10 +1173,17 @@ object GpuCast {
}
}

/**
* Escape quotes and newlines in a string column. Caller is responsible for closing `cv`.
*/
private def escapeJsonString(cv: ColumnVector): ColumnVector = {
// this is a placeholder for implementing string escaping
// https://github.com/NVIDIA/spark-rapids/issues/9514
cv
val chars = Seq("\r", "\n", "\\", "\"")
val escaped = chars.map(StringEscapeUtils.escapeJava)
withResource(ColumnVector.fromStrings(chars: _*)) { search =>
withResource(ColumnVector.fromStrings(escaped: _*)) { replace =>
cv.stringReplace(search, replace)
}
}
}

private[rapids] def castFloatingTypeToString(input: ColumnView): ColumnVector = {
Expand Down