Github user eyalfa commented on a diff in the pull request:
https://github.com/apache/spark/pull/19810#discussion_r218704390
--- Diff:
sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala
---
@@ -52,6 +52,68 @@ object InMemoryRelation {
private[columnar]
case class CachedBatch(numRows: Int, buffers: Array[Array[Byte]], stats:
InternalRow)
+private[columnar] class CachedBatchIterator(
+ rowIterator: Iterator[InternalRow],
+ output: Seq[Attribute],
+ batchSize: Int,
+ useCompression: Boolean,
+ batchStats: LongAccumulator,
+ singleBatchPerPartition: Boolean) extends Iterator[CachedBatch] {
+
+ def next(): CachedBatch = {
+ val columnBuilders = output.map { attribute =>
+ ColumnBuilder(attribute.dataType, batchSize, attribute.name,
useCompression)
+ }.toArray
+
+ var rowCount = 0
+ var totalSize = 0L
+
+ val terminateLoop = (singleBatch: Boolean, rowIter:
Iterator[InternalRow],
+ rowCount: Int, size: Long) => {
+ if (!singleBatch) {
+ rowIter.hasNext && rowCount < batchSize && totalSize <
ColumnBuilder.MAX_BATCH_SIZE_IN_BYTE
+ } else {
+ rowIter.hasNext
--- End diff --
doesn't this run the risk of OOM for large partitions?
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
