Repository: spark Updated Branches: refs/heads/master 5a3533e77 -> 00c310133
[SPARK-15593][SQL] Add DataFrameWriter.foreach to allow the user consuming data in ContinuousQuery ## What changes were proposed in this pull request? * Add DataFrameWriter.foreach to allow the user consuming data in ContinuousQuery * ForeachWriter is the interface for the user to consume partitions of data * Add a type parameter T to DataFrameWriter Usage ```Scala val ds = spark.read....stream().as[String] ds.....write .queryName(...) .option("checkpointLocation", ...) .foreach(new ForeachWriter[Int] { def open(partitionId: Long, version: Long): Boolean = { // prepare some resources for a partition // check `version` if possible and return `false` if this is a duplicated data to skip the data processing. } override def process(value: Int): Unit = { // process data } def close(errorOrNull: Throwable): Unit = { // release resources for a partition // check `errorOrNull` and handle the error if necessary. } }) ``` ## How was this patch tested? New unit tests. Author: Shixiong Zhu <shixi...@databricks.com> Closes #13342 from zsxwing/foreach. Project: http://git-wip-us.apache.org/repos/asf/spark/repo Commit: http://git-wip-us.apache.org/repos/asf/spark/commit/00c31013 Tree: http://git-wip-us.apache.org/repos/asf/spark/tree/00c31013 Diff: http://git-wip-us.apache.org/repos/asf/spark/diff/00c31013 Branch: refs/heads/master Commit: 00c310133df4f3893dd90d801168c2ab9841b102 Parents: 5a3533e Author: Shixiong Zhu <shixi...@databricks.com> Authored: Fri Jun 10 00:11:46 2016 -0700 Committer: Tathagata Das <tathagata.das1...@gmail.com> Committed: Fri Jun 10 00:11:46 2016 -0700 ---------------------------------------------------------------------- .../org/apache/spark/sql/DataFrameWriter.scala | 150 ++++++++++++++----- .../scala/org/apache/spark/sql/Dataset.scala | 2 +- .../org/apache/spark/sql/ForeachWriter.scala | 105 +++++++++++++ .../sql/execution/streaming/ForeachSink.scala | 53 +++++++ .../execution/streaming/ForeachSinkSuite.scala | 141 +++++++++++++++++ .../spark/sql/sources/BucketedReadSuite.scala | 4 +- 6 files changed, 413 insertions(+), 42 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/spark/blob/00c31013/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala ---------------------------------------------------------------------- diff --git a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala index 1dd8818..32e2fdc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala @@ -29,7 +29,7 @@ import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation import org.apache.spark.sql.catalyst.plans.logical.{InsertIntoTable, Project} import org.apache.spark.sql.execution.datasources.{BucketSpec, CreateTableUsingAsSelect, DataSource, HadoopFsRelation} import org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils -import org.apache.spark.sql.execution.streaming.{MemoryPlan, MemorySink, StreamExecution} +import org.apache.spark.sql.execution.streaming._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.streaming.{ContinuousQuery, OutputMode, ProcessingTime, Trigger} import org.apache.spark.util.Utils @@ -40,7 +40,9 @@ import org.apache.spark.util.Utils * * @since 1.4.0 */ -final class DataFrameWriter private[sql](df: DataFrame) { +final class DataFrameWriter[T] private[sql](ds: Dataset[T]) { + + private val df = ds.toDF() /** * Specifies the behavior when data or table already exists. Options include: @@ -51,7 +53,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * * @since 1.4.0 */ - def mode(saveMode: SaveMode): DataFrameWriter = { + def mode(saveMode: SaveMode): DataFrameWriter[T] = { // mode() is used for non-continuous queries // outputMode() is used for continuous queries assertNotStreaming("mode() can only be called on non-continuous queries") @@ -68,7 +70,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * * @since 1.4.0 */ - def mode(saveMode: String): DataFrameWriter = { + def mode(saveMode: String): DataFrameWriter[T] = { // mode() is used for non-continuous queries // outputMode() is used for continuous queries assertNotStreaming("mode() can only be called on non-continuous queries") @@ -93,7 +95,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * @since 2.0.0 */ @Experimental - def outputMode(outputMode: OutputMode): DataFrameWriter = { + def outputMode(outputMode: OutputMode): DataFrameWriter[T] = { assertStreaming("outputMode() can only be called on continuous queries") this.outputMode = outputMode this @@ -109,7 +111,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * @since 2.0.0 */ @Experimental - def outputMode(outputMode: String): DataFrameWriter = { + def outputMode(outputMode: String): DataFrameWriter[T] = { assertStreaming("outputMode() can only be called on continuous queries") this.outputMode = outputMode.toLowerCase match { case "append" => @@ -147,7 +149,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * @since 2.0.0 */ @Experimental - def trigger(trigger: Trigger): DataFrameWriter = { + def trigger(trigger: Trigger): DataFrameWriter[T] = { assertStreaming("trigger() can only be called on continuous queries") this.trigger = trigger this @@ -158,7 +160,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * * @since 1.4.0 */ - def format(source: String): DataFrameWriter = { + def format(source: String): DataFrameWriter[T] = { this.source = source this } @@ -168,7 +170,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * * @since 1.4.0 */ - def option(key: String, value: String): DataFrameWriter = { + def option(key: String, value: String): DataFrameWriter[T] = { this.extraOptions += (key -> value) this } @@ -178,28 +180,28 @@ final class DataFrameWriter private[sql](df: DataFrame) { * * @since 2.0.0 */ - def option(key: String, value: Boolean): DataFrameWriter = option(key, value.toString) + def option(key: String, value: Boolean): DataFrameWriter[T] = option(key, value.toString) /** * Adds an output option for the underlying data source. * * @since 2.0.0 */ - def option(key: String, value: Long): DataFrameWriter = option(key, value.toString) + def option(key: String, value: Long): DataFrameWriter[T] = option(key, value.toString) /** * Adds an output option for the underlying data source. * * @since 2.0.0 */ - def option(key: String, value: Double): DataFrameWriter = option(key, value.toString) + def option(key: String, value: Double): DataFrameWriter[T] = option(key, value.toString) /** * (Scala-specific) Adds output options for the underlying data source. * * @since 1.4.0 */ - def options(options: scala.collection.Map[String, String]): DataFrameWriter = { + def options(options: scala.collection.Map[String, String]): DataFrameWriter[T] = { this.extraOptions ++= options this } @@ -209,7 +211,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * * @since 1.4.0 */ - def options(options: java.util.Map[String, String]): DataFrameWriter = { + def options(options: java.util.Map[String, String]): DataFrameWriter[T] = { this.options(options.asScala) this } @@ -232,7 +234,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * @since 1.4.0 */ @scala.annotation.varargs - def partitionBy(colNames: String*): DataFrameWriter = { + def partitionBy(colNames: String*): DataFrameWriter[T] = { this.partitioningColumns = Option(colNames) this } @@ -246,7 +248,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * @since 2.0 */ @scala.annotation.varargs - def bucketBy(numBuckets: Int, colName: String, colNames: String*): DataFrameWriter = { + def bucketBy(numBuckets: Int, colName: String, colNames: String*): DataFrameWriter[T] = { this.numBuckets = Option(numBuckets) this.bucketColumnNames = Option(colName +: colNames) this @@ -260,7 +262,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * @since 2.0 */ @scala.annotation.varargs - def sortBy(colName: String, colNames: String*): DataFrameWriter = { + def sortBy(colName: String, colNames: String*): DataFrameWriter[T] = { this.sortColumnNames = Option(colName +: colNames) this } @@ -301,7 +303,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { * @since 2.0.0 */ @Experimental - def queryName(queryName: String): DataFrameWriter = { + def queryName(queryName: String): DataFrameWriter[T] = { assertStreaming("queryName() can only be called on continuous queries") this.extraOptions += ("queryName" -> queryName) this @@ -337,16 +339,7 @@ final class DataFrameWriter private[sql](df: DataFrame) { val queryName = extraOptions.getOrElse( "queryName", throw new AnalysisException("queryName must be specified for memory sink")) - val checkpointLocation = extraOptions.get("checkpointLocation").map { userSpecified => - new Path(userSpecified).toUri.toString - }.orElse { - val checkpointConfig: Option[String] = - df.sparkSession.conf.get(SQLConf.CHECKPOINT_LOCATION) - - checkpointConfig.map { location => - new Path(location, queryName).toUri.toString - } - }.getOrElse { + val checkpointLocation = getCheckpointLocation(queryName, failIfNotSet = false).getOrElse { Utils.createTempDir(namePrefix = "memory.stream").getCanonicalPath } @@ -378,21 +371,10 @@ final class DataFrameWriter private[sql](df: DataFrame) { className = source, options = extraOptions.toMap, partitionColumns = normalizedParCols.getOrElse(Nil)) - val queryName = extraOptions.getOrElse("queryName", StreamExecution.nextName) - val checkpointLocation = extraOptions.get("checkpointLocation") - .orElse { - df.sparkSession.sessionState.conf.checkpointLocation.map { l => - new Path(l, queryName).toUri.toString - } - }.getOrElse { - throw new AnalysisException("checkpointLocation must be specified either " + - "through option() or SQLConf") - } - df.sparkSession.sessionState.continuousQueryManager.startQuery( queryName, - checkpointLocation, + getCheckpointLocation(queryName, failIfNotSet = true).get, df, dataSource.createSink(outputMode), outputMode, @@ -401,6 +383,94 @@ final class DataFrameWriter private[sql](df: DataFrame) { } /** + * :: Experimental :: + * Starts the execution of the streaming query, which will continually send results to the given + * [[ForeachWriter]] as as new data arrives. The [[ForeachWriter]] can be used to send the data + * generated by the [[DataFrame]]/[[Dataset]] to an external system. The returned The returned + * [[ContinuousQuery]] object can be used to interact with the stream. + * + * Scala example: + * {{{ + * datasetOfString.write.foreach(new ForeachWriter[String] { + * + * def open(partitionId: Long, version: Long): Boolean = { + * // open connection + * } + * + * def process(record: String) = { + * // write string to connection + * } + * + * def close(errorOrNull: Throwable): Unit = { + * // close the connection + * } + * }) + * }}} + * + * Java example: + * {{{ + * datasetOfString.write().foreach(new ForeachWriter<String>() { + * + * @Override + * public boolean open(long partitionId, long version) { + * // open connection + * } + * + * @Override + * public void process(String value) { + * // write string to connection + * } + * + * @Override + * public void close(Throwable errorOrNull) { + * // close the connection + * } + * }); + * }}} + * + * @since 2.0.0 + */ + @Experimental + def foreach(writer: ForeachWriter[T]): ContinuousQuery = { + assertNotBucketed("foreach") + assertStreaming( + "foreach() can only be called on streaming Datasets/DataFrames.") + + val queryName = extraOptions.getOrElse("queryName", StreamExecution.nextName) + val sink = new ForeachSink[T](ds.sparkSession.sparkContext.clean(writer))(ds.exprEnc) + df.sparkSession.sessionState.continuousQueryManager.startQuery( + queryName, + getCheckpointLocation(queryName, failIfNotSet = false).getOrElse { + Utils.createTempDir(namePrefix = "foreach.stream").getCanonicalPath + }, + df, + sink, + outputMode, + trigger) + } + + /** + * Returns the checkpointLocation for a query. If `failIfNotSet` is `true` but the checkpoint + * location is not set, [[AnalysisException]] will be thrown. If `failIfNotSet` is `false`, `None` + * will be returned if the checkpoint location is not set. + */ + private def getCheckpointLocation(queryName: String, failIfNotSet: Boolean): Option[String] = { + val checkpointLocation = extraOptions.get("checkpointLocation").map { userSpecified => + new Path(userSpecified).toUri.toString + }.orElse { + df.sparkSession.conf.get(SQLConf.CHECKPOINT_LOCATION).map { location => + new Path(location, queryName).toUri.toString + } + } + if (failIfNotSet && checkpointLocation.isEmpty) { + throw new AnalysisException("checkpointLocation must be specified either " + + """through option("checkpointLocation", ...) or """ + + s"""SparkSession.conf.set("${SQLConf.CHECKPOINT_LOCATION.key}", ...)""") + } + checkpointLocation + } + + /** * Inserts the content of the [[DataFrame]] to the specified table. It requires that * the schema of the [[DataFrame]] is the same as the schema of the table. * http://git-wip-us.apache.org/repos/asf/spark/blob/00c31013/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala ---------------------------------------------------------------------- diff --git a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala index 162524a..16bbf30 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala @@ -2400,7 +2400,7 @@ class Dataset[T] private[sql]( * @since 1.6.0 */ @Experimental - def write: DataFrameWriter = new DataFrameWriter(toDF()) + def write: DataFrameWriter[T] = new DataFrameWriter[T](this) /** * Returns the content of the Dataset as a Dataset of JSON strings. http://git-wip-us.apache.org/repos/asf/spark/blob/00c31013/sql/core/src/main/scala/org/apache/spark/sql/ForeachWriter.scala ---------------------------------------------------------------------- diff --git a/sql/core/src/main/scala/org/apache/spark/sql/ForeachWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/ForeachWriter.scala new file mode 100644 index 0000000..09f0742 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/ForeachWriter.scala @@ -0,0 +1,105 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql + +import org.apache.spark.annotation.Experimental +import org.apache.spark.sql.streaming.ContinuousQuery + +/** + * :: Experimental :: + * A class to consume data generated by a [[ContinuousQuery]]. Typically this is used to send the + * generated data to external systems. Each partition will use a new deserialized instance, so you + * usually should do all the initialization (e.g. opening a connection or initiating a transaction) + * in the `open` method. + * + * Scala example: + * {{{ + * datasetOfString.write.foreach(new ForeachWriter[String] { + * + * def open(partitionId: Long, version: Long): Boolean = { + * // open connection + * } + * + * def process(record: String) = { + * // write string to connection + * } + * + * def close(errorOrNull: Throwable): Unit = { + * // close the connection + * } + * }) + * }}} + * + * Java example: + * {{{ + * datasetOfString.write().foreach(new ForeachWriter<String>() { + * + * @Override + * public boolean open(long partitionId, long version) { + * // open connection + * } + * + * @Override + * public void process(String value) { + * // write string to connection + * } + * + * @Override + * public void close(Throwable errorOrNull) { + * // close the connection + * } + * }); + * }}} + * @since 2.0.0 + */ +@Experimental +abstract class ForeachWriter[T] extends Serializable { + + /** + * Called when starting to process one partition of new data in the executor. The `version` is + * for data deduplication when there are failures. When recovering from a failure, some data may + * be generated multiple times but they will always have the same version. + * + * If this method finds using the `partitionId` and `version` that this partition has already been + * processed, it can return `false` to skip the further data processing. However, `close` still + * will be called for cleaning up resources. + * + * @param partitionId the partition id. + * @param version a unique id for data deduplication. + * @return `true` if the corresponding partition and version id should be processed. `false` + * indicates the partition should be skipped. + */ + def open(partitionId: Long, version: Long): Boolean + + /** + * Called to process the data in the executor side. This method will be called only when `open` + * returns `true`. + */ + def process(value: T): Unit + + /** + * Called when stopping to process one partition of new data in the executor side. This is + * guaranteed to be called either `open` returns `true` or `false`. However, + * `close` won't be called in the following cases: + * - JVM crashes without throwing a `Throwable` + * - `open` throws a `Throwable`. + * + * @param errorOrNull the error thrown during processing data or null if there was no error. + */ + def close(errorOrNull: Throwable): Unit +} http://git-wip-us.apache.org/repos/asf/spark/blob/00c31013/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ForeachSink.scala ---------------------------------------------------------------------- diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ForeachSink.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ForeachSink.scala new file mode 100644 index 0000000..14b9b1c --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/ForeachSink.scala @@ -0,0 +1,53 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.streaming + +import org.apache.spark.TaskContext +import org.apache.spark.sql.{DataFrame, Encoder, ForeachWriter} + +/** + * A [[Sink]] that forwards all data into [[ForeachWriter]] according to the contract defined by + * [[ForeachWriter]]. + * + * @param writer The [[ForeachWriter]] to process all data. + * @tparam T The expected type of the sink. + */ +class ForeachSink[T : Encoder](writer: ForeachWriter[T]) extends Sink with Serializable { + + override def addBatch(batchId: Long, data: DataFrame): Unit = { + data.as[T].foreachPartition { iter => + if (writer.open(TaskContext.getPartitionId(), batchId)) { + var isFailed = false + try { + while (iter.hasNext) { + writer.process(iter.next()) + } + } catch { + case e: Throwable => + isFailed = true + writer.close(e) + } + if (!isFailed) { + writer.close(null) + } + } else { + writer.close(null) + } + } + } +} http://git-wip-us.apache.org/repos/asf/spark/blob/00c31013/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/ForeachSinkSuite.scala ---------------------------------------------------------------------- diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/ForeachSinkSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/ForeachSinkSuite.scala new file mode 100644 index 0000000..e1fb3b9 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/ForeachSinkSuite.scala @@ -0,0 +1,141 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.streaming + +import java.util.concurrent.ConcurrentLinkedQueue + +import scala.collection.mutable + +import org.scalatest.BeforeAndAfter + +import org.apache.spark.sql.ForeachWriter +import org.apache.spark.sql.streaming.StreamTest +import org.apache.spark.sql.test.SharedSQLContext + +class ForeachSinkSuite extends StreamTest with SharedSQLContext with BeforeAndAfter { + + import testImplicits._ + + after { + sqlContext.streams.active.foreach(_.stop()) + } + + test("foreach") { + withTempDir { checkpointDir => + val input = MemoryStream[Int] + val query = input.toDS().repartition(2).write + .option("checkpointLocation", checkpointDir.getCanonicalPath) + .foreach(new TestForeachWriter()) + input.addData(1, 2, 3, 4) + query.processAllAvailable() + + val expectedEventsForPartition0 = Seq( + ForeachSinkSuite.Open(partition = 0, version = 0), + ForeachSinkSuite.Process(value = 1), + ForeachSinkSuite.Process(value = 3), + ForeachSinkSuite.Close(None) + ) + val expectedEventsForPartition1 = Seq( + ForeachSinkSuite.Open(partition = 1, version = 0), + ForeachSinkSuite.Process(value = 2), + ForeachSinkSuite.Process(value = 4), + ForeachSinkSuite.Close(None) + ) + + val allEvents = ForeachSinkSuite.allEvents() + assert(allEvents.size === 2) + assert { + allEvents === Seq(expectedEventsForPartition0, expectedEventsForPartition1) || + allEvents === Seq(expectedEventsForPartition1, expectedEventsForPartition0) + } + query.stop() + } + } + + test("foreach with error") { + withTempDir { checkpointDir => + val input = MemoryStream[Int] + val query = input.toDS().repartition(1).write + .option("checkpointLocation", checkpointDir.getCanonicalPath) + .foreach(new TestForeachWriter() { + override def process(value: Int): Unit = { + super.process(value) + throw new RuntimeException("error") + } + }) + input.addData(1, 2, 3, 4) + query.processAllAvailable() + + val allEvents = ForeachSinkSuite.allEvents() + assert(allEvents.size === 1) + assert(allEvents(0)(0) === ForeachSinkSuite.Open(partition = 0, version = 0)) + assert(allEvents(0)(1) === ForeachSinkSuite.Process(value = 1)) + val errorEvent = allEvents(0)(2).asInstanceOf[ForeachSinkSuite.Close] + assert(errorEvent.error.get.isInstanceOf[RuntimeException]) + assert(errorEvent.error.get.getMessage === "error") + query.stop() + } + } +} + +/** A global object to collect events in the executor */ +object ForeachSinkSuite { + + trait Event + + case class Open(partition: Long, version: Long) extends Event + + case class Process[T](value: T) extends Event + + case class Close(error: Option[Throwable]) extends Event + + private val _allEvents = new ConcurrentLinkedQueue[Seq[Event]]() + + def addEvents(events: Seq[Event]): Unit = { + _allEvents.add(events) + } + + def allEvents(): Seq[Seq[Event]] = { + _allEvents.toArray(new Array[Seq[Event]](_allEvents.size())) + } + + def clear(): Unit = { + _allEvents.clear() + } +} + +/** A [[ForeachWriter]] that writes collected events to ForeachSinkSuite */ +class TestForeachWriter extends ForeachWriter[Int] { + ForeachSinkSuite.clear() + + private val events = mutable.ArrayBuffer[ForeachSinkSuite.Event]() + + override def open(partitionId: Long, version: Long): Boolean = { + events += ForeachSinkSuite.Open(partition = partitionId, version = version) + true + } + + override def process(value: Int): Unit = { + events += ForeachSinkSuite.Process(value) + } + + override def close(errorOrNull: Throwable): Unit = { + events += ForeachSinkSuite.Close(error = Option(errorOrNull)) + ForeachSinkSuite.addEvents(events) + } +} http://git-wip-us.apache.org/repos/asf/spark/blob/00c31013/sql/hive/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala ---------------------------------------------------------------------- diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala index bab0092..fc01ff3 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala @@ -238,7 +238,9 @@ class BucketedReadSuite extends QueryTest with SQLTestUtils with TestHiveSinglet shuffleLeft: Boolean, shuffleRight: Boolean): Unit = { withTable("bucketed_table1", "bucketed_table2") { - def withBucket(writer: DataFrameWriter, bucketSpec: Option[BucketSpec]): DataFrameWriter = { + def withBucket( + writer: DataFrameWriter[Row], + bucketSpec: Option[BucketSpec]): DataFrameWriter[Row] = { bucketSpec.map { spec => writer.bucketBy( spec.numBuckets, --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org For additional commands, e-mail: commits-h...@spark.apache.org