From 801d6a92d958f7b9762466e3c6643e54c48eb3a2 Mon Sep 17 00:00:00 2001 From: Zhenhua Wang Date: Sun, 29 Mar 2020 13:30:14 +0900 Subject: [PATCH] [SPARK-31261][SQL] Avoid npe when reading bad csv input with `columnNameCorruptRecord` specified ### What changes were proposed in this pull request? SPARK-25387 avoids npe for bad csv input, but when reading bad csv input with `columnNameCorruptRecord` specified, `getCurrentInput` is called and it still throws npe. ### Why are the changes needed? Bug fix. ### Does this PR introduce any user-facing change? No. ### How was this patch tested? Add a test. Closes #28029 from wzhfy/corrupt_column_npe. Authored-by: Zhenhua Wang Signed-off-by: HyukjinKwon (cherry picked from commit 791d2ba346f3358fc280adbbbe27f2cd50fd3732) Signed-off-by: HyukjinKwon --- .../datasources/csv/UnivocityParser.scala | 3 ++- .../sql/execution/datasources/csv/CSVSuite.scala | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/UnivocityParser.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/UnivocityParser.scala index e847e408c7f19..5579e959c47e7 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/UnivocityParser.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/UnivocityParser.scala @@ -82,7 +82,8 @@ class UnivocityParser( // Retrieve the raw record string. private def getCurrentInput: UTF8String = { - UTF8String.fromString(tokenizer.getContext.currentParsedContent().stripLineEnd) + val currentContent = tokenizer.getContext.currentParsedContent() + if (currentContent == null) null else UTF8String.fromString(currentContent.stripLineEnd) } // This parser first picks some tokens from the input tokens, according to the required schema, diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala index 2ea8f4fa414a2..866d8de2a3286 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala @@ -1822,6 +1822,20 @@ class CSVSuite extends QueryTest with SharedSQLContext with SQLTestUtils with Te assert(spark.read.csv(input).collect().toSet == Set(Row())) } + test("SPARK-31261: bad csv input with `columnNameCorruptRecord` should not cause NPE") { + val schema = StructType( + StructField("a", IntegerType) :: StructField("_corrupt_record", StringType) :: Nil) + val input = spark.createDataset(Seq("\u0000\u0000\u0001234")) + + checkAnswer( + spark.read + .option("columnNameOfCorruptRecord", "_corrupt_record") + .schema(schema) + .csv(input), + Row(null, null)) + assert(spark.read.csv(input).collect().toSet == Set(Row())) + } + test("field names of inferred schema shouldn't compare to the first row") { val input = Seq("1,2").toDS() val df = spark.read.option("enforceSchema", false).csv(input)