Skip to content

Commit

Permalink
[SPARK-50235][SQL] Clean up ColumnVector resource after processing al…
Browse files Browse the repository at this point in the history
…l rows in ColumnarToRowExec

### What changes were proposed in this pull request?

This patch cleans up ColumnVector resource after processing all rows in ColumnarToRowExec. This patch only focus on codeben implementation of ColumnarToRowExec. For non-codegen, it should be relatively rare to use, and currently no good way has proposed, so leaving it to a follow up.

### Why are the changes needed?

Currently we only assign null to ColumnarBatch object but it doesn't release the resources hold by the vectors in the batch. For OnHeapColumnVector, the Java arrays may be automatically collected by JVM, but for OffHeapColumnVector, the allocated off-heap memory will be leaked.

For custom ColumnVector implementations like Arrow-based, it also possibly causes issues on memory safety if the underlying buffers are reused across batches. Because when ColumnarToRowExec begins to fill values for next batch, the arrays in previous batch are still hold.

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

Existing tests.

### Was this patch authored or co-authored using generative AI tooling?

No

Closes #48767 from viirya/close_if_not_writable.

Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Kent Yao <[email protected]>
  • Loading branch information
viirya authored and yaooqinn committed Nov 6, 2024
1 parent 737a65e commit 800faf0
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,18 @@ public abstract class ColumnVector implements AutoCloseable {
@Override
public abstract void close();

/**
* Cleans up memory for this column vector if it's not writable. The column vector is not usable
* after this.
*
* If this is a writable column vector, it is a no-op.
*/
public void closeIfNotWritable() {
// By default, we just call close() for all column vectors. If a column vector is writable, it
// should override this method and do nothing.
close();
}

/**
* Returns true if this column vector contains any null values.
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,16 @@ public void close() {
}
}

/**
* Called to close all the columns if they are not writable. This is used to clean up memory
* allocated during columnar processing.
*/
public void closeIfNotWritable() {
for (ColumnVector c: columns) {
c.closeIfNotWritable();
}
}

/**
* Returns an iterator over the rows in this batch.
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,11 @@ public void close() {
releaseMemory();
}

@Override
public void closeIfNotWritable() {
// no-op
}

public void reserveAdditional(int additionalCapacity) {
reserve(elementsAppended + additionalCapacity);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -194,9 +194,14 @@ case class ColumnarToRowExec(child: SparkPlan) extends ColumnarToRowTransition w
| $shouldStop
| }
| $idx = $numRows;
| $batch.closeIfNotWritable();
| $batch = null;
| $nextBatchFuncName();
|}
|// clean up resources
|if ($batch != null) {
| $batch.close();
|}
""".stripMargin
}

Expand Down

0 comments on commit 800faf0

Please sign in to comment.