diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala
index edcce103d096330239a34b86aff7b190c2e14c14..a9f79da6358dcf039699a1a28a179e494aca7871 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala
@@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
-import org.apache.spark.sql.execution.{SQLExecution, UnsafeKVExternalSorter}
+import org.apache.spark.sql.execution.{QueryExecution, SQLExecution, UnsafeKVExternalSorter}
import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
import org.apache.spark.util.{SerializableConfiguration, Utils}
import org.apache.spark.util.collection.unsafe.sort.UnsafeExternalSorter
@@ -85,7 +85,7 @@ object FileFormatWriter extends Logging {
*/
def write(
sparkSession: SparkSession,
- plan: LogicalPlan,
+ queryExecution: QueryExecution,
fileFormat: FileFormat,
committer: FileCommitProtocol,
outputSpec: OutputSpec,
@@ -101,8 +101,7 @@ object FileFormatWriter extends Logging {
FileOutputFormat.setOutputPath(job, new Path(outputSpec.outputPath))
val partitionSet = AttributeSet(partitionColumns)
- val dataColumns = plan.output.filterNot(partitionSet.contains)
- val queryExecution = Dataset.ofRows(sparkSession, plan).queryExecution
+ val dataColumns = queryExecution.logical.output.filterNot(partitionSet.contains)
// Note: prepareWrite has side effect. It sets "job".
val outputWriterFactory =
@@ -112,7 +111,7 @@ object FileFormatWriter extends Logging {
uuid = UUID.randomUUID().toString,
serializableHadoopConf = new SerializableConfiguration(job.getConfiguration),
outputWriterFactory = outputWriterFactory,
- allColumns = plan.output,
+ allColumns = queryExecution.logical.output,
partitionColumns = partitionColumns,
nonPartitionColumns = dataColumns,
bucketSpec = bucketSpec,
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/InsertIntoHadoopFsRelationCommand.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/InsertIntoHadoopFsRelationCommand.scala
index 28975e1546e79d2bff2a0a8675c817403db81607..a9bde903b3b5e83f5fcd843910a8ce0955b0c487 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/InsertIntoHadoopFsRelationCommand.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/InsertIntoHadoopFsRelationCommand.scala
@@ -100,7 +100,7 @@ case class InsertIntoHadoopFsRelationCommand(
FileFormatWriter.write(
sparkSession = sparkSession,
- plan = query,
+ queryExecution = Dataset.ofRows(sparkSession, query).queryExecution,
fileFormat = fileFormat,
committer = committer,
outputSpec = FileFormatWriter.OutputSpec(
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/FileStreamSink.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/FileStreamSink.scala
index f1c5f9ab5067db8ad73989961746405a1a7fd7f5..0dbe2a71ed3bc9fc7eb73c7125f408539e3defb7 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/FileStreamSink.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/FileStreamSink.scala
@@ -77,7 +77,7 @@ class FileStreamSink(
FileFormatWriter.write(
sparkSession = sparkSession,
- plan = data.logicalPlan,
+ queryExecution = data.queryExecution,
fileFormat = fileFormat,
committer = committer,
outputSpec = FileFormatWriter.OutputSpec(path, Map.empty),
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala
index fa97d9292e55174ab6a5dd1f6fceb99eca08e862..09613ef9e4348c6ba4455f11228ecac13c4406e9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/FileStreamSinkSuite.scala
@@ -21,13 +21,14 @@ import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.execution.DataSourceScanExec
import org.apache.spark.sql.execution.datasources._
import org.apache.spark.sql.execution.streaming.{MemoryStream, MetadataLogFileIndex}
+import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
import org.apache.spark.util.Utils
class FileStreamSinkSuite extends StreamTest {
import testImplicits._
- test("FileStreamSink - unpartitioned writing and batch reading") {
+ test("unpartitioned writing and batch reading") {
val inputData = MemoryStream[Int]
val df = inputData.toDF()
@@ -59,7 +60,7 @@ class FileStreamSinkSuite extends StreamTest {
}
}
- test("FileStreamSink - partitioned writing and batch reading") {
+ test("partitioned writing and batch reading") {
val inputData = MemoryStream[Int]
val ds = inputData.toDS()
@@ -142,16 +143,83 @@ class FileStreamSinkSuite extends StreamTest {
}
}
- test("FileStreamSink - parquet") {
+ // This tests whether FileStreamSink works with aggregations. Specifically, it tests
+ // whether the the correct streaming QueryExecution (i.e. IncrementalExecution) is used to
+ // to execute the trigger for writing data to file sink. See SPARK-18440 for more details.
+ test("writing with aggregation") {
+
+ // Since FileStreamSink currently only supports append mode, we will test FileStreamSink
+ // with aggregations using event time windows and watermark, which allows
+ // aggregation + append mode.
+ val inputData = MemoryStream[Long]
+ val inputDF = inputData.toDF.toDF("time")
+ val outputDf = inputDF
+ .selectExpr("CAST(time AS timestamp) AS timestamp")
+ .withWatermark("timestamp", "10 seconds")
+ .groupBy(window($"timestamp", "5 seconds"))
+ .count()
+ .select("window.start", "window.end", "count")
+
+ val outputDir = Utils.createTempDir(namePrefix = "stream.output").getCanonicalPath
+ val checkpointDir = Utils.createTempDir(namePrefix = "stream.checkpoint").getCanonicalPath
+
+ var query: StreamingQuery = null
+
+ try {
+ query =
+ outputDf.writeStream
+ .option("checkpointLocation", checkpointDir)
+ .format("parquet")
+ .start(outputDir)
+
+
+ def addTimestamp(timestampInSecs: Int*): Unit = {
+ inputData.addData(timestampInSecs.map(_ * 1L): _*)
+ failAfter(streamingTimeout) {
+ query.processAllAvailable()
+ }
+ }
+
+ def check(expectedResult: ((Long, Long), Long)*): Unit = {
+ val outputDf = spark.read.parquet(outputDir)
+ .selectExpr(
+ "CAST(start as BIGINT) AS start",
+ "CAST(end as BIGINT) AS end",
+ "count")
+ checkDataset(
+ outputDf.as[(Long, Long, Long)],
+ expectedResult.map(x => (x._1._1, x._1._2, x._2)): _*)
+ }
+
+ addTimestamp(100) // watermark = None before this, watermark = 100 - 10 = 90 after this
+ check() // nothing emitted yet
+
+ addTimestamp(104, 123) // watermark = 90 before this, watermark = 123 - 10 = 113 after this
+ check() // nothing emitted yet
+
+ addTimestamp(140) // wm = 113 before this, emit results on 100-105, wm = 130 after this
+ check((100L, 105L) -> 2L)
+
+ addTimestamp(150) // wm = 130s before this, emit results on 120-125, wm = 150 after this
+ check((100L, 105L) -> 2L, (120L, 125L) -> 1L)
+
+ } finally {
+ if (query != null) {
+ query.stop()
+ }
+ }
+ }
+
+ test("parquet") {
testFormat(None) // should not throw error as default format parquet when not specified
testFormat(Some("parquet"))
}
- test("FileStreamSink - text") {
+ test("text") {
testFormat(Some("text"))
}
- test("FileStreamSink - json") {
+ test("json") {
testFormat(Some("json"))
}