Fix UT

Author: HeartSaVioR

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousTextSocketSource.scala

@@ -33,7 +33,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.rpc.RpcEndpointRef
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
-import org.apache.spark.sql.catalyst.expressions.UnsafeRow
+import org.apache.spark.sql.catalyst.expressions.{UnsafeProjection, UnsafeRow}
 import org.apache.spark.sql.connector.read.InputPartition
 import org.apache.spark.sql.connector.read.streaming.{ContinuousPartitionReader, ContinuousPartitionReaderFactory, ContinuousStream, Offset, PartitionOffset}
 import org.apache.spark.sql.execution.streaming.{Offset => _, _}
@@ -246,6 +246,8 @@ class TextSocketContinuousPartitionReader(
   private var currentOffset = startOffset
   private var current: Option[InternalRow] = None
 
+  private val projectWithoutTimestamp = UnsafeProjection.create(TextSocketReader.SCHEMA_REGULAR)
+
   override def next(): Boolean = {
     try {
       current = getRecord
@@ -277,8 +279,7 @@ class TextSocketContinuousPartitionReader(
       if (includeTimestamp) {
         rec
       } else {
-        InternalRow(rec.get(0, TextSocketReader.SCHEMA_TIMESTAMP)
-          .asInstanceOf[(String, Timestamp)]._1)
+        projectWithoutTimestamp(rec)
       }
     )
 }

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/sources/TextSocketStreamSuite.scala

@@ -318,7 +318,7 @@ class TextSocketStreamSuite extends StreamTest with SharedSparkSession {
           for (i <- 0 until numRecords / 2) {
             r.next()
             offsets.append(r.getOffset().asInstanceOf[ContinuousRecordPartitionOffset].offset)
-            data.append(r.get().get(0, DataTypes.StringType).asInstanceOf[String].toInt)
+            data.append(r.get().getString(0).toInt)
             // commit the offsets in the middle and validate if processing continues
             if (i == 2) {
               commitOffset(t.partitionId, i + 1)
@@ -381,7 +381,10 @@ class TextSocketStreamSuite extends StreamTest with SharedSparkSession {
         val r = readerFactory.createReader(t).asInstanceOf[TextSocketContinuousPartitionReader]
         for (_ <- 0 until numRecords / 2) {
           r.next()
-          assert(r.get().get(0, TextSocketReader.SCHEMA_TIMESTAMP).isInstanceOf[(_, _)])
+          assert(r.get().numFields === 2)
+          // just try to read columns one by one - it would throw error if the row is corrupted
+          r.get().getString(0)
+          r.get().getLong(1)
         }
       case _ => throw new IllegalStateException("Unexpected task type")
     }