[SPARK-16995][SQL] TreeNodeException when flat mapping RelationalGroupedDataset created from DataFrame containing a column created with lit/expr

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -727,6 +727,13 @@ object FoldablePropagation extends Rule[LogicalPlan] {
         case j @ Join(_, _, LeftOuter | RightOuter | FullOuter, _) =>
           stop = true
           j
+
+        // Operators that operate on objects should only have expressions from encoders, which
+        // should never have foldable expressions.
+        case o: ObjectConsumer => o
+        case o: ObjectProducer => o
+        case a: AppendColumns => a
+
         case p: LogicalPlan if !stop => p.transformExpressions {
           case a: AttributeReference if foldableMap.contains(a) =>
             foldableMap(a)

sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala

@@ -878,6 +878,21 @@ class DatasetSuite extends QueryTest with SharedSQLContext {
     val ds = spark.createDataset(data)(enc)
     checkDataset(ds, (("a", "b"), "c"), (null, "d"))
   }
+
+  test("SPARK-16995: flat mapping on Dataset containing a column created with lit/expr") {
+    val df = Seq("1").toDF("a")
+
+    import df.sparkSession.implicits._
+
+    assertResult(Seq()) {
+      df.withColumn("b", lit(0)).as[ClassData]
+        .groupByKey(_.a).flatMapGroups { case (x, iter) => List[Int]() }.collect().toSeq
+    }
+    assertResult(Seq()) {
+      df.withColumn("b", expr("0")).as[ClassData]
+        .groupByKey(_.a).flatMapGroups { case (x, iter) => List[Int]() }.collect().toSeq
+    }
+  }
 }
 
 case class Generic[T](id: T, value: Double)