feat(amber): columnar hardening (wire versioning, compression, sink)

amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/InputManager.scala

@@ -153,4 +153,9 @@ class InputManager(
     inputBatch = batch
     currentInputIdx = -1
   }
+
+  // Marks the current batch fully consumed (hasUnfinishedInput becomes false).
+  def skipToEnd(): Unit = {
+    currentInputIdx = if (inputBatch == null) -1 else inputBatch.length - 1
+  }
 }

amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManager.scala

@@ -33,11 +33,16 @@ import org.apache.texera.amber.engine.architecture.messaginglayer.OutputManager.
 import org.apache.texera.amber.engine.architecture.sendsemantics.partitioners._
 import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings._
 import org.apache.texera.amber.engine.architecture.worker.managers.{
+  ArrowBatchWriteItem,
   OutputPortStorageWriterThread,
-  PortStorageWriterTerminateSignal
+  PortStorageWriterTerminateSignal,
+  RowWriteItem
 }
 import org.apache.texera.amber.engine.common.AmberLogging
-import org.apache.texera.amber.util.VirtualIdentityUtils
+import org.apache.texera.amber.engine.common.ambermessage.ColumnarFrame
+import org.apache.texera.amber.util.{ArrowUtils, VirtualIdentityUtils}
+import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector.VectorSchemaRoot
 
 import java.net.URI
 import scala.collection.mutable
@@ -204,6 +209,88 @@ class OutputManager(
     saveStateToStorageIfNeeded(state, loopCounter, loopStartId)
   }
 
+  // Allocator for the per-receiver sub-batches built during columnar shuffle.
+  // Created lazily on first shuffle emit, freed by closeColumnarResources().
+  private var columnarEmitAllocator: RootAllocator = _
+  private def emitAllocator(): RootAllocator = {
+    if (columnarEmitAllocator == null) columnarEmitAllocator = new RootAllocator()
+    columnarEmitAllocator
+  }
+  def closeColumnarResources(): Unit = {
+    if (columnarEmitAllocator != null) {
+      columnarEmitAllocator.close()
+      columnarEmitAllocator = null
+    }
+  }
+
+  // Emit an Arrow batch downstream as ColumnarFrames, honoring each link's
+  // partitioning. Single-receiver links and broadcast ship the whole batch;
+  // shuffles slice it into one sub-batch per receiver.
+  def emitColumnarBatch(root: VectorSchemaRoot): Unit = {
+    val rowCount = root.getRowCount
+    val texeraSchema = ArrowUtils.toTexeraSchema(root.getSchema)
+    lazy val wholeBatchBytes = ArrowUtils.serializeRoot(root)
+
+    partitioners.foreach {
+      case (_, partitioner) =>
+        val receivers = partitioner.allReceivers
+        partitioner match {
+          case _ if receivers.size == 1 =>
+            outputGateway.sendTo(
+              receivers.head,
+              ColumnarFrame(wholeBatchBytes, rowCount, texeraSchema)
+            )
+          case _: BroadcastPartitioner =>
+            receivers.foreach(r =>
+              outputGateway.sendTo(r, ColumnarFrame(wholeBatchBytes, rowCount, texeraSchema))
+            )
+          case _ =>
+            // Shuffle/round-robin: assign each row to a receiver via the
+            // partitioner, then ship one Arrow sub-batch per receiver. The
+            // row decode here is only to compute the partition key; downstream
+            // still receives columnar data.
+            val masks = Array.fill(receivers.size)(new Array[Boolean](rowCount))
+            var i = 0
+            while (i < rowCount) {
+              val t = ArrowUtils.getTexeraTuple(i, root)
+              partitioner.getBucketIndex(t).foreach(b => masks(b)(i) = true)
+              i += 1
+            }
+            var b = 0
+            while (b < receivers.size) {
+              val sub = ArrowUtils.selectRows(root, masks(b), emitAllocator())
+              try {
+                if (sub.getRowCount > 0) {
+                  outputGateway.sendTo(
+                    receivers(b),
+                    ColumnarFrame(ArrowUtils.serializeRoot(sub), sub.getRowCount, texeraSchema)
+                  )
+                }
+              } finally sub.close()
+              b += 1
+            }
+        }
+    }
+
+    // Materialize to result storage for the GUI result panel. Hand the whole
+    // Arrow batch to the writer thread(s), which decode off the DP thread. Reuse
+    // the bytes already serialized for the wire when available.
+    if (outputPortResultWriterThreads.nonEmpty && rowCount > 0) {
+      saveArrowBatchToStorageIfNeeded(wholeBatchBytes)
+    }
+  }
+
+  // Columnar batches pending emit, drained one-per-step by the DP loop so that
+  // backpressure applies between batches (unlike a synchronous emit).
+  private var columnarOutputIter: Iterator[VectorSchemaRoot] = Iterator.empty
+  def setColumnarOutput(it: Iterator[VectorSchemaRoot]): Unit = columnarOutputIter = it
+  def hasUnfinishedColumnarOutput: Boolean = columnarOutputIter.hasNext
+  def emitOneColumnarBatch(): Unit = {
+    val root = columnarOutputIter.next()
+    emitColumnarBatch(root)
+    root.close()
+  }
+
   def addPort(portId: PortIdentity, schema: Schema, storageURIBaseOption: Option[URI]): Unit = {
     // each port can only be added and initialized once.
     if (this.ports.contains(portId)) {
@@ -239,10 +326,19 @@ class OutputManager(
     }).foreach({
       case (portId, writerThread) =>
         // write to storage in a separate thread
-        writerThread.queue.put(Left(tuple))
+        writerThread.queue.put(Left(RowWriteItem(tuple)))
     })
   }
 
+  // Columnar sink: hand the whole Arrow batch to each result writer thread,
+  // which decodes it to rows off the DP thread. Avoids the per-row decode on
+  // the hot path and reuses the batch bytes already produced for the wire.
+  def saveArrowBatchToStorageIfNeeded(arrowIpcBytes: Array[Byte]): Unit = {
+    outputPortResultWriterThreads.values.foreach(
+      _.queue.put(Left(ArrowBatchWriteItem(arrowIpcBytes)))
+    )
+  }
+
   private def saveStateToStorageIfNeeded(
       state: State,
       loopCounter: Long,
@@ -254,7 +350,9 @@ class OutputManager(
     // downstream operator (and every worker reading the materialization)
     // needs the full set. The loop envelope is materialized as its own
     // columns so the downstream reader can rebuild it.
-    stateWriterThreads.values.foreach(_.queue.put(Left(state.toTuple(loopCounter, loopStartId))))
+    stateWriterThreads.values.foreach(
+      _.queue.put(Left(RowWriteItem(state.toTuple(loopCounter, loopStartId))))
+    )
   }
 
   /**

amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClient.scala

@@ -124,6 +124,18 @@ class PythonProxyClient(portNumberPromise: Promise[Int], val actorId: ActorVirtu
     dataPayload match {
       case DataFrame(frame) =>
         writeArrowStream(mutable.Queue(ArraySeq.unsafeWrapArray(frame): _*), from, "Data")
+      case ColumnarFrame(bytes, _, _, ver) =>
+        require(
+          ver == ColumnarFrame.CurrentFormatVersion,
+          s"unsupported ColumnarFrame format version $ver " +
+            s"(this worker supports ${ColumnarFrame.CurrentFormatVersion})"
+        )
+        // Columnar wire reaching the Python path: pass the Arrow batch straight
+        // to Flight (no tuple round-trip). deserializeRootFold keeps the root
+        // alive for the put and closes it after.
+        ArrowUtils.deserializeRootFold(bytes, allocator) { root =>
+          writeArrowRoot(root, from, "Data")
+        }
       case StateFrame(state, loopCounter, loopStartId) =>
         // The Arrow wire format for states IS the State row (content +
         // loop_counter + loop_start_id), so the envelope rides its own
@@ -207,6 +219,24 @@ class PythonProxyClient(portNumberPromise: Promise[Int], val actorId: ActorVirtu
     result
   }
 
+  // Send an already-built Arrow batch to the Python worker over Flight, without
+  // decoding to tuples first. Used for the columnar wire passthrough.
+  private def writeArrowRoot(
+      root: VectorSchemaRoot,
+      from: ChannelIdentity,
+      payloadType: String
+  ): Unit = {
+    val descriptor = FlightDescriptor.command(PythonDataHeader(from, payloadType).toByteArray)
+    val flightListener = new SyncPutListener
+    val writer = flightClient.startPut(descriptor, root, flightListener)
+    writer.putNext()
+    writer.completed()
+    val ackMsgBuf: ArrowBuf = flightListener.poll(5, TimeUnit.SECONDS).getApplicationMetadata
+    pythonQueueInMemSize.set(ackMsgBuf.getLong(0))
+    ackMsgBuf.close()
+    flightListener.close()
+  }
+
   private def writeArrowStream(
       tuples: mutable.Queue[Tuple],
       from: ChannelIdentity,

amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/Partitioner.scala

@@ -24,7 +24,8 @@ import org.apache.texera.amber.core.state.State
 import org.apache.texera.amber.core.tuple.Tuple
 import org.apache.texera.amber.core.virtualidentity.ActorVirtualIdentity
 import org.apache.texera.amber.engine.architecture.messaginglayer.NetworkOutputGateway
-import org.apache.texera.amber.engine.common.ambermessage.{DataFrame, StateFrame}
+import org.apache.texera.amber.engine.common.ambermessage.{ColumnarFrame, DataFrame, StateFrame}
+import org.apache.texera.amber.util.ArrowUtils
 
 import scala.collection.mutable.ArrayBuffer
 
@@ -57,9 +58,25 @@ class NetworkOutputBuffer(
 
   def flush(): Unit = {
     if (buffer.nonEmpty) {
-      dataOutputPort.sendTo(to, DataFrame(buffer.toArray))
+      val batch = buffer.toArray
+      // Columnar wire (flagged): send the batch as an Arrow IPC ColumnarFrame.
+      val payload =
+        if (NetworkOutputBuffer.columnarWire)
+          ColumnarFrame(
+            ArrowUtils.serializeTuples(batch.head.getSchema, batch),
+            batch.length,
+            batch.head.getSchema
+          )
+        else DataFrame(batch)
+      dataOutputPort.sendTo(to, payload)
       buffer = new ArrayBuffer[Tuple]()
     }
   }
 
 }
+
+object NetworkOutputBuffer {
+  // Opt-in for the Arrow columnar wire format (default off = row DataFrame),
+  // from application.conf (columnar.enable-columnar-wire), COLUMNAR_WIRE overrides.
+  val columnarWire: Boolean = ApplicationConfig.enableColumnarWire
+}

amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DPThread.scala

@@ -165,7 +165,8 @@ class DPThread(
       var channelId: ChannelIdentity = null
       var msgOpt: Option[WorkflowFIFOMessage] = None
       if (
-        dp.inputManager.hasUnfinishedInput || dp.outputManager.hasUnfinishedOutput || dp.pauseManager.isPaused
+        dp.inputManager.hasUnfinishedInput || dp.outputManager.hasUnfinishedOutput ||
+        dp.outputManager.hasUnfinishedColumnarOutput || dp.pauseManager.isPaused
       ) {
         dp.inputGateway.tryPickControlChannel match {
           case Some(channel) =>

amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessor.scala

@@ -21,7 +21,12 @@ package org.apache.texera.amber.engine.architecture.worker
 
 import com.softwaremill.macwire.wire
 import io.grpc.MethodDescriptor
-import org.apache.texera.amber.core.executor.OperatorExecutor
+import org.apache.texera.amber.core.executor.{
+  ColumnarOperatorExecutor,
+  ColumnarResult,
+  OperatorExecutor
+}
+import org.apache.texera.amber.engine.architecture.sendsemantics.partitioners.NetworkOutputBuffer
 import org.apache.texera.amber.core.state.State
 import org.apache.texera.amber.core.tuple._
 import org.apache.texera.amber.core.virtualidentity.{
@@ -59,6 +64,7 @@ import org.apache.texera.amber.engine.common.ambermessage._
 import org.apache.texera.amber.engine.common.statetransition.WorkerStateManager
 import org.apache.texera.amber.engine.common.virtualidentity.util.COORDINATOR
 import org.apache.texera.amber.error.ErrorUtils.{mkConsoleMessage, safely}
+import org.apache.texera.amber.util.ArrowUtils
 
 import java.util.concurrent.LinkedBlockingQueue
 
@@ -169,6 +175,7 @@ class DataProcessor(
         sendECMToDataChannels(METHOD_END_CHANNEL, PORT_ALIGNMENT)
         // Send Completed signal to worker actor.
         executor.close()
+        outputManager.closeColumnarResources()
         adaptiveBatchingMonitor.stopAdaptiveBatching()
         stateManager.transitTo(COMPLETED)
         logger.debug(
@@ -201,7 +208,9 @@ class DataProcessor(
 
   def continueDataProcessing(): Unit = {
     val dataProcessingStartTime = System.nanoTime()
-    if (outputManager.hasUnfinishedOutput) {
+    if (outputManager.hasUnfinishedColumnarOutput) {
+      outputManager.emitOneColumnarBatch()
+    } else if (outputManager.hasUnfinishedOutput) {
       outputOneTuple()
     } else {
       processInputTuple(inputManager.getNextTuple)
@@ -217,25 +226,79 @@ class DataProcessor(
     val portId = this.inputGateway.getChannel(channelId).getPortId
     dataPayload match {
       case DataFrame(tuples) =>
-        stateManager.conditionalTransitTo(
-          READY,
-          RUNNING,
-          () => {
-            val (state, stateVersion) = stateManager.getStateWithVersion
-            asyncRPCClient.coordinatorInterface.workerStateUpdated(
-              WorkerStateUpdatedRequest(state, stateVersion),
-              asyncRPCClient.mkContext(COORDINATOR)
-            )
-          }
+        processTupleBatch(channelId, portId, tuples)
+      case ColumnarFrame(bytes, _, _, ver) =>
+        require(
+          ver == ColumnarFrame.CurrentFormatVersion,
+          s"unsupported ColumnarFrame format version $ver " +
+            s"(this worker supports ${ColumnarFrame.CurrentFormatVersion})"
         )
-        inputManager.initBatch(channelId, tuples)
-        processInputTuple(inputManager.getNextTuple)
+        executor match {
+          // Native-Arrow path: consume the batch directly, emit a filtered batch.
+          case c: ColumnarOperatorExecutor if NetworkOutputBuffer.columnarWire =>
+            c.processColumnarBatch(bytes, portId.id) match {
+              case ColumnarResult.Emit(result) =>
+                logColumnarModeOnce(active = true, "")
+                statisticsManager.increaseInputStatistics(portId, bytes.length.toLong)
+                outputManager.setColumnarOutput(Iterator.single(result))
+              case ColumnarResult.EmitRows(rows) =>
+                logColumnarModeOnce(active = true, "")
+                statisticsManager.increaseInputStatistics(portId, bytes.length.toLong)
+                outputManager.outputIterator.setTupleOutput(rows)
+              case ColumnarResult.Consumed =>
+                logColumnarModeOnce(active = true, "")
+                statisticsManager.increaseInputStatistics(portId, bytes.length.toLong)
+              case ColumnarResult.Unsupported =>
+                logColumnarModeOnce(
+                  active = false,
+                  "operator has no native columnar path for this batch"
+                )
+                processTupleBatch(channelId, portId, ArrowUtils.deserializeTuples(bytes))
+            }
+          case _ =>
+            processTupleBatch(channelId, portId, ArrowUtils.deserializeTuples(bytes))
+        }
       case StateFrame(state, loopCounter, loopStartId) =>
         processInputState(state, portId.id, loopCounter, loopStartId)
     }
     statisticsManager.increaseDataProcessingTime(System.nanoTime() - dataProcessingStartTime)
   }
 
+  private def processTupleBatch(
+      channelId: ChannelIdentity,
+      portId: PortIdentity,
+      tuples: Array[Tuple]
+  ): Unit = {
+    stateManager.conditionalTransitTo(
+      READY,
+      RUNNING,
+      () => {
+        val (state, stateVersion) = stateManager.getStateWithVersion
+        asyncRPCClient.coordinatorInterface.workerStateUpdated(
+          WorkerStateUpdatedRequest(state, stateVersion),
+          asyncRPCClient.mkContext(COORDINATOR)
+        )
+      }
+    )
+    inputManager.initBatch(channelId, tuples)
+    // Whole-batch path if the executor supports it, else per-tuple; errors fall back.
+    val vectorized: Option[Iterator[(TupleLike, Option[PortIdentity])]] =
+      try executor.processBatchMultiPort(tuples, portId.id)
+      catch safely { case _ => None }
+    vectorized match {
+      case Some(outputIter) =>
+        var i = 0
+        while (i < tuples.length) {
+          statisticsManager.increaseInputStatistics(portId, tuples(i).inMemSize)
+          i += 1
+        }
+        inputManager.skipToEnd()
+        outputManager.outputIterator.setTupleOutput(outputIter)
+      case None =>
+        processInputTuple(inputManager.getNextTuple)
+    }
+  }
+
   def processECM(
       channelId: ChannelIdentity,
       ecm: EmbeddedControlMessage,
@@ -309,6 +372,17 @@ class DataProcessor(
       }
   }
 
+  // Log the columnar-wire decision once per worker so a silent fall back to the
+  // row path is visible in the logs.
+  @transient private var columnarModeLogged = false
+  private[architecture] def logColumnarModeOnce(active: Boolean, reason: String): Unit = {
+    if (!columnarModeLogged) {
+      columnarModeLogged = true
+      if (active) logger.info(s"columnar wire active for $executor")
+      else logger.info(s"columnar wire requested but using row path for $executor: $reason")
+    }
+  }
+
   def handleExecutorException(e: Throwable): Unit = {
     asyncRPCClient.coordinatorInterface.consoleMessageTriggered(
       ConsoleMessageTriggeredRequest(mkConsoleMessage(actorId, e)),