Skip to content

perf: the JVM columnar shuffle insert loop boxes, locks and copies every row more than once #7074

Open

Description

@andygrove

What is the problem the feature request solves?

The JVM columnar shuffle's per-row insert loop does avoidable work for every row. On main:

  • Double partition computation (bypass writer). CometBypassMergeSortShuffleWriter.write calls partitioner.getPartition(key) twice per record and keeps only one result: int partition_id = partitioner.getPartition(key); partitionWriters[partitioner.getPartition(key)].insertRow(..., partition_id).
  • A lock per row. CometDiskBlockWriter.insertRow takes synchronized (CometDiskBlockWriter.this) for each row, even though a single task thread inserts. The lock guards state shared with the spill path, so check what really needs it on the insert side.
  • Serializer round trip and double copy. Each UnsafeRow goes through serOutputStream.writeKey/writeValue into an ExposedByteArrayOutputStream, then Platform.copyMemory copies it again into the page (ArrowIPCWriter.insertRecord). Native later skips the length prefix this adds.
  • Boxing. record._1() boxes the partition id for each record in both writers, and RowPartition stores addresses and sizes in ArrayBuffer[Long] and ArrayBuffer[Int]. Scala's ArrayBuffer is not specialized, so each addRow allocates a java.lang.Long and usually a java.lang.Integer, and getRowAddresses / getRowSizes unbox every element into the arrays handed to JNI. SpillSorter does the same per record while reading back a spill.

The per-row INFO log noted in #5905 is already debug.

Describe the potential solution

  • Reuse the first getPartition result in the bypass writer.
  • Replace RowPartition's buffers with manually grown long[] / int[] and a count (or pass the count over JNI so toArray disappears). This touches the JNI boundary, so it needs the most care.
  • Remove the per-row lock on the insert path (guard only the page and spill state that the spill callback shares).
  • Copy the UnsafeRow bytes straight into the page without the serializer stream.

Measure with CometColumnarShuffleWriteBenchmark or a repartition(n) over a wide table before and after each step. Note #7065 changes SpillSorter.writeSortedFileNative; rebase onto it for the sort-based writer.

Additional context

Finding J4 in #5905, items 3 and 4 of #5198. Related: #7058 and #7061 (per-partition cost in the same writers).

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions