diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/TableCompressionCodec.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/TableCompressionCodec.scala index 1785d31fe3c..56a74366f84 100644 --- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/TableCompressionCodec.scala +++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/TableCompressionCodec.scala @@ -193,6 +193,12 @@ abstract class BatchedTableCompressor(maxBatchMemorySize: Long, stream: Cuda.Str val compressedTables = results.toArray results.clear() + + // Ensure we synchronize on the CUDA stream, because `CompressedTable` instances + // could be copied to host during a spill before we are done. + // TODO: A better way to do this would be via CUDA events, synchronizing on the event + // instead of the whole stream + stream.sync() compressedTables }