Address Codegen support

Author: HeartSaVioR

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

@@ -25,8 +25,8 @@ import scala.collection.JavaConverters._
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.errors._
-import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateUnsafeProjection, Predicate}
+import org.apache.spark.sql.catalyst.expressions.{Predicate => _, _}
+import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.plans.logical.EventTimeWatermark
 import org.apache.spark.sql.catalyst.plans.physical.{AllTuples, ClusteredDistribution, Distribution, Partitioning}
 import org.apache.spark.sql.catalyst.streaming.InternalOutputModes._
@@ -40,7 +40,7 @@ import org.apache.spark.util.{CompletionIterator, NextIterator, Utils}
 case class CountLateRowsExec(
     eventTimeWatermark: Option[Long] = None,
     child: SparkPlan)
-  extends UnaryExecNode with WatermarkSupport {
+  extends UnaryExecNode with WatermarkSupport with CodegenSupport {
 
   // No need to determine key expressions here.
   override def keyExpressions: Seq[Attribute] = Seq.empty
@@ -70,6 +70,50 @@ case class CountLateRowsExec(
       }
     }
   }
+
+  override def inputRDDs(): Seq[RDD[InternalRow]] = {
+    child.asInstanceOf[CodegenSupport].inputRDDs()
+  }
+
+  override protected def doProduce(ctx: CodegenContext): String = {
+    child.asInstanceOf[CodegenSupport].produce(ctx, this)
+  }
+
+  override def doConsume(ctx: CodegenContext, input: Seq[ExprCode], row: ExprCode): String = {
+    val numLateRows = metricTerm(ctx, "numLateRows")
+
+    val generated = watermarkExpression match {
+      case Some(expr) =>
+        val bound = BindReferences.bindReference(expr, child.output)
+        val evaluated = evaluateRequiredVariables(child.output, input, expr.references)
+
+        // Generate the code for the predicate.
+        val ev = ExpressionCanonicalizer.execute(bound).genCode(ctx)
+        val nullCheck = if (bound.nullable) {
+          s"${ev.isNull} || "
+        } else {
+          s""
+        }
+
+        s"""
+           |$evaluated
+           |${ev.code}
+           |if (${nullCheck}${ev.value}) {
+           |  $numLateRows.add(1);
+           |}
+       """.stripMargin
+
+      case None => ""
+    }
+
+    // Note: wrap in "do { } while(false);", so the generated checks can jump out with "continue;"
+    s"""
+       |do {
+       |  $generated
+       |  ${consume(ctx, input)}
+       |} while(false);
+     """.stripMargin
+  }
 }
 
 /** Used to identify the state store for a given operator. */

sql/core/src/test/scala/org/apache/spark/sql/streaming/EventTimeWatermarkSuite.scala

@@ -500,15 +500,19 @@ class EventTimeWatermarkSuite extends StreamTest with BeforeAndAfter with Matche
     testStream(windowedAggregation, OutputMode.Complete)(
       AddData(inputData, 10, 11, 12),
       CheckAnswer((10, 3)),
+      assertNumLateRows(0),
       AddData(inputData, 25),
       CheckAnswer((10, 3), (25, 1)),
+      assertNumLateRows(0),
       AddData(inputData, 25),
       CheckAnswer((10, 3), (25, 2)),
+      assertNumLateRows(0),
       AddData(inputData, 10),
       CheckAnswer((10, 4), (25, 2)),
       assertNumLateRows(1),
       AddData(inputData, 25),
-      CheckAnswer((10, 4), (25, 3))
+      CheckAnswer((10, 4), (25, 3)),
+      assertNumLateRows(0)
     )
   }