Skip to content

[BLAZE-747] Enhance the ArrowFFIExporter.exportNextBatch method to execute conditionally #748

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
richox merged 2 commits into from
Jan 9, 2025
Merged

[BLAZE-747] Enhance the ArrowFFIExporter.exportNextBatch method to execute conditionally #748

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9837a35
[BLAZE-747] Enhance the ArrowFFIExporter.exportNextBatch method to ex…
merrily01 Jan 7, 2025
eb59958
nit: remove useless import
merrily01 Jan 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 27 additions & 6 deletions ...ension/src/main/scala/org/apache/spark/sql/execution/blaze/arrowio/ArrowFFIExporter.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,14 +20,17 @@ import org.apache.arrow.c.ArrowSchema
import org.apache.arrow.c.Data
import org.apache.arrow.vector.VectorSchemaRoot
import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider
import org.apache.hadoop.security.UserGroupInformation
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.execution.blaze.arrowio.util.ArrowUtils
import org.apache.spark.sql.execution.blaze.arrowio.util.ArrowWriter
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.blaze.BlazeConf
import org.apache.spark.sql.blaze.{BlazeConf, NativeHelper}
import org.apache.spark.sql.blaze.util.Using
import org.apache.spark.TaskContext

import java.security.PrivilegedExceptionAction

class ArrowFFIExporter(rowIter: Iterator[InternalRow], schema: StructType) {
private val maxBatchNumRows = BlazeConf.BATCH_SIZE.intConf()
private val maxBatchMemorySize = 1 << 24 // 16MB
Expand All @@ -53,6 +56,9 @@ class ArrowFFIExporter(rowIter: Iterator[InternalRow], schema: StructType) {

def exportNextBatch(exportArrowArrayPtr: Long): Boolean = {
val tc = TaskContext.get()
val currentUserInfo = UserGroupInformation.getCurrentUser
val nativeCurrentUser = NativeHelper.currentUser
val isNativeCurrentUser = currentUserInfo.equals(nativeCurrentUser)

if (tc != null && (tc.isCompleted() || tc.isInterrupted())) return false
if (!rowIter.hasNext) return false
Expand All @@ -64,11 +70,26 @@ class ArrowFFIExporter(rowIter: Iterator[InternalRow], schema: StructType) {
val arrowWriter = ArrowWriter.create(root)
var rowCount = 0

while (rowIter.hasNext
&& rowCount < maxBatchNumRows
&& batchAllocator.getAllocatedMemory < maxBatchMemorySize) {
arrowWriter.write(rowIter.next())
rowCount += 1
def processRows(): Unit = {
while (rowIter.hasNext
&& rowCount < maxBatchNumRows
&& batchAllocator.getAllocatedMemory < maxBatchMemorySize) {
arrowWriter.write(rowIter.next())
rowCount += 1
}
}
// if current user is native user, process rows directly
if (isNativeCurrentUser) {
processRows()
} else {
// otherwise, process rows as native user
nativeCurrentUser.doAs(
new PrivilegedExceptionAction[Unit] {
override def run(): Unit = {
processRows()
}
}
)
}
arrowWriter.finish()

Expand Down