fix issue without whole stage codegen

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala

@@ -377,6 +377,9 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro
 
         public SpecificUnsafeProjection(Object[] references) {
           this.references = references;
+        }
+
+        public void initializeStatesForPartition(int partitionIndex) {
           ${ctx.initMutableStates()}
         }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/package.scala

@@ -67,11 +67,11 @@ package object expressions  {
   abstract class Projection extends (InternalRow => InternalRow) {
 
     /**
-     * Initialize internal state given the current partition index.
-     * This is used by non-deterministic expressions to set the initial state
+     * Initialize internal states given the current partition index.
+     * This is used by non-deterministic expressions to set initial states.
      * The default implementation does nothing.
      */
-    def initializeStateForPartition(partitionIndex: Int): Unit = {}
+    def initializeStatesForPartition(partitionIndex: Int): Unit = {}
   }
 
   /**

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

@@ -70,9 +70,10 @@ case class ProjectExec(projectList: Seq[NamedExpression], child: SparkPlan)
   }
 
   protected override def doExecute(): RDD[InternalRow] = {
-    child.execute().mapPartitionsInternal { iter =>
+    child.execute().mapPartitionsWithIndexInternal { (index, iter) =>
       val project = UnsafeProjection.create(projectList, child.output,
         subexpressionEliminationEnabled)
+      project.initializeStatesForPartition(index)
       iter.map(project)
     }
   }