BryanCutler commented on a change in pull request #24795: [SPARK-27945][SQL]
Minimal changes to support columnar processing
URL: https://github.com/apache/spark/pull/24795#discussion_r294933928
##########
File path:
sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
##########
@@ -251,6 +286,388 @@ object MyExtensions {
(_: Seq[Expression]) => Literal(5, IntegerType))
}
+case class CloseableColumnBatchIterator(itr: Iterator[ColumnarBatch],
+ f: ColumnarBatch => ColumnarBatch) extends Iterator[ColumnarBatch] {
+ var cb: ColumnarBatch = null
+
+ private def closeCurrentBatch(): Unit = {
+ if (cb != null) {
+ cb.close
+ cb = null
+ }
+ }
+
+ TaskContext.get().addTaskCompletionListener[Unit]((tc: TaskContext) => {
+ closeCurrentBatch()
+ })
+
+ override def hasNext: Boolean = {
+ closeCurrentBatch()
+ itr.hasNext
+ }
+
+ override def next(): ColumnarBatch = {
+ closeCurrentBatch()
+ cb = f(itr.next())
+ cb
+ }
+}
+
+object NoCloseColumnVector extends Logging {
+ def wrapIfNeeded(cv: ColumnVector): NoCloseColumnVector = cv match {
+ case ref: NoCloseColumnVector =>
+ ref
+ case vec => NoCloseColumnVector(vec)
+ }
+}
+
+/**
+ * Provide a ColumnVector so ColumnarExpression can close temporary values
without
+ * having to guess what type it really is.
+ */
+case class NoCloseColumnVector(wrapped: ColumnVector) extends
ColumnVector(wrapped.dataType) {
+ private var refCount = 1
+
+ /**
+ * Don't actually close the ColumnVector this wraps. The producer of the
vector will take
+ * care of that.
+ */
+ override def close(): Unit = {
+ // Empty
+ }
+
+
Review comment:
nit: can you use 1 newline here and below instead of 2?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
