bug fix

Author: uncleGen

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala

@@ -848,7 +848,7 @@ case class Sample(
 case class ReservoirSample(
     keys: Seq[Attribute],
     child: LogicalPlan,
-    k: Int,
+    reservoirSize: Int,
     streaming: Boolean = false)
   extends UnaryNode {
   override def maxRows: Option[Long] = child.maxRows

sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala

@@ -2029,9 +2029,9 @@ class Dataset[T] private[sql](
    */
   @Experimental
   @InterfaceStability.Evolving
-  def reservoir(k: Int): Dataset[T] = withTypedPlan {
+  def reservoir(reservoirSize: Int): Dataset[T] = withTypedPlan {
     val allColumns = queryExecution.analyzed.output
-    ReservoirSample(allColumns, logicalPlan, k, isStreaming)
+    ReservoirSample(allColumns, logicalPlan, reservoirSize, isStreaming)
   }
 
   /**

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala

@@ -259,8 +259,8 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
    */
   object ReservoirSampleStrategy extends Strategy {
     override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
-      case ReservoirSample(keys, child, k, true) =>
-        StreamingReservoirSampleExec(keys, PlanLater(child), k) :: Nil
+      case ReservoirSample(keys, child, reservoirSize, true) =>
+        StreamingReservoirSampleExec(keys, PlanLater(child), reservoirSize) :: Nil
 
       case _ => Nil
     }
@@ -421,8 +421,8 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
         execution.window.WindowExec(windowExprs, partitionSpec, orderSpec, planLater(child)) :: Nil
       case logical.Sample(lb, ub, withReplacement, seed, child) =>
         execution.SampleExec(lb, ub, withReplacement, seed, planLater(child)) :: Nil
-      case logical.ReservoirSample(keys, child, k, false) =>
-        execution.ReservoirSampleExec(k, PlanLater(child)) :: Nil
+      case logical.ReservoirSample(keys, child, reservoirSize, false) =>
+        execution.ReservoirSampleExec(reservoirSize, PlanLater(child)) :: Nil
       case logical.LocalRelation(output, data) =>
         LocalTableScanExec(output, data) :: Nil
       case logical.LocalLimit(IntegerLiteral(limit), child) =>

sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala

@@ -658,19 +658,19 @@ object SubqueryExec {
     ThreadUtils.newDaemonCachedThreadPool("subquery", 16))
 }
 
-case class ReservoirSampleExec(k: Int, child: SparkPlan) extends UnaryExecNode {
+case class ReservoirSampleExec(reservoirSize: Int, child: SparkPlan) extends UnaryExecNode {
   override def output: Seq[Attribute] = child.output
 
   override def outputPartitioning: Partitioning = child.outputPartitioning
 
   protected override def doExecute(): RDD[InternalRow] = {
     child.execute()
       .mapPartitions(it => {
-        val (sample, count) = SamplingUtils.reservoirSampleAndCount(it, k)
+        val (sample, count) = SamplingUtils.reservoirSampleAndCount(it, reservoirSize)
         sample.map((_, count)).toIterator
       })
       .repartition(1)
       .mapPartitions(it => {
-        SamplingUtils.reservoirSampleWithWeight(it, k).iterator})
+        SamplingUtils.reservoirSampleWithWeight(it, reservoirSize).iterator})
   }
 }

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

@@ -103,13 +103,18 @@ class IncrementalExecution(
           child,
           Some(stateId),
           Some(offsetSeqMetadata.batchWatermarkMs))
-      
-      case StreamingReservoirSampleExec(k, keys, child, None, None, None) =>
+
+      case StreamingReservoirSampleExec(keys, child, reservoirSize, None, None, None) =>
         val stateId =
           OperatorStateId(checkpointLocation, operatorId.getAndIncrement(), currentBatchId)
         StreamingReservoirSampleExec(
-          k, keys, child, Some(stateId), Some(currentEventTimeWatermark), Some(outputMode))
-      
+          keys,
+          child,
+          reservoirSize,
+          Some(stateId),
+          Some(offsetSeqMetadata.batchWatermarkMs),
+          Some(outputMode))
+
       case m: FlatMapGroupsWithStateExec =>
         val stateId =
           OperatorStateId(checkpointLocation, operatorId.getAndIncrement(), currentBatchId)

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

@@ -330,12 +330,12 @@ object StreamingDeduplicateExec {
 /**
  * Physical operator for executing streaming Sampling.
  *
- * @param k random sample k elements.
+ * @param reservoirSize number of random sample elements.
  */
 case class StreamingReservoirSampleExec(
     keyExpressions: Seq[Attribute],
     child: SparkPlan,
-    k: Int,
+    reservoirSize: Int,
     stateId: Option[OperatorStateId] = None,
     eventTimeWatermark: Option[Long] = None,
     outputMode: Option[OutputMode] = None)
@@ -378,13 +378,13 @@ case class StreamingReservoirSampleExec(
 
       baseIterator.foreach { r =>
         count += 1
-        if (numSamples < k) {
+        if (numSamples < reservoirSize) {
           numSamples += 1
           store.put(enc.toRow(numSamples.toString).asInstanceOf[UnsafeRow],
             r.asInstanceOf[UnsafeRow])
         } else {
           val randomIdx = (rand.nextDouble() * (numRecordsInPart + count)).toLong
-          if (randomIdx <= k) {
+          if (randomIdx <= reservoirSize) {
             val replacementIdx = enc.toRow(randomIdx.toString).asInstanceOf[UnsafeRow]
             store.put(replacementIdx, r.asInstanceOf[UnsafeRow])
           }
@@ -421,7 +421,7 @@ case class StreamingReservoirSampleExec(
       }
     }.repartition(1).mapPartitions(it => {
       SamplingUtils.reservoirSampleWithWeight(
-        it.map(item => (item, item.getLong(keyExpressions.size))), k)
+        it.map(item => (item, item.getLong(keyExpressions.size))), reservoirSize)
         .map(row =>
           UnsafeProjection.create(fieldTypes)
             .apply(InternalRow.fromSeq(row.toSeq(fieldTypes)))