Update code

Author: beliefer

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFileFormat.scala

@@ -22,14 +22,14 @@ import java.net.URI
 
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FileStatus, Path}
+import org.apache.hadoop.io.WritableComparable
 import org.apache.hadoop.mapred.JobConf
 import org.apache.hadoop.mapreduce._
 import org.apache.hadoop.mapreduce.lib.input.FileSplit
 import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl
 import org.apache.orc.{OrcUtils => _, _}
 import org.apache.orc.OrcConf.COMPRESS
 import org.apache.orc.mapred.OrcStruct
-import org.apache.orc.mapreduce._
 
 import org.apache.spark.TaskContext
 import org.apache.spark.sql.SparkSession
@@ -155,7 +155,7 @@ class OrcFileFormat
         if (orcFilterPushDown && filters.nonEmpty) {
           OrcUtils.readCatalystSchema(filePath, conf, ignoreCorruptFiles).foreach { fileSchema =>
             OrcFilters.createFilter(fileSchema, filters).foreach { f =>
-              OrcInputFormat.setSearchArgument(conf, f, fileSchema.fieldNames)
+              mapreduce.OrcInputFormat.setSearchArgument(conf, f, fileSchema.fieldNames)
             }
           }
         }
@@ -193,8 +193,8 @@ class OrcFileFormat
 
           iter.asInstanceOf[Iterator[InternalRow]]
         } else {
-          val orcRecordReader = new OrcInputFormat[OrcStruct]
-            .createRecordReader(fileSplit, taskAttemptContext)
+          val orcRecordReader: mapreduce.OrcMapreduceRecordReader[OrcStruct] =
+            createRecordReader[OrcStruct](fileSplit, taskAttemptContext)
           val iter = new RecordReaderIterator[OrcStruct](orcRecordReader)
           Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ => iter.close()))
 
@@ -214,6 +214,19 @@ class OrcFileFormat
     }
   }
 
+  private def createRecordReader[V <: WritableComparable[_]](
+      inputSplit: InputSplit,
+      taskAttemptContext: TaskAttemptContext): mapreduce.OrcMapreduceRecordReader[V] = {
+    val split = inputSplit.asInstanceOf[FileSplit]
+    val conf = taskAttemptContext.getConfiguration()
+    val readOptions = OrcFile.readerOptions(conf)
+      .maxLength(OrcConf.MAX_FILE_LENGTH.getLong(conf)).useUTCTimestamp(true)
+    val file = OrcFile.createReader(split.getPath(), readOptions)
+    val options = org.apache.orc.mapred.OrcInputFormat.buildOptions(
+      conf, file, split.getStart(), split.getLength()).useSelected(true)
+    new mapreduce.OrcMapreduceRecordReader(file, options)
+  }
+
   override def supportDataType(dataType: DataType): Boolean = dataType match {
     case _: AtomicType => true
 

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcQuerySuite.scala

@@ -21,6 +21,8 @@ import java.io.File
 import java.nio.charset.StandardCharsets
 import java.sql.Timestamp
 import java.time.{LocalDateTime, ZoneOffset}
+import java.util.TimeZone
+
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 import org.apache.hadoop.mapreduce.{JobID, TaskAttemptID, TaskID, TaskType}
@@ -30,6 +32,7 @@ import org.apache.orc.{OrcConf, OrcFile}
 import org.apache.orc.OrcConf.COMPRESS
 import org.apache.orc.mapred.OrcStruct
 import org.apache.orc.mapreduce.OrcInputFormat
+
 import org.apache.spark.{SparkConf, SparkException}
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.TableIdentifier
@@ -41,8 +44,6 @@ import org.apache.spark.sql.test.SharedSparkSession
 import org.apache.spark.sql.types._
 import org.apache.spark.util.Utils
 
-import java.util.TimeZone
-
 case class AllDataTypesWithNonPrimitiveType(
     stringField: String,
     intField: Int,
@@ -832,23 +833,27 @@ abstract class OrcQuerySuite extends OrcQueryTest with SharedSparkSession {
   }
 
   test("SPARK-37463: read/write Timestamp ntz or ltz to Orc uses UTC timestamp") {
-    TimeZone.setDefault(TimeZone.getTimeZone("America/Los_Angeles"))
-    sql("set spark.sql.session.timeZone = America/Los_Angeles")
+    val localTimeZone = TimeZone.getDefault
+    try {
+      TimeZone.setDefault(TimeZone.getTimeZone("America/Los_Angeles"))
 
-    val df =
-      sql("select timestamp_ntz '2021-06-01 00:00:00' ts_ntz, timestamp_ltz '2021-06-01 00:00:00' ts_ltz")
+      val df = sql("""
+                     |select timestamp_ntz '2021-06-01 00:00:00' ts_ntz,
+                     |timestamp_ltz '2021-06-01 00:00:00' ts_ltz
+                     |""".stripMargin)
 
-    df.write.mode("overwrite").orc("ts_ntz_orc")
-    df.write.mode("overwrite").parquet("ts_ntz_parquet")
+      df.write.mode("overwrite").orc("ts_ntz_orc")
 
-    val queryOrc = "select * from `orc`.`ts_ntz_orc`"
-    val queryParquet = "select * from `parquet`.`ts_ntz_parquet`"
+      val query = "select * from `orc`.`ts_ntz_orc`"
 
-    val tzs = Seq("America/Los_Angeles", "UTC", "Europe/Amsterdam")
-    for (tz <- tzs) {
-      TimeZone.setDefault(TimeZone.getTimeZone(tz))
-      sql(s"set spark.sql.session.timeZone = $tz")
-      sql(queryOrc).collect().equals(sql(queryParquet).collect())
+      Seq("America/Los_Angeles", "UTC", "Europe/Amsterdam").foreach { tz =>
+        TimeZone.setDefault(TimeZone.getTimeZone(tz))
+        withAllOrcReaders {
+          checkAnswer(sql(query), df)
+        }
+      }
+    } finally {
+      TimeZone.setDefault(localTimeZone)
     }
   }
 }