You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2024/04/17 07:47:39 UTC
(spark) branch master updated: [SPARK-47867][SQL] Support variant in JSON scan
This is an automated email from the ASF dual-hosted git repository.
wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 4913c344756d [SPARK-47867][SQL] Support variant in JSON scan
4913c344756d is described below
commit 4913c344756df8c0e382eea762386c1ec797d7f4
Author: Chenhao Li <ch...@databricks.com>
AuthorDate: Wed Apr 17 15:47:25 2024 +0800
[SPARK-47867][SQL] Support variant in JSON scan
### What changes were proposed in this pull request?
This PR adds support for the variant type in the JSON scan.
As part of this PR we introduce one new JSON option: `spark.read.format("json").option("singleVariantColumn", "colName")`. Setting this option specifies that each JSON document should be ingested into a single variant column called `colName`. When this option is used, the user must not specify a schema, and the schema is inferred as `colName variant`.
### Example 1 (multiple variant fields)
JSON files can be ingested into variant fields, e.g.
```
spark.read.format("json").schema("i int, var variant, arr ARRAY<variant>").load("a.json").show(false)
```
for a file with the following data:
```
{"i": 1, "var": {"d": "+94875-04-12", "string":"string1","int":1,"array":[1,2,3],"dict": {"key": "value1"}}, "arr": [{"a": 1}, {"b": 2}, {"c": 3, "d": [1, 2, 3]}]}
{"i": 2, "var": {"string":"string2","int":2,"array":[2,4,6],"dict": {"key": "value2"}}}
{}
{"i": 3}
```
### Example 2 (one variant field)
Here's another example with a single variant field:
```
spark.read.format("json").schema("var variant").load("a.json").show(false)
```
for a file with the following data:
```
{"var": {"d": "+94875-04-12", "string":"string1","int":1,"array":[1,2,3],"dict": {"key": "value1"}}}
{"var": {"string":"string2","int":2,"array":[2,4,6],"dict": {"key": "value2"}}}
{}
```
### Example 3 (singleVariantColumn option)
Each JSON document can also be ingested into a single variant column, e.g.
```
spark.read.format("json").option("singleVariantColumn", "var").load("a.json").show(false)
```
for a file with the following data:
```
{"i": 1, "var": {"d": "+94875-04-12", "string":"string1","int":1,"array":[1,2,3],"dict": {"key": "value1"}}, "arr": [{"a": 1}, {"b": 2}, {"c": 3, "d": [1, 2, 3]}]}
{"i": 2, "var": {"string":"string2","int":2,"array":[2,4,6],"dict": {"key": "value2"}}}
{}
{"i": 3}
```
### Why are the changes needed?
It allows Spark to ingest variant values directly from the JSON data source. Previously, the `parse_json` expression can only operate on a string column that is already in an existing table.
### Does this PR introduce _any_ user-facing change?
Yes.
### How was this patch tested?
Unit tests that verify the result and error reporting in JSON scan.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #46071 from chenhao-db/json_scan_variant.
Authored-by: Chenhao Li <ch...@databricks.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../src/main/resources/error/error-conditions.json | 6 ++
docs/sql-error-conditions.md | 6 ++
.../spark/sql/catalyst/json/JSONOptions.scala | 6 ++
.../spark/sql/catalyst/json/JacksonParser.scala | 18 +++++-
.../spark/sql/errors/QueryCompilationErrors.scala | 6 ++
.../org/apache/spark/sql/DataFrameReader.scala | 14 +++++
.../datasources/json/JsonFileFormat.scala | 8 ++-
.../scala/org/apache/spark/sql/VariantSuite.scala | 66 ++++++++++++++++++++++
.../sql/execution/datasources/json/JsonSuite.scala | 3 +-
9 files changed, 128 insertions(+), 5 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json
index e1c8c881f98f..9636ddbf73bc 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -2621,6 +2621,12 @@
],
"sqlState" : "42000"
},
+ "INVALID_SINGLE_VARIANT_COLUMN" : {
+ "message" : [
+ "The `singleVariantColumn` option cannot be used if there is also a user specified schema."
+ ],
+ "sqlState" : "42613"
+ },
"INVALID_SQL_ARG" : {
"message" : [
"The argument <name> of `sql()` is invalid. Consider to replace it either by a SQL literal or by collection constructor functions such as `map()`, `array()`, `struct()`."
diff --git a/docs/sql-error-conditions.md b/docs/sql-error-conditions.md
index 36d7cf58f09f..6b1d7c168c21 100644
--- a/docs/sql-error-conditions.md
+++ b/docs/sql-error-conditions.md
@@ -1437,6 +1437,12 @@ For more details see [INVALID_SCHEMA](sql-error-conditions-invalid-schema-error-
Expected format is 'SET', 'SET key', or 'SET key=value'. If you want to include special characters in key, or include semicolon in value, please use backquotes, e.g., SET `key`=`value`.
+### INVALID_SINGLE_VARIANT_COLUMN
+
+[SQLSTATE: 42613](sql-error-conditions-sqlstates.html#class-42-syntax-error-or-access-rule-violation)
+
+The `singleVariantColumn` option cannot be used if there is also a user specified schema.
+
### INVALID_SQL_ARG
[SQLSTATE: 42K08](sql-error-conditions-sqlstates.html#class-42-syntax-error-or-access-rule-violation)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JSONOptions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JSONOptions.scala
index e5aa0bb6d2c0..945b6e7de8b7 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JSONOptions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JSONOptions.scala
@@ -188,6 +188,11 @@ class JSONOptions(
val writeNonAsciiCharacterAsCodePoint: Boolean =
parameters.get(WRITE_NON_ASCII_CHARACTER_AS_CODEPOINT).map(_.toBoolean).getOrElse(false)
+ // This option takes in a column name and specifies that the entire JSON record should be stored
+ // as a single VARIANT type column in the table with the given column name.
+ // E.g. spark.read.format("json").option("singleVariantColumn", "colName")
+ val singleVariantColumn: Option[String] = parameters.get(SINGLE_VARIANT_COLUMN)
+
/** Build a Jackson [[JsonFactory]] using JSON options. */
def buildJsonFactory(): JsonFactory = {
val streamReadConstraints = StreamReadConstraints
@@ -282,6 +287,7 @@ object JSONOptions extends DataSourceOptions {
val COLUMN_NAME_OF_CORRUPTED_RECORD = newOption("columnNameOfCorruptRecord")
val TIME_ZONE = newOption("timeZone")
val WRITE_NON_ASCII_CHARACTER_AS_CODEPOINT = newOption("writeNonAsciiCharacterAsCodePoint")
+ val SINGLE_VARIANT_COLUMN = newOption("singleVariantColumn")
// Options with alternative
val ENCODING = "encoding"
val CHARSET = "charset"
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JacksonParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JacksonParser.scala
index d3f33a70323f..f8318aa7ce0a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JacksonParser.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/json/JacksonParser.scala
@@ -36,7 +36,8 @@ import org.apache.spark.sql.errors.{ExecutionErrors, QueryExecutionErrors}
import org.apache.spark.sql.internal.{LegacyBehaviorPolicy, SQLConf}
import org.apache.spark.sql.sources.Filter
import org.apache.spark.sql.types._
-import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String}
+import org.apache.spark.types.variant._
+import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String, VariantVal}
import org.apache.spark.util.Utils
/**
@@ -105,12 +106,25 @@ class JacksonParser(
*/
private def makeRootConverter(dt: DataType): JsonParser => Iterable[InternalRow] = {
dt match {
+ case _: StructType if options.singleVariantColumn.isDefined => (parser: JsonParser) => {
+ Some(InternalRow(parseVariant(parser)))
+ }
case st: StructType => makeStructRootConverter(st)
case mt: MapType => makeMapRootConverter(mt)
case at: ArrayType => makeArrayRootConverter(at)
}
}
+ protected final def parseVariant(parser: JsonParser): VariantVal = {
+ try {
+ val v = VariantBuilder.parseJson(parser)
+ new VariantVal(v.getValue, v.getMetadata)
+ } catch {
+ case _: VariantSizeLimitException =>
+ throw QueryExecutionErrors.variantSizeLimitError(VariantUtil.SIZE_LIMIT, "JacksonParser")
+ }
+ }
+
private def makeStructRootConverter(st: StructType): JsonParser => Iterable[InternalRow] = {
val elementConverter = makeConverter(st)
val fieldConverters = st.map(_.dataType).map(makeConverter).toArray
@@ -380,6 +394,8 @@ class JacksonParser(
case _ => null
}
+ case _: VariantType => parseVariant
+
// We don't actually hit this exception though, we keep it for understandability
case _ => throw ExecutionErrors.unsupportedDataTypeError(dataType)
}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
index 6ad5624d4730..24c85b76437c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
@@ -3227,6 +3227,12 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat
)
}
+ def invalidSingleVariantColumn(): Throwable = {
+ new AnalysisException(
+ errorClass = "INVALID_SINGLE_VARIANT_COLUMN",
+ messageParameters = Map.empty)
+ }
+
def writeWithSaveModeUnsupportedBySourceError(source: String, createMode: String): Throwable = {
new AnalysisException(
errorClass = "UNSUPPORTED_DATA_SOURCE_SAVE_MODE",
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
index 9992d8cbba07..9d7a765a24c9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameReader.scala
@@ -77,6 +77,7 @@ class DataFrameReader private[sql](sparkSession: SparkSession) extends Logging {
if (schema != null) {
val replaced = CharVarcharUtils.failIfHasCharVarchar(schema).asInstanceOf[StructType]
this.userSpecifiedSchema = Option(replaced)
+ validateSingleVariantColumn()
}
this
}
@@ -106,6 +107,7 @@ class DataFrameReader private[sql](sparkSession: SparkSession) extends Logging {
*/
def option(key: String, value: String): DataFrameReader = {
this.extraOptions = this.extraOptions + (key -> value)
+ validateSingleVariantColumn()
this
}
@@ -149,6 +151,7 @@ class DataFrameReader private[sql](sparkSession: SparkSession) extends Logging {
*/
def options(options: scala.collection.Map[String, String]): DataFrameReader = {
this.extraOptions ++= options
+ validateSingleVariantColumn()
this
}
@@ -766,6 +769,17 @@ class DataFrameReader private[sql](sparkSession: SparkSession) extends Logging {
}
}
+ /**
+ * Ensure that the `singleVariantColumn` option cannot be used if there is also a user specified
+ * schema.
+ */
+ private def validateSingleVariantColumn(): Unit = {
+ if (extraOptions.get(JSONOptions.SINGLE_VARIANT_COLUMN).isDefined &&
+ userSpecifiedSchema.isDefined) {
+ throw QueryCompilationErrors.invalidSingleVariantColumn()
+ }
+ }
+
///////////////////////////////////////////////////////////////////////////////////////
// Builder pattern config options
///////////////////////////////////////////////////////////////////////////////////////
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonFileFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonFileFormat.scala
index 7fb6e98fb046..6174c017f604 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonFileFormat.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonFileFormat.scala
@@ -55,8 +55,10 @@ class JsonFileFormat extends TextBasedFileFormat with DataSourceRegister {
options,
sparkSession.sessionState.conf.sessionLocalTimeZone,
sparkSession.sessionState.conf.columnNameOfCorruptRecord)
- JsonDataSource(parsedOptions).inferSchema(
- sparkSession, files, parsedOptions)
+ parsedOptions.singleVariantColumn match {
+ case Some(columnName) => Some(StructType(Array(StructField(columnName, VariantType))))
+ case None => JsonDataSource(parsedOptions).inferSchema(sparkSession, files, parsedOptions)
+ }
}
override def prepareWrite(
@@ -134,7 +136,7 @@ class JsonFileFormat extends TextBasedFileFormat with DataSourceRegister {
override def equals(other: Any): Boolean = other.isInstanceOf[JsonFileFormat]
override def supportDataType(dataType: DataType): Boolean = dataType match {
- case _: VariantType => false
+ case _: VariantType => true
case _: AtomicType => true
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala
index d276ec4428b9..d2d12920b68a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala
@@ -18,6 +18,8 @@
package org.apache.spark.sql
import java.io.File
+import java.nio.charset.StandardCharsets
+import java.nio.file.Files
import scala.collection.mutable
import scala.jdk.CollectionConverters._
@@ -272,6 +274,70 @@ class VariantSuite extends QueryTest with SharedSparkSession {
}
}
+ test("json option constraints") {
+ withTempDir { dir =>
+ val file = new File(dir, "file.json")
+ Files.write(file.toPath, "0".getBytes(StandardCharsets.UTF_8))
+
+ // Ensure that we get an error when setting the singleVariantColumn JSON option while also
+ // specifying a schema.
+ checkError(
+ exception = intercept[AnalysisException] {
+ spark.read.format("json").option("singleVariantColumn", "var").schema("var variant")
+ },
+ errorClass = "INVALID_SINGLE_VARIANT_COLUMN",
+ parameters = Map.empty
+ )
+ checkError(
+ exception = intercept[AnalysisException] {
+ spark.read.format("json").option("singleVariantColumn", "another_name")
+ .schema("var variant").json(file.getAbsolutePath).collect()
+ },
+ errorClass = "INVALID_SINGLE_VARIANT_COLUMN",
+ parameters = Map.empty
+ )
+ }
+ }
+
+ test("json scan") {
+ val content = Seq(
+ "true",
+ """{"a": [], "b": null}""",
+ """{"a": 1}""",
+ "[1, 2, 3]"
+ ).mkString("\n").getBytes(StandardCharsets.UTF_8)
+
+ withTempDir { dir =>
+ val file = new File(dir, "file.json")
+ Files.write(file.toPath, content)
+
+ checkAnswer(
+ spark.read.format("json").option("singleVariantColumn", "var")
+ .load(file.getAbsolutePath)
+ .selectExpr("to_json(var)"),
+ Seq(Row("true"), Row("""{"a":[],"b":null}"""), Row("""{"a":1}"""), Row("[1,2,3]"))
+ )
+
+ checkAnswer(
+ spark.read.format("json").schema("a variant, b variant")
+ .load(file.getAbsolutePath).selectExpr("to_json(a)", "to_json(b)"),
+ Seq(Row(null, null), Row("[]", "null"), Row("1", null), Row(null, null))
+ )
+ }
+
+ // Test scan with partitions.
+ withTempDir { dir =>
+ new File(dir, "a=1/b=2/").mkdirs()
+ Files.write(new File(dir, "a=1/b=2/file.json").toPath, content)
+ checkAnswer(
+ spark.read.format("json").option("singleVariantColumn", "var")
+ .load(dir.getAbsolutePath).selectExpr("a", "b", "to_json(var)"),
+ Seq(Row(1, 2, "true"), Row(1, 2, """{"a":[],"b":null}"""), Row(1, 2, """{"a":1}"""),
+ Row(1, 2, "[1,2,3]"))
+ )
+ }
+ }
+
test("group/order/join variant are disabled") {
var ex = intercept[AnalysisException] {
spark.sql("select parse_json('') group by 1")
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
index f3c332bab183..c17a25be8e2a 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala
@@ -3728,7 +3728,7 @@ abstract class JsonSuite
}
test("SPARK-40667: validate JSON Options") {
- assert(JSONOptions.getAllOptions.size == 28)
+ assert(JSONOptions.getAllOptions.size == 29)
// Please add validation on any new Json options here
assert(JSONOptions.isValidOption("samplingRatio"))
assert(JSONOptions.isValidOption("primitivesAsString"))
@@ -3756,6 +3756,7 @@ abstract class JsonSuite
assert(JSONOptions.isValidOption("columnNameOfCorruptRecord"))
assert(JSONOptions.isValidOption("timeZone"))
assert(JSONOptions.isValidOption("writeNonAsciiCharacterAsCodePoint"))
+ assert(JSONOptions.isValidOption("singleVariantColumn"))
assert(JSONOptions.isValidOption("encoding"))
assert(JSONOptions.isValidOption("charset"))
// Please add validation on any new Json options with alternative here
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org