You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@druid.apache.org by gi...@apache.org on 2026/01/14 18:21:01 UTC
(druid) branch master updated: WireTransferable result batches. (#18909)
This is an automated email from the ASF dual-hosted git repository.
gian pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 10e56bf005e WireTransferable result batches. (#18909)
10e56bf005e is described below
commit 10e56bf005e2d73e884babbfaaf124829db728ef
Author: Gian Merlino <gi...@gmail.com>
AuthorDate: Wed Jan 14 10:20:54 2026 -0800
WireTransferable result batches. (#18909)
This patch introduces a "WireTransferable" interface that is used to
serialize batches of results represented by RowsAndColumns objects.
This allows query implementations to control the exact manner in which
their result batches are modeled, serialized, and deserialized, which
enables greater efficiency.
For example, this would allow an MSQ implementation of TopN (which does
not currently exist) to use a data structure between leaf and merge stages
that matches what it uses in its native implementation. It would also
allow this data structure to be serialized efficiently, without needing
to use JSON/Smile.
WireTransferable is used in two places:
1) In ObjectMappers, when they read or write RowsAndColumns objects.
2) In MSQ channels, when frames or other batches are read or written to
disk or network.
Prior to this patch, the ObjectMapper path only handled FrameRowsAndColumns,
and the channel path only handled Frames, in both cases using a
Frame-specific serialization format. These Frame-specific serialization
formats are still used by default. The new FrameWireTransferable format
can be enabled in both places by setting
druid.serde.rac.useLegacyFrameSerialization = false.
As part of this change, the ReadableFrameChannels and WritableFrameChannels
used by MSQ are generalized to allow transmitting any RowsAndColumns,
not just Frames. When Frames are transmitted, they are enclosed in a thin
FrameRowsAndColumns wrapper. In the interest of minimizing the size of the
diff, the names of the classes are not changed.
---
.../frame/FrameChannelMergerBenchmark.java | 2 +-
.../apache/druid/msq/counters/ChannelCounters.java | 11 +-
.../druid/msq/dart/worker/DartFrameContext.java | 7 +
.../org/apache/druid/msq/exec/ControllerImpl.java | 9 +-
.../msq/exec/ControllerQueryResultsReader.java | 2 +-
.../org/apache/druid/msq/exec/FrameContext.java | 7 +
.../org/apache/druid/msq/exec/RunWorkOrder.java | 18 +-
.../java/org/apache/druid/msq/exec/WorkerImpl.java | 19 +-
.../msq/exec/std/StandardShuffleOperations.java | 4 +-
.../msq/indexing/CountingReadableFrameChannel.java | 10 +-
.../msq/indexing/CountingWritableFrameChannel.java | 8 +-
.../druid/msq/indexing/IndexerFrameContext.java | 7 +
.../KeyStatisticsCollectionProcessor.java | 8 +-
.../processor/SegmentGeneratorFrameProcessor.java | 2 +-
.../BroadcastJoinSegmentMapFnProcessor.java | 2 +-
.../WindowOperatorQueryFrameProcessor.java | 5 +-
.../querykit/common/OffsetLimitFrameProcessor.java | 5 +-
.../common/SortMergeJoinFrameProcessor.java | 5 +-
.../groupby/GroupByPostShuffleFrameProcessor.java | 5 +-
.../groupby/GroupByPreShuffleFrameProcessor.java | 5 +-
.../results/ExportResultsFrameProcessor.java | 2 +-
.../results/QueryResultsFrameProcessor.java | 8 +-
.../msq/querykit/scan/ScanQueryFrameProcessor.java | 5 +-
.../input/DurableStorageInputChannelFactory.java | 25 +-
...ableStorageQueryResultsInputChannelFactory.java | 7 +-
.../DurableStorageStageInputChannelFactory.java | 7 +-
.../shuffle/input/WorkerInputChannelFactory.java | 14 +-
.../shuffle/output/ChannelStageOutputReader.java | 11 +-
.../output/DurableStorageOutputChannelFactory.java | 23 +-
...bleStorageQueryResultsOutputChannelFactory.java | 18 +-
.../DurableStorageTaskOutputChannelFactory.java | 27 +-
.../msq/shuffle/output/FileStageOutputReader.java | 10 +-
.../shuffle/output/FutureReadableFrameChannel.java | 4 +-
.../msq/shuffle/output/StageOutputHolder.java | 14 +-
.../msq/sql/resources/SqlStatementResource.java | 14 +-
.../BroadcastJoinSegmentMapFnProcessorTest.java | 12 +-
.../msq/querykit/ChainedProcessorManagerTest.java | 8 +-
.../querykit/NonFailingWritableFrameChannel.java | 6 +-
.../druid/msq/rpc/BaseWorkerClientImplTest.java | 16 +-
.../DurableStorageOutputChannelFactoryTest.java | 4 +-
.../output/ChannelStageOutputReaderTest.java | 20 +-
.../druid/msq/sql/MSQTaskQueryMakerTest.java | 8 +-
.../resources/SqlMSQStatementResourcePostTest.java | 6 +-
.../sql/resources/SqlStatementResourceTest.java | 3 +-
.../org/apache/druid/msq/test/MSQTestBase.java | 10 +-
.../druid/msq/test/MSQTestWorkerContext.java | 7 +
.../msq/util/SqlStatementResourceHelperTest.java | 11 +-
.../main/java/org/apache/druid/frame/Frame.java | 195 ++++---------
.../frame/channel/BlockingQueueFrameChannel.java | 19 +-
.../channel/ComposingReadableFrameChannel.java | 4 +-
.../channel/ComposingWritableFrameChannel.java | 15 +-
...ableStoragePartitionedReadableFrameChannel.java | 12 +-
.../druid/frame/channel/FrameChannelSequence.java | 10 +-
.../channel/ReadableByteChunksFrameChannel.java | 92 +++++--
.../frame/channel/ReadableConcatFrameChannel.java | 4 +-
.../frame/channel/ReadableFileFrameChannel.java | 23 +-
.../druid/frame/channel/ReadableFrameChannel.java | 50 +++-
.../channel/ReadableInputStreamFrameChannel.java | 17 +-
.../frame/channel/ReadableNilFrameChannel.java | 4 +-
.../druid/frame/channel/WritableFrameChannel.java | 47 +++-
.../frame/channel/WritableFrameFileChannel.java | 5 +-
.../org/apache/druid/frame/file/FrameFile.java | 85 ++++--
.../apache/druid/frame/file/FrameFileWriter.java | 112 ++++++--
.../frame/processor/FileOutputChannelFactory.java | 20 +-
.../druid/frame/processor/FrameChannelBatcher.java | 2 +-
.../processor/FrameChannelHashPartitioner.java | 2 +-
.../druid/frame/processor/FrameChannelMerger.java | 23 +-
.../druid/frame/processor/FrameChannelMixer.java | 8 +-
.../druid/frame/processor/OutputChannel.java | 11 +-
.../apache/druid/frame/processor/SuperSorter.java | 7 +-
.../apache/druid/frame/wire/FrameCompression.java | 221 +++++++++++++++
.../druid/frame/wire/FrameWireTransferable.java | 157 +++++++++++
.../apache/druid/guice/DruidSecondaryModule.java | 115 +++++++-
.../jackson/DruidDefaultSerializersModule.java | 4 -
.../druid/java/util/common/ByteBufferUtils.java | 25 ++
.../query/rowsandcols/AppendableMapOfColumns.java | 10 +
.../query/rowsandcols/ConcatRowsAndColumns.java | 7 +
.../rowsandcols/CursorFactoryRowsAndColumns.java | 2 +-
.../query/rowsandcols/EmptyRowsAndColumns.java | 8 +
.../query/rowsandcols/LimitedRowsAndColumns.java | 8 +
.../rowsandcols/MapOfColumnsRowsAndColumns.java | 2 +-
.../rowsandcols/RearrangedRowsAndColumns.java | 7 +
.../druid/query/rowsandcols/RowsAndColumns.java | 84 +-----
.../concrete/AbstractFrameRowsAndColumns.java | 52 ++--
.../concrete/ColumnBasedFrameRowsAndColumns.java | 6 +-
.../rowsandcols/concrete/FrameRowsAndColumns.java | 18 ++
.../concrete/RowBasedFrameRowsAndColumns.java | 6 +-
.../rowsandcols/semantic/WireTransferable.java | 157 +++++++++++
.../serde/RowsAndColumnsDeserializer.java | 93 +++++++
.../serde/RowsAndColumnsSerializer.java | 101 +++++++
.../rowsandcols/serde/WireTransferableContext.java | 107 ++++++++
.../java/org/apache/druid/frame/FrameTest.java | 7 +-
.../channel/ComposingWritableFrameChannelTest.java | 14 +-
.../ReadableByteChunksFrameChannelTest.java | 249 ++++++++++++++---
.../ReadableInputStreamFrameChannelTest.java | 23 +-
.../file/FrameFileHttpResponseHandlerTest.java | 2 +-
.../org/apache/druid/frame/file/FrameFileTest.java | 102 +++++--
.../frame/file/FrameFileWireTransferableTest.java | 114 ++++++++
.../druid/frame/file/FrameFileWriterTest.java | 16 +-
.../ComposingOutputChannelFactoryTest.java | 14 +-
.../processor/FileOutputChannelFactoryTest.java | 3 +-
.../processor/FrameProcessorExecutorTest.java | 14 +-
.../frame/processor/OutputChannelFactoryTest.java | 7 +-
.../processor/ReadableFileFrameChannelTest.java | 4 +-
.../druid/frame/processor/SuperSorterTest.java | 20 +-
.../test/AlwaysAsyncReadableFrameChannel.java | 6 +-
.../processor/test/ChompingFrameProcessor.java | 2 +-
.../processor/test/SuperBlasterFrameProcessor.java | 8 +-
.../apache/druid/frame/testutil/FrameTestUtil.java | 46 +++-
.../druid/frame/wire/FrameCompressionTest.java | 302 +++++++++++++++++++++
.../druid/guice/DruidSecondaryModuleTest.java | 145 +++++++++-
.../druid/jackson/DefaultObjectMapperTest.java | 27 +-
.../query/rowsandcols/NoAsRowsAndColumns.java | 15 +
.../serde/FrameWireTransferableTest.java | 156 +++++++++++
.../serde/WireTransferableContextTest.java | 126 +++++++++
115 files changed, 3122 insertions(+), 698 deletions(-)
diff --git a/benchmarks/src/test/java/org/apache/druid/benchmark/frame/FrameChannelMergerBenchmark.java b/benchmarks/src/test/java/org/apache/druid/benchmark/frame/FrameChannelMergerBenchmark.java
index ce723e23af6..b24dc2f2a66 100644
--- a/benchmarks/src/test/java/org/apache/druid/benchmark/frame/FrameChannelMergerBenchmark.java
+++ b/benchmarks/src/test/java/org/apache/druid/benchmark/frame/FrameChannelMergerBenchmark.java
@@ -367,7 +367,7 @@ public class FrameChannelMergerBenchmark
while (!outputChannel.readable().isFinished()) {
FutureUtils.getUnchecked(outputChannel.readable().readabilityFuture(), false);
if (outputChannel.readable().canRead()) {
- final Frame frame = outputChannel.readable().read();
+ final Frame frame = outputChannel.readable().readFrame();
blackhole.consume(frame);
}
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/counters/ChannelCounters.java b/multi-stage-query/src/main/java/org/apache/druid/msq/counters/ChannelCounters.java
index 1cf017635a6..c51ec0cad13 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/counters/ChannelCounters.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/counters/ChannelCounters.java
@@ -27,6 +27,7 @@ import com.google.errorprone.annotations.concurrent.GuardedBy;
import it.unimi.dsi.fastutil.longs.LongArrayList;
import it.unimi.dsi.fastutil.longs.LongList;
import org.apache.druid.frame.Frame;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import javax.annotation.Nullable;
import java.util.Arrays;
@@ -78,9 +79,15 @@ public class ChannelCounters implements QueryCounter
add(NO_PARTITION, nRows, nBytes, 0, 1);
}
- public void addFrame(final int partitionNumber, final Frame frame)
+ /**
+ * Increment counts by one frame, and if this {@link RowsAndColumns} is a {@link Frame}, also increment
+ * bytes by {@link Frame#numBytes()}.
+ */
+ public void addRAC(final RowsAndColumns rac, final int partitionNumber)
{
- add(partitionNumber, frame.numRows(), frame.numBytes(), 1, 0);
+ final Frame frame = rac.as(Frame.class);
+ final long numBytes = frame != null ? frame.numBytes() : 0;
+ add(partitionNumber, rac.numRows(), numBytes, 1, 0);
}
public ChannelCounters setTotalFiles(final long nFiles)
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartFrameContext.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartFrameContext.java
index a1fcdf18be7..3ac9d6bb402 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartFrameContext.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartFrameContext.java
@@ -33,6 +33,7 @@ import org.apache.druid.msq.exec.WorkerStorageParameters;
import org.apache.druid.msq.kernel.StageId;
import org.apache.druid.query.groupby.GroupingEngine;
import org.apache.druid.query.policy.PolicyEnforcer;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.segment.IndexIO;
import org.apache.druid.segment.IndexMerger;
import org.apache.druid.segment.SegmentWrangler;
@@ -135,6 +136,12 @@ public class DartFrameContext implements FrameContext
return workerContext.jsonMapper();
}
+ @Override
+ public WireTransferableContext wireTransferableContext()
+ {
+ return workerContext.injector().getInstance(WireTransferableContext.class);
+ }
+
@Override
public IndexIO indexIO()
{
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
index 6c058ef0df4..bf02ebebd82 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
@@ -171,6 +171,7 @@ import org.apache.druid.query.QueryContext;
import org.apache.druid.query.QueryContexts;
import org.apache.druid.query.aggregation.AggregatorFactory;
import org.apache.druid.query.groupby.GroupByQuery;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.segment.IndexSpec;
import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.column.RowSignature;
@@ -2781,16 +2782,20 @@ public class ControllerImpl implements Controller
final InputChannelFactory inputChannelFactory;
+ final WireTransferableContext wireTransferableContext =
+ context.injector().getInstance(WireTransferableContext.class);
+
if (queryKernelConfig.isDurableStorage()
|| MSQControllerTask.writeFinalStageResultsToDurableStorage(querySpec.getDestination())) {
inputChannelFactory = DurableStorageInputChannelFactory.createStandardImplementation(
queryId(),
MSQTasks.makeStorageConnector(context.injector()),
closer,
- MSQControllerTask.writeFinalStageResultsToDurableStorage(querySpec.getDestination())
+ MSQControllerTask.writeFinalStageResultsToDurableStorage(querySpec.getDestination()),
+ wireTransferableContext
);
} else {
- inputChannelFactory = new WorkerInputChannelFactory(netClient, () -> taskIds);
+ inputChannelFactory = new WorkerInputChannelFactory(netClient, () -> taskIds, wireTransferableContext);
}
final FrameProcessorExecutor resultReaderExec = createResultReaderExec(queryId());
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerQueryResultsReader.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerQueryResultsReader.java
index ae24704c9d8..6ef4b743c21 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerQueryResultsReader.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerQueryResultsReader.java
@@ -121,7 +121,7 @@ public class ControllerQueryResultsReader implements FrameProcessor<Void>
queryListener.onResultsComplete();
return ReturnOrAwait.returnObject(null);
} else {
- final Frame frame = in.read();
+ final Frame frame = in.readFrame();
Yielder<Object[]> rowYielder = Yielders.each(
SqlStatementResourceHelper.getResultSequence(
frame,
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/FrameContext.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/FrameContext.java
index 501616b1190..4826872b35d 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/FrameContext.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/FrameContext.java
@@ -24,6 +24,8 @@ import org.apache.druid.client.coordinator.CoordinatorClient;
import org.apache.druid.msq.kernel.WorkOrder;
import org.apache.druid.query.groupby.GroupingEngine;
import org.apache.druid.query.policy.PolicyEnforcer;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.segment.IndexIO;
import org.apache.druid.segment.IndexMerger;
import org.apache.druid.segment.SegmentWrangler;
@@ -72,6 +74,11 @@ public interface FrameContext extends Closeable
ObjectMapper jsonMapper();
+ /**
+ * Context for {@link WireTransferable} serde.
+ */
+ WireTransferableContext wireTransferableContext();
+
IndexIO indexIO();
File persistDir();
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java
index 7fc6b83d119..f73ffa010fc 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/RunWorkOrder.java
@@ -431,7 +431,12 @@ public class RunWorkOrder
case LOCAL_STORAGE:
final File fileChannelDirectory =
frameContext.tempDir(StringUtils.format("output_stage_%06d", workOrder.getStageNumber()));
- outputChannelFactory = new FileOutputChannelFactory(fileChannelDirectory, frameSize, null);
+ outputChannelFactory = new FileOutputChannelFactory(
+ fileChannelDirectory,
+ frameSize,
+ null,
+ frameContext.wireTransferableContext()
+ );
break;
case DURABLE_STORAGE_INTERMEDIATE:
@@ -464,8 +469,12 @@ public class RunWorkOrder
final int frameSize = frameContext.memoryParameters().getFrameSize();
final File fileChannelDirectory =
new File(tempDir, StringUtils.format("intermediate-stage-%06d", workOrder.getStageNumber()));
- final FileOutputChannelFactory fileOutputChannelFactory =
- new FileOutputChannelFactory(fileChannelDirectory, frameSize, intermediateByteTracker);
+ final FileOutputChannelFactory fileOutputChannelFactory = new FileOutputChannelFactory(
+ fileChannelDirectory,
+ frameSize,
+ intermediateByteTracker,
+ frameContext.wireTransferableContext()
+ );
if (workOrder.getOutputChannelMode().isDurable()
&& frameContext.storageParameters().isIntermediateStorageLimitConfigured()) {
@@ -497,7 +506,8 @@ public class RunWorkOrder
frameSize,
MSQTasks.makeStorageConnector(workerContext.injector()),
tmpDir,
- isQueryResults
+ isQueryResults,
+ frameContext.wireTransferableContext()
);
}
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java
index e75d096e398..ce4792cb72f 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java
@@ -82,6 +82,7 @@ import org.apache.druid.query.PrioritizedCallable;
import org.apache.druid.query.PrioritizedRunnable;
import org.apache.druid.query.QueryContext;
import org.apache.druid.query.QueryProcessingPool;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.server.DruidNode;
import org.joda.time.Interval;
@@ -839,7 +840,7 @@ public class WorkerImpl implements Worker
return new WorkerOrLocalInputChannelFactory(
id(),
workerIds,
- new WorkerInputChannelFactory(workerClient, workerIds),
+ new WorkerInputChannelFactory(workerClient, workerIds, getWireTransferableContext()),
this::getOrCreateStageOutputHolder
);
@@ -849,7 +850,8 @@ public class WorkerImpl implements Worker
task.getControllerTaskId(),
MSQTasks.makeStorageConnector(context.injector()),
closer,
- outputChannelMode == OutputChannelMode.DURABLE_STORAGE_QUERY_RESULTS
+ outputChannelMode == OutputChannelMode.DURABLE_STORAGE_QUERY_RESULTS,
+ getWireTransferableContext()
);
default:
@@ -974,8 +976,17 @@ public class WorkerImpl implements Worker
private StageOutputHolder getOrCreateStageOutputHolder(final StageId stageId, final int partitionNumber)
{
- return stageOutputs.computeIfAbsent(stageId, ignored1 -> new ConcurrentHashMap<>())
- .computeIfAbsent(partitionNumber, ignored -> new StageOutputHolder());
+ return stageOutputs
+ .computeIfAbsent(stageId, ignored1 -> new ConcurrentHashMap<>())
+ .computeIfAbsent(partitionNumber, ignored -> new StageOutputHolder(getWireTransferableContext()));
+ }
+
+ /**
+ * Retrieve {@link WireTransferableContext} from our injector.
+ */
+ private WireTransferableContext getWireTransferableContext()
+ {
+ return context.injector().getInstance(WireTransferableContext.class);
}
/**
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/std/StandardShuffleOperations.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/std/StandardShuffleOperations.java
index d3daec7d606..10cbdabfd66 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/std/StandardShuffleOperations.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/std/StandardShuffleOperations.java
@@ -28,7 +28,7 @@ import org.apache.druid.common.guava.FutureUtils;
import org.apache.druid.frame.FrameType;
import org.apache.druid.frame.allocation.ArenaMemoryAllocatorFactory;
import org.apache.druid.frame.channel.BlockingQueueFrameChannel;
-import org.apache.druid.frame.channel.FrameWithPartition;
+import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.key.ClusterByPartitions;
import org.apache.druid.frame.processor.Bouncer;
import org.apache.druid.frame.processor.FrameChannelHashPartitioner;
@@ -159,7 +159,7 @@ public class StandardShuffleOperations
final OutputChannel outputChannel = OutputChannel.readOnly(
channel.readable(),
- FrameWithPartition.NO_PARTITION
+ WritableFrameChannel.NO_PARTITION
);
return new ResultAndChannels<>(
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/CountingReadableFrameChannel.java b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/CountingReadableFrameChannel.java
index 6462832acbf..2f0ed8c1f07 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/CountingReadableFrameChannel.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/CountingReadableFrameChannel.java
@@ -20,9 +20,9 @@
package org.apache.druid.msq.indexing;
import com.google.common.util.concurrent.ListenableFuture;
-import org.apache.druid.frame.Frame;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
public class CountingReadableFrameChannel implements ReadableFrameChannel
{
@@ -54,11 +54,11 @@ public class CountingReadableFrameChannel implements ReadableFrameChannel
}
@Override
- public Frame read()
+ public RowsAndColumns read()
{
- final Frame frame = baseChannel.read();
- channelCounters.addFrame(partitionNumber, frame);
- return frame;
+ final RowsAndColumns rac = baseChannel.read();
+ channelCounters.addRAC(rac, partitionNumber);
+ return rac;
}
@Override
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/CountingWritableFrameChannel.java b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/CountingWritableFrameChannel.java
index 1b4963ac674..17528ee4fc2 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/CountingWritableFrameChannel.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/CountingWritableFrameChannel.java
@@ -20,9 +20,9 @@
package org.apache.druid.msq.indexing;
import com.google.common.util.concurrent.ListenableFuture;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.msq.counters.ChannelCounters;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import javax.annotation.Nullable;
import java.io.IOException;
@@ -45,10 +45,10 @@ public class CountingWritableFrameChannel implements WritableFrameChannel
}
@Override
- public void write(FrameWithPartition frame) throws IOException
+ public void write(RowsAndColumns rac, int partitionNum) throws IOException
{
- baseChannel.write(frame);
- channelCounters.addFrame(partitionNumber == null ? frame.partition() : partitionNumber, frame.frame());
+ baseChannel.write(rac, partitionNum);
+ channelCounters.addRAC(rac, partitionNumber == null ? partitionNum : partitionNumber);
}
@Override
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerFrameContext.java b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerFrameContext.java
index 9e215a09e91..4af15e79aee 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerFrameContext.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerFrameContext.java
@@ -32,6 +32,7 @@ import org.apache.druid.msq.exec.WorkerStorageParameters;
import org.apache.druid.msq.kernel.StageId;
import org.apache.druid.query.groupby.GroupingEngine;
import org.apache.druid.query.policy.PolicyEnforcer;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.segment.IndexIO;
import org.apache.druid.segment.IndexMerger;
import org.apache.druid.segment.SegmentWrangler;
@@ -138,6 +139,12 @@ public class IndexerFrameContext implements FrameContext
return context.jsonMapper();
}
+ @Override
+ public WireTransferableContext wireTransferableContext()
+ {
+ return context.injector().getInstance(WireTransferableContext.class);
+ }
+
@Override
public IndexIO indexIO()
{
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/processor/KeyStatisticsCollectionProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/processor/KeyStatisticsCollectionProcessor.java
index ecb689e6c65..8cc9f799df6 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/processor/KeyStatisticsCollectionProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/processor/KeyStatisticsCollectionProcessor.java
@@ -25,7 +25,6 @@ import org.apache.datasketches.memory.Memory;
import org.apache.druid.frame.Frame;
import org.apache.druid.frame.FrameType;
import org.apache.druid.frame.allocation.MemoryRange;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.key.ClusterBy;
@@ -108,7 +107,7 @@ public class KeyStatisticsCollectionProcessor implements FrameProcessor<ClusterB
return ReturnOrAwait.returnObject(clusterByStatisticsCollector);
}
- final Frame frame = inputChannel.read();
+ final Frame frame = inputChannel.readFrame();
final Cursor cursor = FrameProcessors.makeCursor(frame, frameReader);
final IntSupplier rowWeightSupplier =
makeRowWeightSupplier(frameReader, frame.type(), cursor.getColumnSelectorFactory());
@@ -119,8 +118,9 @@ public class KeyStatisticsCollectionProcessor implements FrameProcessor<ClusterB
clusterByStatisticsCollector.add(key, rowWeightSupplier.getAsInt());
}
- // Clears partition info (uses NO_PARTITION), but that's OK, because it isn't needed downstream of this processor.
- outputChannel.write(new FrameWithPartition(frame, FrameWithPartition.NO_PARTITION));
+ // Clears partition info (writes without partition number), but that's OK, because it isn't needed
+ // downstream of this processor.
+ outputChannel.write(frame);
return ReturnOrAwait.awaitAll(1);
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/processor/SegmentGeneratorFrameProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/processor/SegmentGeneratorFrameProcessor.java
index a499a93c4eb..d42ce9455fd 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/processor/SegmentGeneratorFrameProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/processor/SegmentGeneratorFrameProcessor.java
@@ -161,7 +161,7 @@ public class SegmentGeneratorFrameProcessor implements FrameProcessor<DataSegmen
log.debug("Received first frame for segment [%s].", segmentIdWithShardSpec.asSegmentId());
}
- addFrame(inChannel.read());
+ addFrame(inChannel.readFrame());
return ReturnOrAwait.awaitAll(1);
}
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BroadcastJoinSegmentMapFnProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BroadcastJoinSegmentMapFnProcessor.java
index 55d0daa50e5..d84deea0ad4 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BroadcastJoinSegmentMapFnProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BroadcastJoinSegmentMapFnProcessor.java
@@ -244,7 +244,7 @@ public class BroadcastJoinSegmentMapFnProcessor implements FrameProcessor<Segmen
while (inputChannelIterator.hasNext()) {
final int channelNumber = inputChannelIterator.nextInt();
if (sideChannelNumbers.contains(channelNumber) && channels.get(channelNumber).canRead()) {
- final Frame frame = channels.get(channelNumber).read();
+ final Frame frame = channels.get(channelNumber).readFrame();
memoryUsed += frame.numBytes();
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java
index 674e3ea4602..33c79947fb7 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java
@@ -23,7 +23,6 @@ import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.Iterables;
import it.unimi.dsi.fastutil.ints.IntSet;
import org.apache.druid.frame.Frame;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.processor.FrameProcessor;
@@ -138,7 +137,7 @@ public class WindowOperatorQueryFrameProcessor implements FrameProcessor<Object>
}
if (inputChannel.canRead()) {
- final Frame frame = inputChannel.read();
+ final Frame frame = inputChannel.readFrame();
LazilyDecoratedRowsAndColumns ldrc = convertRowFrameToRowsAndColumns(frame);
frameRowsAndColsBuilder.add(ldrc);
@@ -294,7 +293,7 @@ public class WindowOperatorQueryFrameProcessor implements FrameProcessor<Object>
return 0;
} else {
final Frame frame = Frame.wrap(frameWriter.toByteArray());
- Iterables.getOnlyElement(outputChannels()).write(new FrameWithPartition(frame, FrameWithPartition.NO_PARTITION));
+ Iterables.getOnlyElement(outputChannels()).write(frame);
frameWriter.close();
frameWriter = null;
return frame.numRows();
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/common/OffsetLimitFrameProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/common/OffsetLimitFrameProcessor.java
index 435f6a0fc84..2ed98c46dde 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/common/OffsetLimitFrameProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/common/OffsetLimitFrameProcessor.java
@@ -21,7 +21,6 @@ package org.apache.druid.msq.querykit.common;
import it.unimi.dsi.fastutil.ints.IntSet;
import org.apache.druid.frame.Frame;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.processor.FrameProcessor;
@@ -95,11 +94,11 @@ public class OffsetLimitFrameProcessor implements FrameProcessor<Object>
return ReturnOrAwait.returnObject(Unit.instance());
}
- final Frame frame = inputChannel.read();
+ final Frame frame = inputChannel.readFrame();
final Frame truncatedFrame = chopAndProcess(frame, frameReader);
if (truncatedFrame != null) {
- outputChannel.write(new FrameWithPartition(truncatedFrame, FrameWithPartition.NO_PARTITION));
+ outputChannel.write(truncatedFrame);
}
if (rowsProcessedSoFar == offset + limit) {
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/common/SortMergeJoinFrameProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/common/SortMergeJoinFrameProcessor.java
index d62f5277a0c..21b0ecd1d5f 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/common/SortMergeJoinFrameProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/common/SortMergeJoinFrameProcessor.java
@@ -25,7 +25,6 @@ import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
import it.unimi.dsi.fastutil.ints.IntSet;
import org.apache.druid.frame.Frame;
import org.apache.druid.frame.FrameType;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.key.FrameComparisonWidget;
@@ -496,7 +495,7 @@ public class SortMergeJoinFrameProcessor implements FrameProcessor<Object>
} else if (!tracker.canBufferMoreFrames()) {
return false;
} else {
- final Frame frame = channel.read();
+ final Frame frame = channel.readFrame();
if (frame.numRows() == 0) {
// Skip, read next.
@@ -531,7 +530,7 @@ public class SortMergeJoinFrameProcessor implements FrameProcessor<Object>
final Frame frame = Frame.wrap(frameWriter.toByteArray());
frameWriter.close();
frameWriter = null;
- outputChannel.write(new FrameWithPartition(frame, FrameWithPartition.NO_PARTITION));
+ outputChannel.write(frame);
}
}
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPostShuffleFrameProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPostShuffleFrameProcessor.java
index 3cbf8e7dd18..ac308adfaef 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPostShuffleFrameProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPostShuffleFrameProcessor.java
@@ -22,7 +22,6 @@ package org.apache.druid.msq.querykit.groupby;
import com.fasterxml.jackson.databind.ObjectMapper;
import it.unimi.dsi.fastutil.ints.IntSet;
import org.apache.druid.frame.Frame;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.key.ClusterBy;
@@ -145,7 +144,7 @@ public class GroupByPostShuffleFrameProcessor implements FrameProcessor<Object>
writeEmptyAggregationsFrameIfNeeded();
return ReturnOrAwait.returnObject(Unit.instance());
} else {
- final Frame frame = inputChannel.read();
+ final Frame frame = inputChannel.readFrame();
frameCursor = FrameProcessors.makeCursor(frame, frameReader);
final ColumnSelectorFactory frameColumnSelectorFactory = frameCursor.getColumnSelectorFactory();
@@ -264,7 +263,7 @@ public class GroupByPostShuffleFrameProcessor implements FrameProcessor<Object>
{
if (frameWriter != null && frameWriter.getNumRows() > 0) {
final Frame frame = Frame.wrap(frameWriter.toByteArray());
- outputChannel.write(new FrameWithPartition(frame, FrameWithPartition.NO_PARTITION));
+ outputChannel.write(frame);
frameWriter.close();
frameWriter = null;
outputRows += frame.numRows();
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleFrameProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleFrameProcessor.java
index 8a52d8704d8..03524b47d47 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleFrameProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleFrameProcessor.java
@@ -26,7 +26,6 @@ import org.apache.druid.collections.ResourceHolder;
import org.apache.druid.common.guava.FutureUtils;
import org.apache.druid.error.DruidException;
import org.apache.druid.frame.Frame;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.processor.FrameProcessor;
@@ -216,7 +215,7 @@ public class GroupByPreShuffleFrameProcessor extends BaseLeafFrameProcessor
closeAndDiscardResultYielder();
if (inputChannel.canRead()) {
- final Frame frame = inputChannel.read();
+ final Frame frame = inputChannel.readFrame();
final FrameSegment frameSegment = new FrameSegment(frame, inputFrameReader);
final Segment mappedSegment = mapUnmanagedSegment(frameSegment);
@@ -292,7 +291,7 @@ public class GroupByPreShuffleFrameProcessor extends BaseLeafFrameProcessor
{
if (frameWriter != null && frameWriter.getNumRows() > 0) {
final Frame frame = Frame.wrap(frameWriter.toByteArray());
- Iterables.getOnlyElement(outputChannels()).write(new FrameWithPartition(frame, FrameWithPartition.NO_PARTITION));
+ Iterables.getOnlyElement(outputChannels()).write(frame);
frameWriter.close();
frameWriter = null;
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/results/ExportResultsFrameProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/results/ExportResultsFrameProcessor.java
index 1d712012cbd..9c9a9028892 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/results/ExportResultsFrameProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/results/ExportResultsFrameProcessor.java
@@ -145,7 +145,7 @@ public class ExportResultsFrameProcessor implements FrameProcessor<Object>
exportWriter.writeResponseEnd();
return ReturnOrAwait.returnObject(exportFilePath);
} else {
- exportFrame(inputChannel.read());
+ exportFrame(inputChannel.readFrame());
return ReturnOrAwait.awaitAll(1);
}
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/results/QueryResultsFrameProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/results/QueryResultsFrameProcessor.java
index 7871af4bcfb..fc9240df913 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/results/QueryResultsFrameProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/results/QueryResultsFrameProcessor.java
@@ -20,13 +20,13 @@
package org.apache.druid.msq.querykit.results;
import it.unimi.dsi.fastutil.ints.IntSet;
-import org.apache.druid.frame.Frame;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.processor.FrameProcessor;
import org.apache.druid.frame.processor.FrameProcessors;
import org.apache.druid.frame.processor.ReturnOrAwait;
import org.apache.druid.java.util.common.Unit;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import java.io.IOException;
import java.util.Collections;
@@ -67,7 +67,7 @@ public class QueryResultsFrameProcessor implements FrameProcessor<Object>
if (inChannel.isFinished()) {
return ReturnOrAwait.returnObject(Unit.instance());
}
- writeFrame(inChannel.read());
+ writeRAC(inChannel.read());
return ReturnOrAwait.awaitAll(1);
}
@@ -77,8 +77,8 @@ public class QueryResultsFrameProcessor implements FrameProcessor<Object>
FrameProcessors.closeAll(inputChannels(), outputChannels());
}
- private void writeFrame(final Frame frame) throws IOException
+ private void writeRAC(final RowsAndColumns rac) throws IOException
{
- outChannel.write(frame);
+ outChannel.write(rac);
}
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java
index 2302c937944..ec24a9fc12c 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java
@@ -29,7 +29,6 @@ import org.apache.druid.collections.ResourceHolder;
import org.apache.druid.common.guava.FutureUtils;
import org.apache.druid.error.DruidException;
import org.apache.druid.frame.Frame;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.processor.FrameProcessor;
@@ -374,7 +373,7 @@ public class ScanQueryFrameProcessor extends BaseLeafFrameProcessor
{
if (cursor == null || cursor.isDone()) {
if (inputChannel.canRead()) {
- final Frame frame = inputChannel.read();
+ final Frame frame = inputChannel.readFrame();
final FrameSegment frameSegment = new FrameSegment(frame, inputFrameReader);
final Segment mappedSegment = mapUnmanagedSegment(frameSegment);
@@ -506,7 +505,7 @@ public class ScanQueryFrameProcessor extends BaseLeafFrameProcessor
{
if (frameWriter != null && frameWriter.getNumRows() > 0) {
final Frame frame = Frame.wrap(frameWriter.toByteArray());
- Iterables.getOnlyElement(outputChannels()).write(new FrameWithPartition(frame, FrameWithPartition.NO_PARTITION));
+ Iterables.getOnlyElement(outputChannels()).write(frame);
frameWriter.close();
frameWriter = null;
return frame.numRows();
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageInputChannelFactory.java b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageInputChannelFactory.java
index 601b4a4b3ad..7faa7c99c8e 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageInputChannelFactory.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageInputChannelFactory.java
@@ -32,8 +32,10 @@ import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.msq.exec.InputChannelFactory;
import org.apache.druid.msq.kernel.StageId;
import org.apache.druid.msq.shuffle.output.DurableStorageOutputChannelFactory;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.storage.StorageConnector;
+import javax.annotation.Nullable;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
@@ -52,15 +54,20 @@ public abstract class DurableStorageInputChannelFactory implements InputChannelF
private final ExecutorService remoteInputStreamPool;
private final String controllerTaskId;
+ @Nullable
+ private final WireTransferableContext wireTransferableContext;
+
public DurableStorageInputChannelFactory(
final String controllerTaskId,
final StorageConnector storageConnector,
- final ExecutorService remoteInputStreamPool
+ final ExecutorService remoteInputStreamPool,
+ @Nullable final WireTransferableContext wireTransferableContext
)
{
this.controllerTaskId = Preconditions.checkNotNull(controllerTaskId, "controllerTaskId");
this.storageConnector = Preconditions.checkNotNull(storageConnector, "storageConnector");
this.remoteInputStreamPool = Preconditions.checkNotNull(remoteInputStreamPool, "remoteInputStreamPool");
+ this.wireTransferableContext = wireTransferableContext;
}
/**
@@ -71,7 +78,8 @@ public abstract class DurableStorageInputChannelFactory implements InputChannelF
final String controllerTaskId,
final StorageConnector storageConnector,
final Closer closer,
- final boolean isQueryResults
+ final boolean isQueryResults,
+ @Nullable final WireTransferableContext wireTransferableContext
)
{
final String threadNameFormat =
@@ -84,10 +92,16 @@ public abstract class DurableStorageInputChannelFactory implements InputChannelF
return new DurableStorageQueryResultsInputChannelFactory(
controllerTaskId,
storageConnector,
- remoteInputStreamPool
+ remoteInputStreamPool,
+ wireTransferableContext
);
}
- return new DurableStorageStageInputChannelFactory(controllerTaskId, storageConnector, remoteInputStreamPool);
+ return new DurableStorageStageInputChannelFactory(
+ controllerTaskId,
+ storageConnector,
+ remoteInputStreamPool,
+ wireTransferableContext
+ );
}
@@ -124,7 +138,8 @@ public abstract class DurableStorageInputChannelFactory implements InputChannelF
inputStream,
remotePartitionPath,
remoteInputStreamPool,
- false
+ false,
+ wireTransferableContext
);
}
catch (Exception e) {
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageQueryResultsInputChannelFactory.java b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageQueryResultsInputChannelFactory.java
index 620b2271e25..ea7b4f80abe 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageQueryResultsInputChannelFactory.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageQueryResultsInputChannelFactory.java
@@ -20,8 +20,10 @@
package org.apache.druid.msq.shuffle.input;
import org.apache.druid.frame.util.DurableStorageUtils;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.storage.StorageConnector;
+import javax.annotation.Nullable;
import java.util.concurrent.ExecutorService;
/**
@@ -33,10 +35,11 @@ public class DurableStorageQueryResultsInputChannelFactory extends DurableStorag
public DurableStorageQueryResultsInputChannelFactory(
String controllerTaskId,
StorageConnector storageConnector,
- ExecutorService remoteInputStreamPool
+ ExecutorService remoteInputStreamPool,
+ @Nullable WireTransferableContext wireTransferableContext
)
{
- super(controllerTaskId, storageConnector, remoteInputStreamPool);
+ super(controllerTaskId, storageConnector, remoteInputStreamPool, wireTransferableContext);
}
@Override
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageStageInputChannelFactory.java b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageStageInputChannelFactory.java
index 1034b402679..982ecdff776 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageStageInputChannelFactory.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/DurableStorageStageInputChannelFactory.java
@@ -20,8 +20,10 @@
package org.apache.druid.msq.shuffle.input;
import org.apache.druid.frame.util.DurableStorageUtils;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.storage.StorageConnector;
+import javax.annotation.Nullable;
import java.util.concurrent.ExecutorService;
/**
@@ -33,10 +35,11 @@ public class DurableStorageStageInputChannelFactory extends DurableStorageInputC
public DurableStorageStageInputChannelFactory(
String controllerTaskId,
StorageConnector storageConnector,
- ExecutorService remoteInputStreamPool
+ ExecutorService remoteInputStreamPool,
+ @Nullable WireTransferableContext wireTransferableContext
)
{
- super(controllerTaskId, storageConnector, remoteInputStreamPool);
+ super(controllerTaskId, storageConnector, remoteInputStreamPool, wireTransferableContext);
}
@Override
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/WorkerInputChannelFactory.java b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/WorkerInputChannelFactory.java
index e55f63a1188..b3a5ed16e4d 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/WorkerInputChannelFactory.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/input/WorkerInputChannelFactory.java
@@ -30,7 +30,9 @@ import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.msq.exec.InputChannelFactory;
import org.apache.druid.msq.exec.WorkerClient;
import org.apache.druid.msq.kernel.StageId;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
+import javax.annotation.Nullable;
import java.util.List;
import java.util.function.Supplier;
@@ -42,10 +44,18 @@ public class WorkerInputChannelFactory implements InputChannelFactory
private final WorkerClient workerClient;
private final Supplier<List<String>> taskList;
- public WorkerInputChannelFactory(final WorkerClient workerClient, final Supplier<List<String>> taskList)
+ @Nullable
+ private final WireTransferableContext wtContext;
+
+ public WorkerInputChannelFactory(
+ final WorkerClient workerClient,
+ final Supplier<List<String>> taskList,
+ @Nullable final WireTransferableContext wtContext
+ )
{
this.workerClient = Preconditions.checkNotNull(workerClient, "workerClient");
this.taskList = Preconditions.checkNotNull(taskList, "taskList");
+ this.wtContext = wtContext;
}
@Override
@@ -53,7 +63,7 @@ public class WorkerInputChannelFactory implements InputChannelFactory
{
final String taskId = taskList.get().get(workerNumber);
final ReadableByteChunksFrameChannel channel =
- ReadableByteChunksFrameChannel.create(makeChannelId(taskId, stageId, partitionNumber), false);
+ ReadableByteChunksFrameChannel.create(makeChannelId(taskId, stageId, partitionNumber), false, wtContext);
fetch(taskId, stageId, partitionNumber, 0, channel);
return channel;
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/ChannelStageOutputReader.java b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/ChannelStageOutputReader.java
index ec95ca7af6a..aee5ec68a5d 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/ChannelStageOutputReader.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/ChannelStageOutputReader.java
@@ -33,7 +33,9 @@ import org.apache.druid.frame.file.FrameFileWriter;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.msq.exec.OutputChannelMode;
import org.apache.druid.msq.kernel.controller.ControllerQueryKernelUtils;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
+import javax.annotation.Nullable;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
@@ -92,10 +94,13 @@ public class ChannelStageOutputReader implements StageOutputReader
@GuardedBy("this")
private boolean didCloseWriter;
- public ChannelStageOutputReader(final ReadableFrameChannel channel)
+ public ChannelStageOutputReader(
+ final ReadableFrameChannel channel,
+ @Nullable final WireTransferableContext wtContext
+ )
{
this.channel = channel;
- this.writer = FrameFileWriter.open(new ChunkAcceptor(), null, ByteTracker.unboundedTracker());
+ this.writer = FrameFileWriter.open(new ChunkAcceptor(), null, ByteTracker.unboundedTracker(), wtContext);
}
/**
@@ -153,7 +158,7 @@ public class ChannelStageOutputReader implements StageOutputReader
continue;
} else if (channel.canRead()) {
try {
- writer.writeFrame(channel.read(), FrameFileWriter.NO_PARTITION);
+ writer.write(channel.read(), FrameFileWriter.NO_PARTITION);
}
catch (Exception e) {
try {
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageOutputChannelFactory.java b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageOutputChannelFactory.java
index 293f1758427..ba07a9b18ea 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageOutputChannelFactory.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageOutputChannelFactory.java
@@ -27,6 +27,7 @@ import org.apache.druid.frame.processor.OutputChannelFactory;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.concurrent.Execs;
import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.storage.StorageConnector;
import java.io.File;
@@ -48,6 +49,7 @@ public abstract class DurableStorageOutputChannelFactory implements OutputChanne
protected final StorageConnector storageConnector;
protected final File tmpDir;
protected final ExecutorService remoteInputStreamPool;
+ protected final WireTransferableContext wireTransferableContext;
public DurableStorageOutputChannelFactory(
final String controllerTaskId,
@@ -56,7 +58,8 @@ public abstract class DurableStorageOutputChannelFactory implements OutputChanne
final String taskId,
final int frameSize,
final StorageConnector storageConnector,
- final File tmpDir
+ final File tmpDir,
+ final WireTransferableContext wireTransferableContext
)
{
this.controllerTaskId = Preconditions.checkNotNull(controllerTaskId, "controllerTaskId");
@@ -68,6 +71,7 @@ public abstract class DurableStorageOutputChannelFactory implements OutputChanne
this.tmpDir = Preconditions.checkNotNull(tmpDir, "tmpDir is null");
this.remoteInputStreamPool =
Executors.newCachedThreadPool(Execs.makeThreadFactory("-remote-fetcher-%d"));
+ this.wireTransferableContext = wireTransferableContext;
}
/**
@@ -82,7 +86,8 @@ public abstract class DurableStorageOutputChannelFactory implements OutputChanne
final int frameSize,
final StorageConnector storageConnector,
final File tmpDir,
- final boolean isQueryResults
+ final boolean isQueryResults,
+ final WireTransferableContext wireTransferableContext
)
{
if (isQueryResults) {
@@ -93,7 +98,8 @@ public abstract class DurableStorageOutputChannelFactory implements OutputChanne
taskId,
frameSize,
storageConnector,
- tmpDir
+ tmpDir,
+ wireTransferableContext
);
} else {
return new DurableStorageTaskOutputChannelFactory(
@@ -103,7 +109,8 @@ public abstract class DurableStorageOutputChannelFactory implements OutputChanne
taskId,
frameSize,
storageConnector,
- tmpDir
+ tmpDir,
+ wireTransferableContext
);
}
}
@@ -139,8 +146,12 @@ public abstract class DurableStorageOutputChannelFactory implements OutputChanne
final String fileName = getFileNameWithPathForPartition(partitionNumber);
// As tasks dependent on output of this partition will forever block if no file is present in RemoteStorage. Hence, writing a dummy frame.
try {
- FrameFileWriter.open(Channels.newChannel(storageConnector.write(fileName)), null, ByteTracker.unboundedTracker())
- .close();
+ FrameFileWriter.open(
+ Channels.newChannel(storageConnector.write(fileName)),
+ null,
+ ByteTracker.unboundedTracker(),
+ wireTransferableContext
+ ).close();
return OutputChannel.nil(partitionNumber);
}
catch (IOException e) {
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageQueryResultsOutputChannelFactory.java b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageQueryResultsOutputChannelFactory.java
index 271129aa435..86824553751 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageQueryResultsOutputChannelFactory.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageQueryResultsOutputChannelFactory.java
@@ -28,6 +28,7 @@ import org.apache.druid.frame.processor.OutputChannel;
import org.apache.druid.frame.processor.PartitionedOutputChannel;
import org.apache.druid.frame.util.DurableStorageUtils;
import org.apache.druid.java.util.common.UOE;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.storage.StorageConnector;
import java.io.File;
@@ -47,10 +48,20 @@ public class DurableStorageQueryResultsOutputChannelFactory extends DurableStora
String taskId,
int frameSize,
StorageConnector storageConnector,
- File tmpDir
+ File tmpDir,
+ WireTransferableContext wireTransferableContext
)
{
- super(controllerTaskId, workerNumber, stageNumber, taskId, frameSize, storageConnector, tmpDir);
+ super(
+ controllerTaskId,
+ workerNumber,
+ stageNumber,
+ taskId,
+ frameSize,
+ storageConnector,
+ tmpDir,
+ wireTransferableContext
+ );
}
@Override
@@ -83,7 +94,8 @@ public class DurableStorageQueryResultsOutputChannelFactory extends DurableStora
FrameFileWriter.open(
Channels.newChannel(storageConnector.write(fileName)),
null,
- ByteTracker.unboundedTracker()
+ ByteTracker.unboundedTracker(),
+ wireTransferableContext
)
);
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageTaskOutputChannelFactory.java b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageTaskOutputChannelFactory.java
index 7ece0dd3155..174ba552bb8 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageTaskOutputChannelFactory.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/DurableStorageTaskOutputChannelFactory.java
@@ -39,6 +39,7 @@ import org.apache.druid.java.util.common.FileUtils;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.MappedByteBufferHandler;
import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.storage.StorageConnector;
import java.io.File;
@@ -65,10 +66,20 @@ public class DurableStorageTaskOutputChannelFactory
String taskId,
int frameSize,
StorageConnector storageConnector,
- File tmpDir
+ File tmpDir,
+ WireTransferableContext wireTransferableContext
)
{
- super(controllerTaskId, workerNumber, stageNumber, taskId, frameSize, storageConnector, tmpDir);
+ super(
+ controllerTaskId,
+ workerNumber,
+ stageNumber,
+ taskId,
+ frameSize,
+ storageConnector,
+ tmpDir,
+ wireTransferableContext
+ );
}
@Override
@@ -99,7 +110,8 @@ public class DurableStorageTaskOutputChannelFactory
FrameFileWriter.open(
Channels.newChannel(storageConnector.write(fileName)),
null,
- ByteTracker.unboundedTracker()
+ ByteTracker.unboundedTracker(),
+ wireTransferableContext
)
);
@@ -120,7 +132,8 @@ public class DurableStorageTaskOutputChannelFactory
storageConnector.read(fileName),
fileName,
remoteInputStreamPool,
- false
+ false,
+ wireTransferableContext
);
}
catch (IOException e) {
@@ -147,7 +160,8 @@ public class DurableStorageTaskOutputChannelFactory
FrameFileWriter.open(
Channels.newChannel(countingOutputStream),
ByteBuffer.allocate(Frame.compressionBufferSize(frameSize)),
- ByteTracker.unboundedTracker()
+ ByteTracker.unboundedTracker(),
+ wireTransferableContext
)
);
@@ -201,7 +215,8 @@ public class DurableStorageTaskOutputChannelFactory
frameFileFooterSupplier,
fileName,
remoteInputStreamPool,
- footerFile
+ footerFile,
+ wireTransferableContext
)
);
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/FileStageOutputReader.java b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/FileStageOutputReader.java
index 29fb7b17ee7..d67fc032b1f 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/FileStageOutputReader.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/FileStageOutputReader.java
@@ -25,6 +25,7 @@ import it.unimi.dsi.fastutil.bytes.ByteArrays;
import org.apache.druid.frame.channel.ReadableFileFrameChannel;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.file.FrameFile;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
import java.io.ByteArrayInputStream;
import java.io.IOException;
@@ -38,10 +39,15 @@ import java.nio.channels.Channels;
public class FileStageOutputReader implements StageOutputReader
{
private final FrameFile frameFile;
+ private final WireTransferable.ConcreteDeserializer deserializer;
- public FileStageOutputReader(FrameFile frameFile)
+ public FileStageOutputReader(
+ final FrameFile frameFile,
+ final WireTransferable.ConcreteDeserializer deserializer
+ )
{
this.frameFile = frameFile;
+ this.deserializer = deserializer;
}
/**
@@ -83,7 +89,7 @@ public class FileStageOutputReader implements StageOutputReader
@Override
public ReadableFrameChannel readLocally()
{
- return new ReadableFileFrameChannel(frameFile.newReference());
+ return new ReadableFileFrameChannel(frameFile.newReference(), deserializer);
}
/**
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/FutureReadableFrameChannel.java b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/FutureReadableFrameChannel.java
index 8dcb8786713..9ba3bca304f 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/FutureReadableFrameChannel.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/FutureReadableFrameChannel.java
@@ -21,9 +21,9 @@ package org.apache.druid.msq.shuffle.output;
import com.google.common.util.concurrent.ListenableFuture;
import org.apache.druid.common.guava.FutureUtils;
-import org.apache.druid.frame.Frame;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import java.util.NoSuchElementException;
@@ -64,7 +64,7 @@ public class FutureReadableFrameChannel implements ReadableFrameChannel
}
@Override
- public Frame read()
+ public RowsAndColumns read()
{
if (populateChannel()) {
return channel.read();
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java
index 3823f03b110..e4089b88599 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java
@@ -31,6 +31,7 @@ import org.apache.druid.java.util.common.ISE;
import org.apache.druid.msq.exec.Worker;
import org.apache.druid.msq.kernel.StageId;
import org.apache.druid.msq.rpc.WorkerResource;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.utils.CloseableUtils;
import javax.servlet.http.HttpServletRequest;
@@ -45,10 +46,10 @@ public class StageOutputHolder implements Closeable
private final SettableFuture<ReadableFrameChannel> channelFuture;
private final ListenableFuture<StageOutputReader> readerFuture;
- public StageOutputHolder()
+ public StageOutputHolder(final WireTransferableContext wireTransferableContext)
{
this.channelFuture = SettableFuture.create();
- this.readerFuture = FutureUtils.transform(channelFuture, StageOutputHolder::createReader);
+ this.readerFuture = FutureUtils.transform(channelFuture, channel -> createReader(channel, wireTransferableContext));
}
/**
@@ -112,7 +113,10 @@ public class StageOutputHolder implements Closeable
}
}
- private static StageOutputReader createReader(final ReadableFrameChannel channel)
+ private static StageOutputReader createReader(
+ final ReadableFrameChannel channel,
+ final WireTransferableContext wireTransferableContext
+ )
{
if (channel == null) {
// Happens if close() was called before the channel resolved.
@@ -133,11 +137,11 @@ public class StageOutputHolder implements Closeable
// Close original channel, so we don't leak a frame file reference.
channel.close();
- return new FileStageOutputReader(frameFile);
+ return new FileStageOutputReader(frameFile, wireTransferableContext.concreteDeserializer());
}
}
// Generic implementation for any other type of channel.
- return new ChannelStageOutputReader(channel);
+ return new ChannelStageOutputReader(channel, wireTransferableContext);
}
}
diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java b/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java
index 28e5f257f09..d3363b25498 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java
@@ -37,6 +37,7 @@ import org.apache.druid.error.Forbidden;
import org.apache.druid.error.InvalidInput;
import org.apache.druid.error.NotFound;
import org.apache.druid.error.QueryExceptionCompat;
+import org.apache.druid.frame.Frame;
import org.apache.druid.frame.channel.FrameChannelSequence;
import org.apache.druid.indexer.TaskStatusPlus;
import org.apache.druid.java.util.common.ISE;
@@ -73,6 +74,7 @@ import org.apache.druid.query.ExecutionMode;
import org.apache.druid.query.QueryContext;
import org.apache.druid.query.QueryContexts;
import org.apache.druid.query.QueryException;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.rpc.HttpResponseException;
import org.apache.druid.rpc.indexing.OverlordClient;
import org.apache.druid.server.QueryResponse;
@@ -136,6 +138,7 @@ public class SqlStatementResource
private final AuthorizerMapper authorizerMapper;
private final DefaultQueryConfig defaultQueryConfig;
private final ServerConfig serverConfig;
+ private final WireTransferableContext wireTransferableContext;
@Inject
public SqlStatementResource(
@@ -145,7 +148,8 @@ public class SqlStatementResource
final @MultiStageQuery StorageConnectorProvider storageConnectorProvider,
final AuthorizerMapper authorizerMapper,
final DefaultQueryConfig defaultQueryConfig,
- final ServerConfig serverConfig
+ final ServerConfig serverConfig,
+ final WireTransferableContext wireTransferableContext
)
{
this.msqSqlStatementFactory = msqSqlStatementFactory;
@@ -155,6 +159,7 @@ public class SqlStatementResource
this.authorizerMapper = authorizerMapper;
this.defaultQueryConfig = defaultQueryConfig;
this.serverConfig = serverConfig;
+ this.wireTransferableContext = wireTransferableContext;
}
/**
@@ -813,7 +818,8 @@ public class SqlStatementResource
msqControllerTask.getId(),
storageConnector,
closer,
- true
+ true,
+ wireTransferableContext
);
results = Optional.of(Yielders.each(
Sequences.concat(pages.stream()
@@ -838,9 +844,9 @@ public class SqlStatementResource
}
})
.collect(Collectors.toList()))
- .flatMap(frame ->
+ .flatMap(rac ->
SqlStatementResourceHelper.getResultSequence(
- frame,
+ rac.as(Frame.class),
finalStage.getFrameReader(),
msqControllerTask.getQuerySpec().getColumnMappings(),
new ResultsContext(
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/BroadcastJoinSegmentMapFnProcessorTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/BroadcastJoinSegmentMapFnProcessorTest.java
index 85365c1fcca..3eebc2d92d3 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/BroadcastJoinSegmentMapFnProcessorTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/BroadcastJoinSegmentMapFnProcessorTest.java
@@ -26,7 +26,6 @@ import it.unimi.dsi.fastutil.ints.Int2IntMap;
import it.unimi.dsi.fastutil.ints.Int2IntOpenHashMap;
import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
import it.unimi.dsi.fastutil.ints.IntSet;
-import org.apache.druid.frame.Frame;
import org.apache.druid.frame.FrameType;
import org.apache.druid.frame.allocation.ArenaMemoryAllocator;
import org.apache.druid.frame.channel.ByteTracker;
@@ -47,6 +46,7 @@ import org.apache.druid.query.JoinDataSource;
import org.apache.druid.query.Query;
import org.apache.druid.query.QueryContext;
import org.apache.druid.query.policy.NoopPolicyEnforcer;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import org.apache.druid.segment.CursorFactory;
import org.apache.druid.segment.QueryableIndexCursorFactory;
import org.apache.druid.segment.TestIndex;
@@ -120,8 +120,8 @@ public class BroadcastJoinSegmentMapFnProcessorTest extends InitializedNullHandl
final List<ReadableFrameChannel> channels = new ArrayList<>();
channels.add(new ExplodingReadableFrameChannel());
- channels.add(new ReadableFileFrameChannel(FrameFile.open(testDataFile1, null)));
- channels.add(new ReadableFileFrameChannel(FrameFile.open(testDataFile2, null)));
+ channels.add(new ReadableFileFrameChannel(FrameFile.open(testDataFile1, null), null));
+ channels.add(new ReadableFileFrameChannel(FrameFile.open(testDataFile2, null), null));
final List<FrameReader> channelReaders = new ArrayList<>();
channelReaders.add(null);
@@ -215,7 +215,7 @@ public class BroadcastJoinSegmentMapFnProcessorTest extends InitializedNullHandl
sideStageChannelNumberMap.put(0, 0);
final List<ReadableFrameChannel> channels = new ArrayList<>();
- channels.add(new ReadableFileFrameChannel(FrameFile.open(testDataFile1, ByteTracker.unboundedTracker())));
+ channels.add(new ReadableFileFrameChannel(FrameFile.open(testDataFile1, ByteTracker.unboundedTracker()), null));
final List<FrameReader> channelReaders = new ArrayList<>();
channelReaders.add(frameReader1);
@@ -256,7 +256,7 @@ public class BroadcastJoinSegmentMapFnProcessorTest extends InitializedNullHandl
sideStageChannelNumberMap.put(0, 0);
final List<ReadableFrameChannel> channels = new ArrayList<>();
- channels.add(new ReadableFileFrameChannel(FrameFile.open(testDataFile1, ByteTracker.unboundedTracker())));
+ channels.add(new ReadableFileFrameChannel(FrameFile.open(testDataFile1, ByteTracker.unboundedTracker()), null));
final List<FrameReader> channelReaders = new ArrayList<>();
channelReaders.add(frameReader1);
@@ -316,7 +316,7 @@ public class BroadcastJoinSegmentMapFnProcessorTest extends InitializedNullHandl
}
@Override
- public Frame read()
+ public RowsAndColumns read()
{
throw new IllegalStateException();
}
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/ChainedProcessorManagerTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/ChainedProcessorManagerTest.java
index 4fadb7700f7..ddadbf37897 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/ChainedProcessorManagerTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/ChainedProcessorManagerTest.java
@@ -129,7 +129,7 @@ public class ChainedProcessorManagerTest extends FrameProcessorExecutorTest.Base
final HashSet<Long> actualValues = new HashSet<>();
try (ReadableFrameChannel readable = outputChannel.readable()) {
while (readable.canRead()) {
- actualValues.add(extractColumnValue(readable.read(), 1));
+ actualValues.add(extractColumnValue(readable.readFrame(), 1));
}
}
Assert.assertEquals(new HashSet<>(expectedValues), actualValues);
@@ -168,7 +168,7 @@ public class ChainedProcessorManagerTest extends FrameProcessorExecutorTest.Base
final HashSet<Long> actualValues = new HashSet<>();
try (ReadableFrameChannel readable = outputChannel.readable()) {
while (readable.canRead()) {
- actualValues.add(extractColumnValue(readable.read(), 1));
+ actualValues.add(extractColumnValue(readable.readFrame(), 1));
}
}
Assert.assertEquals(expectedValues, actualValues);
@@ -205,7 +205,7 @@ public class ChainedProcessorManagerTest extends FrameProcessorExecutorTest.Base
final HashSet<Long> actualValues = new HashSet<>();
try (ReadableFrameChannel readable = outputChannel.readable()) {
while (readable.canRead()) {
- actualValues.add(extractColumnValue(readable.read(), 1));
+ actualValues.add(extractColumnValue(readable.readFrame(), 1));
}
}
Assert.assertEquals(expectedValues, actualValues);
@@ -252,7 +252,7 @@ public class ChainedProcessorManagerTest extends FrameProcessorExecutorTest.Base
final HashSet<Long> actualValues = new HashSet<>();
try (ReadableFrameChannel readable = outputChannel.readable()) {
while (readable.canRead()) {
- actualValues.add(extractColumnValue(readable.read(), 1));
+ actualValues.add(extractColumnValue(readable.readFrame(), 1));
}
}
Assert.assertEquals(expectedValues, actualValues);
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/NonFailingWritableFrameChannel.java b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/NonFailingWritableFrameChannel.java
index 07b1dc2af7a..9a32513303b 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/NonFailingWritableFrameChannel.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/querykit/NonFailingWritableFrameChannel.java
@@ -20,8 +20,8 @@
package org.apache.druid.msq.querykit;
import com.google.common.util.concurrent.ListenableFuture;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.WritableFrameChannel;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import javax.annotation.Nullable;
import java.io.IOException;
@@ -36,9 +36,9 @@ public class NonFailingWritableFrameChannel implements WritableFrameChannel
}
@Override
- public void write(FrameWithPartition frameWithPartition) throws IOException
+ public void write(RowsAndColumns rac, int partitionNumber) throws IOException
{
- delegate.write(frameWithPartition);
+ delegate.write(rac, partitionNumber);
}
@Override
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/rpc/BaseWorkerClientImplTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/rpc/BaseWorkerClientImplTest.java
index 705369a2b4c..79cc1495c48 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/rpc/BaseWorkerClientImplTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/rpc/BaseWorkerClientImplTest.java
@@ -154,7 +154,7 @@ public class BaseWorkerClientImplTest extends InitializedNullHandlingTest
// Perform the test.
final StageId stageId = new StageId("xyz", 1);
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("testChannel", false);
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("testChannel", false, null);
final Future<List<List<Object>>> framesFuture = readChannelAsync(channel);
Assert.assertFalse(workerClient.fetchChannelData(WORKER_ID, stageId, 2, 0, channel).get());
@@ -180,7 +180,7 @@ public class BaseWorkerClientImplTest extends InitializedNullHandlingTest
// Perform the test.
final StageId stageId = new StageId("xyz", 1);
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("testChannel", false);
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("testChannel", false, null);
channel.close(); // ReadableFrameChannel's close() method.
final ExecutionException e = Assert.assertThrows(
@@ -219,7 +219,7 @@ public class BaseWorkerClientImplTest extends InitializedNullHandlingTest
// Perform the test.
final StageId stageId = new StageId("xyz", 1);
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("testChannel", false);
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("testChannel", false, null);
final Future<List<List<Object>>> framesFuture = readChannelAsync(channel);
Assert.assertFalse(workerClient.fetchChannelData(WORKER_ID, stageId, 2, 0, channel).get());
@@ -244,7 +244,7 @@ public class BaseWorkerClientImplTest extends InitializedNullHandlingTest
// Perform the test.
final StageId stageId = new StageId("xyz", 1);
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("testChannel", false);
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("testChannel", false, null);
final ExecutionException e = Assert.assertThrows(
ExecutionException.class,
@@ -281,7 +281,7 @@ public class BaseWorkerClientImplTest extends InitializedNullHandlingTest
// Perform the test.
final StageId stageId = new StageId("xyz", 1);
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("testChannel", false);
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("testChannel", false, null);
final Future<List<List<Object>>> framesFuture = readChannelAsync(channel);
Assert.assertFalse(workerClient.fetchChannelData(WORKER_ID, stageId, 2, 0, channel).get());
@@ -302,7 +302,7 @@ public class BaseWorkerClientImplTest extends InitializedNullHandlingTest
FutureUtils.getUnchecked(channel.readabilityFuture(), false);
if (channel.canRead()) {
- final Frame frame = channel.read();
+ final Frame frame = channel.readFrame();
retVal.addAll(FrameTestUtil.readRowsFromCursorFactory(FRAME_READER.makeCursorFactory(frame)).toList());
}
}
@@ -318,10 +318,10 @@ public class BaseWorkerClientImplTest extends InitializedNullHandlingTest
{
final ByteArrayOutputStream baos = new ByteArrayOutputStream();
final FrameFileWriter writer =
- FrameFileWriter.open(Channels.newChannel(baos), null, ByteTracker.unboundedTracker());
+ FrameFileWriter.open(Channels.newChannel(baos), null, ByteTracker.unboundedTracker(), FrameTestUtil.WT_CONTEXT_LEGACY);
frames.forEach(frame -> {
try {
- writer.writeFrame(frame, FrameFileWriter.NO_PARTITION);
+ writer.write(frame.asRAC(), FrameFileWriter.NO_PARTITION);
}
catch (IOException e) {
throw new RuntimeException(e);
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/shuffle/DurableStorageOutputChannelFactoryTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/shuffle/DurableStorageOutputChannelFactoryTest.java
index 38d99dfe03d..645de1692ae 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/shuffle/DurableStorageOutputChannelFactoryTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/shuffle/DurableStorageOutputChannelFactoryTest.java
@@ -20,6 +20,7 @@
package org.apache.druid.msq.shuffle;
import org.apache.druid.frame.processor.OutputChannelFactoryTest;
+import org.apache.druid.frame.testutil.FrameTestUtil;
import org.apache.druid.msq.shuffle.output.DurableStorageTaskOutputChannelFactory;
import org.apache.druid.storage.local.LocalFileStorageConnector;
import org.junit.ClassRule;
@@ -43,7 +44,8 @@ public class DurableStorageOutputChannelFactoryTest extends OutputChannelFactory
"0",
100,
new LocalFileStorageConnector(folder.newFolder()),
- folder.newFolder()
+ folder.newFolder(),
+ FrameTestUtil.WT_CONTEXT_LEGACY
),
100
);
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/shuffle/output/ChannelStageOutputReaderTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/shuffle/output/ChannelStageOutputReaderTest.java
index 789f4f67f9b..0454e38b281 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/shuffle/output/ChannelStageOutputReaderTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/shuffle/output/ChannelStageOutputReaderTest.java
@@ -91,7 +91,7 @@ public class ChannelStageOutputReaderTest extends InitializedNullHandlingTest
public void setUp() throws Exception
{
channel = BlockingQueueFrameChannel.minimal();
- channelReader = new ChannelStageOutputReader(channel.readable());
+ channelReader = new ChannelStageOutputReader(channel.readable(), FrameTestUtil.WT_CONTEXT_LEGACY);
tmpFile = temporaryFolder.newFile();
tmpOut = Files.newOutputStream(tmpFile.toPath());
}
@@ -127,7 +127,7 @@ public class ChannelStageOutputReaderTest extends InitializedNullHandlingTest
}
Assert.assertEquals(1, tmpFrameFile.numFrames());
- Assert.assertEquals(frame.numBytes(), tmpFrameFile.frame(0).numBytes());
+ Assert.assertEquals(frame.numBytes(), tmpFrameFile.rac(0, null).as(Frame.class).numBytes());
}
@Test
@@ -144,7 +144,7 @@ public class ChannelStageOutputReaderTest extends InitializedNullHandlingTest
}
Assert.assertEquals(1, tmpFrameFile.numFrames());
- Assert.assertEquals(frame.numBytes(), tmpFrameFile.frame(0).numBytes());
+ Assert.assertEquals(frame.numBytes(), tmpFrameFile.rac(0, null).as(Frame.class).numBytes());
}
@Test
@@ -159,7 +159,7 @@ public class ChannelStageOutputReaderTest extends InitializedNullHandlingTest
}
Assert.assertEquals(1, tmpFrameFile.numFrames());
- Assert.assertEquals(frame.numBytes(), tmpFrameFile.frame(0).numBytes());
+ Assert.assertEquals(frame.numBytes(), tmpFrameFile.rac(0, null).as(Frame.class).numBytes());
}
@Test
@@ -219,9 +219,9 @@ public class ChannelStageOutputReaderTest extends InitializedNullHandlingTest
}
Assert.assertEquals(3, tmpFrameFile.numFrames());
- Assert.assertEquals(frame.numBytes(), tmpFrameFile.frame(0).numBytes());
- Assert.assertEquals(frame.numBytes(), tmpFrameFile.frame(1).numBytes());
- Assert.assertEquals(frame.numBytes(), tmpFrameFile.frame(2).numBytes());
+ Assert.assertEquals(frame.numBytes(), tmpFrameFile.rac(0, null).as(Frame.class).numBytes());
+ Assert.assertEquals(frame.numBytes(), tmpFrameFile.rac(1, null).as(Frame.class).numBytes());
+ Assert.assertEquals(frame.numBytes(), tmpFrameFile.rac(2, null).as(Frame.class).numBytes());
}
/**
@@ -280,7 +280,7 @@ public class ChannelStageOutputReaderTest extends InitializedNullHandlingTest
private static final int EXPECTED_NUM_ROWS = 1209;
private final BlockingQueueFrameChannel channel = new BlockingQueueFrameChannel(MAX_FRAMES);
- private final ChannelStageOutputReader reader = new ChannelStageOutputReader(channel.readable());
+ private final ChannelStageOutputReader reader = new ChannelStageOutputReader(channel.readable(), FrameTestUtil.WT_CONTEXT_LEGACY);
@Rule
public TemporaryFolder temporaryFolder = new TemporaryFolder();
@@ -387,7 +387,7 @@ public class ChannelStageOutputReaderTest extends InitializedNullHandlingTest
final FrameFile frameFile = FrameFile.open(tmpFile, null);
final int numRows =
- FrameTestUtil.readRowsFromFrameChannel(new ReadableFileFrameChannel(frameFile), frameReader).toList().size();
+ FrameTestUtil.readRowsFromFrameChannel(new ReadableFileFrameChannel(frameFile, null), frameReader).toList().size();
Assert.assertEquals(EXPECTED_NUM_ROWS, numRows);
}
@@ -425,7 +425,7 @@ public class ChannelStageOutputReaderTest extends InitializedNullHandlingTest
final FrameFile frameFile = FrameFile.open(tmpFile, null);
final int numRows =
- FrameTestUtil.readRowsFromFrameChannel(new ReadableFileFrameChannel(frameFile), frameReader).toList().size();
+ FrameTestUtil.readRowsFromFrameChannel(new ReadableFileFrameChannel(frameFile, null), frameReader).toList().size();
Assert.assertEquals(EXPECTED_NUM_ROWS, numRows);
}
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java
index 64c481e75bc..be911d59638 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/sql/MSQTaskQueryMakerTest.java
@@ -31,6 +31,7 @@ import com.google.inject.testing.fieldbinder.BoundFieldModule;
import com.google.inject.util.Modules;
import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
import org.apache.druid.client.ImmutableSegmentLoadInfo;
+import org.apache.druid.frame.testutil.FrameTestUtil;
import org.apache.druid.guice.ConfigModule;
import org.apache.druid.guice.DruidGuiceExtensions;
import org.apache.druid.guice.DruidSecondaryModule;
@@ -80,6 +81,7 @@ import org.apache.druid.query.policy.NoopPolicyEnforcer;
import org.apache.druid.query.policy.PolicyEnforcer;
import org.apache.druid.query.policy.RestrictAllTablesPolicyEnforcer;
import org.apache.druid.query.policy.RowFilterPolicy;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.query.scan.ScanQuery;
import org.apache.druid.rpc.indexing.OverlordClient;
import org.apache.druid.segment.IndexIO;
@@ -115,6 +117,7 @@ import org.mockito.junit.MockitoRule;
import javax.annotation.Nullable;
import java.util.ArrayList;
+import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
@@ -213,10 +216,11 @@ public class MSQTaskQueryMakerTest
new ConfigModule(),
new SegmentWranglerModule(),
new LookylooModule(),
- new MSQIndexingModule()
+ new MSQIndexingModule(),
+ binder -> binder.bind(WireTransferableContext.class).toInstance(FrameTestUtil.WT_CONTEXT_LEGACY)
);
Injector injector = Guice.createInjector(defaultModule, BoundFieldModule.of(this));
- DruidSecondaryModule.setupJackson(injector, objectMapper);
+ DruidSecondaryModule.setupJackson(injector, objectMapper, Collections.emptyMap(), true);
// Populate loadedSegmentMetadata from walker segments so CoordinatorClient.fetchSegment() can find them
List<ImmutableSegmentLoadInfo> loadedSegmentMetadata = new ArrayList<>();
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlMSQStatementResourcePostTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlMSQStatementResourcePostTest.java
index 73502be0b81..bfc6fd3ef8a 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlMSQStatementResourcePostTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlMSQStatementResourcePostTest.java
@@ -79,7 +79,8 @@ public class SqlMSQStatementResourcePostTest extends MSQTestBase
s -> localFileStorageConnector,
authorizerMapper,
new DefaultQueryConfig(Map.of("debug", "false")),
- new ServerConfig()
+ new ServerConfig(),
+ null
);
}
@@ -335,7 +336,8 @@ public class SqlMSQStatementResourcePostTest extends MSQTestBase
s -> NilStorageConnector.getInstance(),
authorizerMapper,
DefaultQueryConfig.NIL,
- new ServerConfig()
+ new ServerConfig(),
+ null
);
String errorMessage = "The sql statement api cannot read from the select destination [durableStorage] provided in "
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlStatementResourceTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlStatementResourceTest.java
index d71bc4501c4..8f757afe0ca 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlStatementResourceTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlStatementResourceTest.java
@@ -702,7 +702,8 @@ public class SqlStatementResourceTest extends MSQTestBase
tempDir -> localFileStorageConnector,
authorizerMapper,
new DefaultQueryConfig(Map.of("debug", "true")),
- new ServerConfig()
+ new ServerConfig(),
+ null
);
}
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
index cda596cbc30..5d70be1737b 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java
@@ -43,6 +43,7 @@ import org.apache.druid.data.input.impl.DimensionsSpec;
import org.apache.druid.data.input.impl.LongDimensionSchema;
import org.apache.druid.data.input.impl.StringDimensionSchema;
import org.apache.druid.discovery.NodeRole;
+import org.apache.druid.frame.Frame;
import org.apache.druid.frame.FrameType;
import org.apache.druid.frame.channel.FrameChannelSequence;
import org.apache.druid.frame.processor.Bouncer;
@@ -809,7 +810,7 @@ public class MSQTestBase extends BaseCalciteQueryTest
"compaction"
)
).registerSubtypes(ExternalDataSource.class));
- DruidSecondaryModule.setupJackson(injector, mapper);
+ DruidSecondaryModule.setupJackson(injector, mapper, Collections.emptyMap(), true);
mapper.registerSubtypes(new NamedType(LocalLoadSpec.class, "local"));
@@ -1650,14 +1651,15 @@ public class MSQTestBase extends BaseCalciteQueryTest
controllerId,
localFileStorageConnector,
closer,
- true
+ true,
+ null
);
rows.addAll(new FrameChannelSequence(inputChannelFactory.openChannel(
finalStage.getId(),
pageInformation.getWorker() == null ? 0 : pageInformation.getWorker(),
pageInformation.getPartition() == null ? 0 : pageInformation.getPartition()
- )).flatMap(frame -> SqlStatementResourceHelper.getResultSequence(
- frame,
+ )).flatMap(rac -> SqlStatementResourceHelper.getResultSequence(
+ rac.as(Frame.class),
finalStage.getFrameReader(),
msqControllerTask.getQuerySpec().getColumnMappings(),
new ResultsContext(msqControllerTask.getSqlTypeNames(), msqControllerTask.getSqlResultsContext()),
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java
index b5bdbf5e5ef..f1d9d0f72bc 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java
@@ -45,6 +45,7 @@ import org.apache.druid.msq.kernel.WorkOrder;
import org.apache.druid.msq.util.MultiStageQueryContext;
import org.apache.druid.query.groupby.GroupingEngine;
import org.apache.druid.query.policy.PolicyEnforcer;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.segment.IndexIO;
import org.apache.druid.segment.IndexMerger;
import org.apache.druid.segment.IndexMergerV9;
@@ -272,6 +273,12 @@ public class MSQTestWorkerContext implements WorkerContext
return mapper;
}
+ @Override
+ public WireTransferableContext wireTransferableContext()
+ {
+ return injector.getInstance(WireTransferableContext.class);
+ }
+
@Override
public IndexIO indexIO()
{
diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/util/SqlStatementResourceHelperTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/util/SqlStatementResourceHelperTest.java
index b467ea02cf9..9a16f831230 100644
--- a/multi-stage-query/src/test/java/org/apache/druid/msq/util/SqlStatementResourceHelperTest.java
+++ b/multi-stage-query/src/test/java/org/apache/druid/msq/util/SqlStatementResourceHelperTest.java
@@ -36,6 +36,7 @@ import org.apache.druid.msq.indexing.report.MSQStatusReport;
import org.apache.druid.msq.indexing.report.MSQTaskReportPayload;
import org.apache.druid.msq.indexing.report.MSQTaskReportTest;
import org.apache.druid.msq.sql.entity.PageInformation;
+import org.apache.druid.query.rowsandcols.concrete.FrameRowsAndColumns;
import org.easymock.EasyMock;
import org.junit.Assert;
import org.junit.Test;
@@ -448,7 +449,7 @@ public class SqlStatementResourceHelperTest
if (prev > current) {
throw new IllegalArgumentException("Channel numbers should be in increasing order");
}
- channelCounters.addFrame(current, createFrame(current * 10 + 1, 100L));
+ channelCounters.addRAC(createFrame(current * 10 + 1, 100L).asRAC(), current);
prev = current;
}
return channelCounters;
@@ -460,6 +461,14 @@ public class SqlStatementResourceHelperTest
Frame frame = EasyMock.mock(Frame.class);
EasyMock.expect(frame.numRows()).andReturn(numRows).anyTimes();
EasyMock.expect(frame.numBytes()).andReturn(numBytes).anyTimes();
+
+ // Mock asRAC() to return a mock FrameRowsAndColumns
+ FrameRowsAndColumns rac = EasyMock.mock(FrameRowsAndColumns.class);
+ EasyMock.expect(rac.numRows()).andReturn(numRows).anyTimes();
+ EasyMock.expect(rac.as(Frame.class)).andReturn(frame).anyTimes();
+ EasyMock.replay(rac);
+
+ EasyMock.expect(frame.asRAC()).andReturn(rac).anyTimes();
EasyMock.replay(frame);
return frame;
}
diff --git a/processing/src/main/java/org/apache/druid/frame/Frame.java b/processing/src/main/java/org/apache/druid/frame/Frame.java
index 7a2366b9769..64cd496bca9 100644
--- a/processing/src/main/java/org/apache/druid/frame/Frame.java
+++ b/processing/src/main/java/org/apache/druid/frame/Frame.java
@@ -19,23 +19,23 @@
package org.apache.druid.frame;
-import com.google.common.primitives.Ints;
-import net.jpountz.lz4.LZ4Compressor;
-import net.jpountz.lz4.LZ4Factory;
-import net.jpountz.lz4.LZ4SafeDecompressor;
import org.apache.datasketches.memory.Memory;
import org.apache.datasketches.memory.WritableMemory;
-import org.apache.druid.frame.channel.ByteTracker;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.frame.wire.FrameCompression;
import org.apache.druid.io.Channels;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.ISE;
-import org.apache.druid.segment.data.CompressionStrategy;
+import org.apache.druid.query.rowsandcols.concrete.ColumnBasedFrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.FrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowsAndColumns;
import javax.annotation.Nullable;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.channels.WritableByteChannel;
+import java.util.Objects;
/**
* A data frame.
@@ -66,17 +66,7 @@ import java.nio.channels.WritableByteChannel;
* as little-endian longs.
* - NNN bytes: regions, back-to-back.
*
- * There is also a compressed frame format. Compressed frames are written by {@link #writeTo} when "compress" is
- * true, and decompressed by {@link #decompress}. Format:
- *
- * - 1 byte: compression type: {@link CompressionStrategy#getId()}. Currently, only LZ4 is supported.
- * - 8 bytes: compressed frame length, little-endian long
- * - 8 bytes: uncompressed frame length (numBytes), little-endian long
- * - NNN bytes: LZ4-compressed frame
- * - 8 bytes: 64-bit xxhash checksum of prior content, including 16-byte header and compressed frame, little-endian long
- *
- * Note to developers: if we end up needing to add more fields here, consider introducing a Smile (or Protobuf, etc)
- * header to make it simpler to add more fields.
+ * There is also a compressed frame format. See {@link FrameCompression}.
*/
public class Frame
{
@@ -84,20 +74,9 @@ public class Frame
Byte.BYTES /* version */ +
Long.BYTES /* total size */ +
Integer.BYTES /* number of rows */ +
- Integer.BYTES /* number of columns */ +
+ Integer.BYTES /* number of regions */ +
Byte.BYTES /* permuted flag */;
- // Compression type, compressed length, uncompressed length
- public static final int COMPRESSED_FRAME_HEADER_SIZE = Byte.BYTES + Long.BYTES * 2;
- public static final int COMPRESSED_FRAME_TRAILER_SIZE = Long.BYTES; // Checksum
- public static final int COMPRESSED_FRAME_ENVELOPE_SIZE = COMPRESSED_FRAME_HEADER_SIZE
- + COMPRESSED_FRAME_TRAILER_SIZE;
-
- private static final LZ4Compressor LZ4_COMPRESSOR = LZ4Factory.fastestInstance().fastCompressor();
- private static final LZ4SafeDecompressor LZ4_DECOMPRESSOR = LZ4Factory.fastestInstance().safeDecompressor();
-
- private static final int CHECKSUM_SEED = 0;
-
private final Memory memory;
private final FrameType frameType;
private final long numBytes;
@@ -180,73 +159,7 @@ public class Frame
*/
public static Frame decompress(final Memory memory, final long position, final long length)
{
- if (memory.getCapacity() < position + length) {
- throw new ISE("Provided position, length is out of bounds");
- }
-
- if (length < COMPRESSED_FRAME_ENVELOPE_SIZE) {
- throw new ISE("Region too short");
- }
-
- // Verify checksum.
- final long expectedChecksum = memory.getLong(position + length - COMPRESSED_FRAME_TRAILER_SIZE);
- final long actualChecksum = memory.xxHash64(position, length - COMPRESSED_FRAME_TRAILER_SIZE, CHECKSUM_SEED);
-
- if (expectedChecksum != actualChecksum) {
- throw new ISE("Checksum mismatch");
- }
-
- final byte compressionTypeId = memory.getByte(position);
- final CompressionStrategy compressionStrategy = CompressionStrategy.forId(compressionTypeId);
-
- if (compressionStrategy != CompressionStrategy.LZ4) {
- throw new ISE("Unsupported compression strategy [%s]", compressionStrategy);
- }
-
- final int compressedFrameLength = Ints.checkedCast(memory.getLong(position + Byte.BYTES));
- final int uncompressedFrameLength = Ints.checkedCast(memory.getLong(position + Byte.BYTES + Long.BYTES));
- final int compressedFrameLengthFromRegionLength = Ints.checkedCast(length - COMPRESSED_FRAME_ENVELOPE_SIZE);
- final long frameStart = position + COMPRESSED_FRAME_HEADER_SIZE;
-
- // Verify length.
- if (compressedFrameLength != compressedFrameLengthFromRegionLength) {
- throw new ISE(
- "Compressed sizes disagree: [%d] (embedded) vs [%d] (region length)",
- compressedFrameLength,
- compressedFrameLengthFromRegionLength
- );
- }
-
- if (memory.hasByteBuffer()) {
- // Decompress directly out of the ByteBuffer.
- final ByteBuffer srcBuffer = memory.getByteBuffer();
- final ByteBuffer dstBuffer = ByteBuffer.allocate(uncompressedFrameLength);
- final int numBytesDecompressed =
- LZ4_DECOMPRESSOR.decompress(
- srcBuffer,
- Ints.checkedCast(memory.getRegionOffset() + frameStart),
- compressedFrameLength,
- dstBuffer,
- 0,
- uncompressedFrameLength
- );
-
- // Sanity check.
- if (numBytesDecompressed != uncompressedFrameLength) {
- throw new ISE(
- "Expected to decompress [%d] bytes but got [%d] bytes",
- uncompressedFrameLength,
- numBytesDecompressed
- );
- }
-
- return Frame.wrap(dstBuffer);
- } else {
- // Copy first, then decompress.
- final byte[] compressedFrame = new byte[compressedFrameLength];
- memory.getByteArray(frameStart, compressedFrame, 0, compressedFrameLength);
- return Frame.wrap(LZ4_DECOMPRESSOR.decompress(compressedFrame, uncompressedFrameLength));
- }
+ return Frame.wrap(FrameCompression.decompress(memory, position, length));
}
/**
@@ -254,7 +167,7 @@ public class Frame
*/
public static int compressionBufferSize(final long frameBytes)
{
- return COMPRESSED_FRAME_ENVELOPE_SIZE + LZ4_COMPRESSOR.maxCompressedLength(Ints.checkedCast(frameBytes));
+ return FrameCompression.compressionBufferSize(frameBytes);
}
public FrameType type()
@@ -331,6 +244,17 @@ public class Frame
return memory.region(regionStartPosition, regionEndPosition - regionStartPosition);
}
+ /**
+ * Direct, readable access to this frame's memory.
+ * Most callers should use {@link #region} and {@link #physicalRow}, rather than this direct-access method.
+ *
+ * @throws IllegalStateException if this frame wraps non-writable memory
+ */
+ public Memory readableMemory()
+ {
+ return memory;
+ }
+
/**
* Direct, writable access to this frame's memory. Used by operations that modify the frame in-place, like
* {@link org.apache.druid.frame.write.FrameSort}.
@@ -358,8 +282,7 @@ public class Frame
public long writeTo(
final WritableByteChannel channel,
final boolean compress,
- @Nullable final ByteBuffer compressionBuffer,
- ByteTracker byteTracker
+ @Nullable final ByteBuffer compressionBuffer
) throws IOException
{
if (compress) {
@@ -373,54 +296,30 @@ public class Frame
);
}
- final ByteBuffer frameBuffer;
- final int frameBufferPosition;
- final int compressedFrameLength;
-
- if (memory.hasByteBuffer()) {
- // Optimized path when Memory is backed by ByteBuffer.
- frameBuffer = memory.getByteBuffer();
- frameBufferPosition = Ints.checkedCast(memory.getRegionOffset());
- } else {
- // Copy to byte array first, then decompress.
- final byte[] frameBytes = new byte[Ints.checkedCast(numBytes)];
- memory.getByteArray(0, frameBytes, 0, Ints.checkedCast(numBytes));
- frameBuffer = ByteBuffer.wrap(frameBytes);
- frameBufferPosition = 0;
- }
-
- compressedFrameLength = LZ4_COMPRESSOR.compress(
- frameBuffer,
- frameBufferPosition,
- Ints.checkedCast(numBytes),
- compressionBuffer,
- COMPRESSED_FRAME_HEADER_SIZE,
- compressionBuffer.capacity() - COMPRESSED_FRAME_ENVELOPE_SIZE
- );
+ final ByteBuffer compressedBuffer = FrameCompression.compress(memory, numBytes, compressionBuffer);
+ Channels.writeFully(channel, compressedBuffer);
- compressionBuffer.order(ByteOrder.LITTLE_ENDIAN)
- .limit(COMPRESSED_FRAME_ENVELOPE_SIZE + compressedFrameLength)
- .position(0);
-
- compressionBuffer.put(0, CompressionStrategy.LZ4.getId())
- .putLong(Byte.BYTES, compressedFrameLength)
- .putLong(Byte.BYTES + Long.BYTES, numBytes);
-
- final long checksum = Memory.wrap(compressionBuffer)
- .xxHash64(0, COMPRESSED_FRAME_HEADER_SIZE + compressedFrameLength, CHECKSUM_SEED);
-
- compressionBuffer.putLong(COMPRESSED_FRAME_HEADER_SIZE + compressedFrameLength, checksum);
- byteTracker.reserve(compressionBuffer.remaining());
- Channels.writeFully(channel, compressionBuffer);
-
- return COMPRESSED_FRAME_ENVELOPE_SIZE + compressedFrameLength;
+ return compressedBuffer.limit();
} else {
- byteTracker.reserve(numBytes);
memory.writeTo(0, numBytes, channel);
return numBytes;
}
}
+ /**
+ * Returns this frame as a {@link FrameRowsAndColumns} without a signature.
+ */
+ public FrameRowsAndColumns asRAC()
+ {
+ if (frameType.isRowBased()) {
+ return new RowBasedFrameRowsAndColumns(this, null);
+ } else if (frameType.isColumnar()) {
+ return new ColumnBasedFrameRowsAndColumns(this, null);
+ } else {
+ throw DruidException.defensive("FrameType[%s] was neither row-based nor columnar", frameType);
+ }
+ }
+
/**
* Returns this object if {@link #type()} is row-based, otherwise throws an error.
*/
@@ -445,6 +344,22 @@ public class Frame
return this;
}
+ @Override
+ public boolean equals(Object o)
+ {
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ Frame frame = (Frame) o;
+ return Objects.equals(memory, frame.memory);
+ }
+
+ @Override
+ public int hashCode()
+ {
+ return Objects.hashCode(memory);
+ }
+
/**
* Perform basic frame validations, ensuring that the length of frame and region locations are correct.
*/
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/BlockingQueueFrameChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/BlockingQueueFrameChannel.java
index 12e14d0dfea..7a41a781147 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/BlockingQueueFrameChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/BlockingQueueFrameChannel.java
@@ -23,10 +23,11 @@ import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import com.google.errorprone.annotations.concurrent.GuardedBy;
-import org.apache.druid.frame.Frame;
+import it.unimi.dsi.fastutil.objects.ObjectIntPair;
import org.apache.druid.java.util.common.Either;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import javax.annotation.Nullable;
import java.util.ArrayDeque;
@@ -43,7 +44,7 @@ import java.util.Optional;
public class BlockingQueueFrameChannel
{
@SuppressWarnings("OptionalUsedAsFieldOrParameterType")
- private static final Optional<Either<Throwable, FrameWithPartition>> END_MARKER = Optional.empty();
+ private static final Optional<Either<Throwable, ObjectIntPair<RowsAndColumns>>> END_MARKER = Optional.empty();
private final int maxQueuedFrames;
private final Object lock = new Object();
@@ -52,7 +53,7 @@ public class BlockingQueueFrameChannel
private final Readable readable;
@GuardedBy("lock")
- private final ArrayDeque<Optional<Either<Throwable, FrameWithPartition>>> queue;
+ private final ArrayDeque<Optional<Either<Throwable, ObjectIntPair<RowsAndColumns>>>> queue;
@GuardedBy("lock")
private SettableFuture<?> readyForWritingFuture = null;
@@ -129,7 +130,7 @@ public class BlockingQueueFrameChannel
private class Writable implements WritableFrameChannel
{
@Override
- public void write(FrameWithPartition frame)
+ public void write(RowsAndColumns rac, int partitionNumber)
{
synchronized (lock) {
if (isFinished()) {
@@ -138,7 +139,7 @@ public class BlockingQueueFrameChannel
// Caller should have checked if this channel was ready for writing.
throw new ISE("Channel has no capacity");
} else {
- if (!queue.offer(Optional.of(Either.value(frame)))) {
+ if (!queue.offer(Optional.of(Either.value(ObjectIntPair.of(rac, partitionNumber))))) {
// If this happens, it's a bug in this class's capacity-counting.
throw new ISE("Channel had capacity, but could not add frame");
}
@@ -196,7 +197,7 @@ public class BlockingQueueFrameChannel
public boolean isClosed()
{
synchronized (lock) {
- final Optional<Either<Throwable, FrameWithPartition>> lastElement = queue.peekLast();
+ final Optional<Either<Throwable, ObjectIntPair<RowsAndColumns>>> lastElement = queue.peekLast();
return END_MARKER.equals(lastElement);
}
}
@@ -219,9 +220,9 @@ public class BlockingQueueFrameChannel
}
@Override
- public Frame read()
+ public RowsAndColumns read()
{
- final Optional<Either<Throwable, FrameWithPartition>> next;
+ final Optional<Either<Throwable, ObjectIntPair<RowsAndColumns>>> next;
synchronized (lock) {
if (isFinished()) {
@@ -237,7 +238,7 @@ public class BlockingQueueFrameChannel
notifyWriter();
}
- return next.get().valueOrThrow().frame();
+ return next.get().valueOrThrow().left();
}
@Override
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/ComposingReadableFrameChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/ComposingReadableFrameChannel.java
index db43ee249a1..76aed14fe76 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/ComposingReadableFrameChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/ComposingReadableFrameChannel.java
@@ -23,7 +23,7 @@ import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
-import org.apache.druid.frame.Frame;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import javax.annotation.concurrent.NotThreadSafe;
import java.util.ArrayList;
@@ -103,7 +103,7 @@ public class ComposingReadableFrameChannel implements ReadableFrameChannel
}
@Override
- public Frame read()
+ public RowsAndColumns read()
{
return currentChannel.read();
}
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/ComposingWritableFrameChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/ComposingWritableFrameChannel.java
index 03a2cf1a8a0..96c5271dcf8 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/ComposingWritableFrameChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/ComposingWritableFrameChannel.java
@@ -28,6 +28,7 @@ import org.apache.druid.frame.processor.PartitionedOutputChannel;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.query.ResourceLimitExceededException;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import javax.annotation.Nullable;
import java.io.IOException;
@@ -87,25 +88,25 @@ public class ComposingWritableFrameChannel implements WritableFrameChannel
}
@Override
- public void write(FrameWithPartition frameWithPartition) throws IOException
+ public void write(RowsAndColumns rac, int partitionNum) throws IOException
{
if (currentIndex >= writableChannelSuppliers.size()) {
throw new ISE("No more channels available to write. Total available channels : " + writableChannelSuppliers.size());
}
if (partitionNumber != null
- && frameWithPartition.partition() != FrameWithPartition.NO_PARTITION
- && frameWithPartition.partition() != partitionNumber) {
+ && partitionNum != WritableFrameChannel.NO_PARTITION
+ && partitionNum != partitionNumber) {
throw DruidException.defensive(
"Invalid partition number[%d], expected[%d]",
- frameWithPartition.partition(),
+ partitionNum,
partitionNumber
);
}
try {
- final int writtenPartition = partitionNumber != null ? partitionNumber : frameWithPartition.partition();
- writableChannelSuppliers.get(currentIndex).get().write(frameWithPartition);
+ final int writtenPartition = partitionNumber != null ? partitionNumber : partitionNum;
+ writableChannelSuppliers.get(currentIndex).get().write(rac, partitionNum);
partitionToChannelMap.computeIfAbsent(writtenPartition, k -> Sets.newHashSetWithExpectedSize(1))
.add(currentIndex);
}
@@ -124,7 +125,7 @@ public class ComposingWritableFrameChannel implements WritableFrameChannel
if (currentIndex >= writableChannelSuppliers.size()) {
throw rlee;
}
- write(frameWithPartition);
+ write(rac, partitionNum);
}
}
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/DurableStoragePartitionedReadableFrameChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/DurableStoragePartitionedReadableFrameChannel.java
index 0e0e6680df3..6242da83cbe 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/DurableStoragePartitionedReadableFrameChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/DurableStoragePartitionedReadableFrameChannel.java
@@ -21,8 +21,10 @@ package org.apache.druid.frame.channel;
import org.apache.druid.frame.file.FrameFileFooter;
import org.apache.druid.frame.file.FrameFileWriter;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.storage.StorageConnector;
+import javax.annotation.Nullable;
import java.io.File;
import java.io.IOException;
import java.io.UncheckedIOException;
@@ -37,12 +39,16 @@ public class DurableStoragePartitionedReadableFrameChannel implements Partitione
private final ExecutorService remoteInputStreamPool;
private final File footerFile;
+ @Nullable
+ private final WireTransferableContext wireTransferableContext;
+
public DurableStoragePartitionedReadableFrameChannel(
StorageConnector storageConnector,
Supplier<FrameFileFooter> frameFileFooterSupplier,
String frameFileFullPath,
ExecutorService remoteInputStreamPool,
- File footerFile
+ File footerFile,
+ @Nullable WireTransferableContext wireTransferableContext
)
{
this.storageConnector = storageConnector;
@@ -50,6 +56,7 @@ public class DurableStoragePartitionedReadableFrameChannel implements Partitione
this.frameFileFullPath = frameFileFullPath;
this.remoteInputStreamPool = remoteInputStreamPool;
this.footerFile = footerFile;
+ this.wireTransferableContext = wireTransferableContext;
}
@Override
@@ -71,7 +78,8 @@ public class DurableStoragePartitionedReadableFrameChannel implements Partitione
storageConnector.readRange(frameFileFullPath, startByte, endByte - startByte),
frameFileFullPath,
remoteInputStreamPool,
- true
+ true,
+ wireTransferableContext
);
}
catch (IOException e) {
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/FrameChannelSequence.java b/processing/src/main/java/org/apache/druid/frame/channel/FrameChannelSequence.java
index d56e7a2aa88..83075b0e48e 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/FrameChannelSequence.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/FrameChannelSequence.java
@@ -20,8 +20,8 @@
package org.apache.druid.frame.channel;
import org.apache.druid.common.guava.FutureUtils;
-import org.apache.druid.frame.Frame;
import org.apache.druid.java.util.common.guava.BaseSequence;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import java.io.Closeable;
import java.util.Iterator;
@@ -29,12 +29,12 @@ import java.util.NoSuchElementException;
/**
* Adapter that converts a {@link ReadableFrameChannel} into a {@link org.apache.druid.java.util.common.guava.Sequence}
- * of {@link Frame}.
+ * of {@link RowsAndColumns}.
*
* This class does blocking reads on the channel, rather than nonblocking reads. Therefore, it is preferable to use
* {@link ReadableFrameChannel} directly whenever nonblocking reads are desired.
*/
-public class FrameChannelSequence extends BaseSequence<Frame, FrameChannelSequence.FrameChannelIterator>
+public class FrameChannelSequence extends BaseSequence<RowsAndColumns, FrameChannelSequence.FrameChannelIterator>
{
public FrameChannelSequence(final ReadableFrameChannel channel)
{
@@ -56,7 +56,7 @@ public class FrameChannelSequence extends BaseSequence<Frame, FrameChannelSequen
);
}
- public static class FrameChannelIterator implements Iterator<Frame>, Closeable
+ public static class FrameChannelIterator implements Iterator<RowsAndColumns>, Closeable
{
private final ReadableFrameChannel channel;
@@ -74,7 +74,7 @@ public class FrameChannelSequence extends BaseSequence<Frame, FrameChannelSequen
}
@Override
- public Frame next()
+ public RowsAndColumns next()
{
// Side effect of hasNext(): we also call await() to ensure the next frame is available.
if (!hasNext()) {
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/ReadableByteChunksFrameChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/ReadableByteChunksFrameChannel.java
index 71e844f2a83..4397b780575 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/ReadableByteChunksFrameChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/ReadableByteChunksFrameChannel.java
@@ -29,12 +29,17 @@ import com.google.common.util.concurrent.SettableFuture;
import com.google.errorprone.annotations.concurrent.GuardedBy;
import org.apache.datasketches.memory.Memory;
import org.apache.datasketches.memory.WritableMemory;
+import org.apache.druid.error.DruidException;
import org.apache.druid.frame.Frame;
import org.apache.druid.frame.file.FrameFileWriter;
+import org.apache.druid.frame.wire.FrameCompression;
+import org.apache.druid.java.util.common.ByteBufferUtils;
import org.apache.druid.java.util.common.Either;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import javax.annotation.Nullable;
import java.nio.ByteOrder;
@@ -60,9 +65,10 @@ public class ReadableByteChunksFrameChannel implements ReadableFrameChannel
private static final long MAX_FRAME_SIZE_BYTES = 100_000_000;
private static final int UNKNOWN_LENGTH = -1;
+ private static final byte NO_MARKER = 0;
private static final int FRAME_MARKER_BYTES = Byte.BYTES;
private static final int FRAME_MARKER_AND_COMPRESSED_ENVELOPE_BYTES =
- Byte.BYTES + Frame.COMPRESSED_FRAME_ENVELOPE_SIZE;
+ Byte.BYTES + FrameCompression.COMPRESSED_DATA_ENVELOPE_SIZE;
private enum StreamPart
{
@@ -99,28 +105,53 @@ public class ReadableByteChunksFrameChannel implements ReadableFrameChannel
@GuardedBy("lock")
private long nextCompressedFrameLength = UNKNOWN_LENGTH;
+ /**
+ * The marker type of the next frame to be read. Set when {@link #nextCompressedFrameLength} is determined.
+ * Will be one of {@link FrameFileWriter#MARKER_FRAME}, {@link FrameFileWriter#MARKER_RAC}, or {@link #NO_MARKER}
+ * if no frame is currently ready.
+ */
+ @GuardedBy("lock")
+ private byte nextMarkerType = NO_MARKER;
+
@GuardedBy("lock")
private StreamPart streamPart;
private final boolean framesOnly;
- private ReadableByteChunksFrameChannel(String id, long bytesLimit, boolean framesOnly)
+ @Nullable
+ private final WireTransferableContext wtContext;
+
+ private ReadableByteChunksFrameChannel(
+ String id,
+ long bytesLimit,
+ boolean framesOnly,
+ @Nullable WireTransferableContext wtContext
+ )
{
this.id = Preconditions.checkNotNull(id, "id");
this.bytesLimit = bytesLimit;
this.streamPart = framesOnly ? StreamPart.FRAMES : StreamPart.MAGIC;
this.framesOnly = framesOnly;
+ this.wtContext = wtContext;
}
/**
* Create a channel that aims to limit its memory footprint to one frame. The channel exerts backpressure
* from {@link #addChunk} immediately once a full frame has been buffered.
+ *
+ * @param id identifier for this channel
+ * @param framesOnly if true, expects stream to include only FRAMES part (no header or footer)
+ * @param wireTransferableContext context for wire transfer serde
*/
- public static ReadableByteChunksFrameChannel create(final String id, boolean framesOnly)
+ public static ReadableByteChunksFrameChannel create(
+ final String id,
+ boolean framesOnly,
+ @Nullable WireTransferableContext wireTransferableContext
+ )
{
// Set byte limit to 1, so backpressure will be exerted as soon as we have a full frame buffered.
// (The bytesLimit is soft: it will be exceeded if needed to store a complete frame.)
- return new ReadableByteChunksFrameChannel(id, 1, framesOnly);
+ return new ReadableByteChunksFrameChannel(id, 1, framesOnly, wireTransferableContext);
}
/**
@@ -188,6 +219,7 @@ public class ReadableByteChunksFrameChannel implements ReadableFrameChannel
chunks.clear();
chunks.add(Either.error(t));
nextCompressedFrameLength = UNKNOWN_LENGTH;
+ nextMarkerType = NO_MARKER;
doneWriting();
}
}
@@ -229,7 +261,7 @@ public class ReadableByteChunksFrameChannel implements ReadableFrameChannel
}
@Override
- public Frame read()
+ public RowsAndColumns read()
{
synchronized (lock) {
if (canReadError()) {
@@ -238,11 +270,12 @@ public class ReadableByteChunksFrameChannel implements ReadableFrameChannel
Throwables.propagateIfPossible(t);
throw new RuntimeException(t);
} else if (canReadFrame()) {
- return nextFrame();
+ return nextRAC();
} else if (noMoreWrites) {
// The last few chunks are an incomplete or missing frame.
chunks.clear();
nextCompressedFrameLength = UNKNOWN_LENGTH;
+ nextMarkerType = NO_MARKER;
throw new ISE(
"Incomplete or missing frame at end of stream (id = %s, position = %d)",
@@ -278,6 +311,7 @@ public class ReadableByteChunksFrameChannel implements ReadableFrameChannel
synchronized (lock) {
chunks.clear();
nextCompressedFrameLength = UNKNOWN_LENGTH;
+ nextMarkerType = NO_MARKER;
// Setting "noMoreWrites" causes the upstream entity to realize this channel has closed the next time
// it calls "addChunk".
@@ -312,21 +346,23 @@ public class ReadableByteChunksFrameChannel implements ReadableFrameChannel
}
}
- private Frame nextFrame()
+ private RowsAndColumns nextRAC()
{
- final Memory frameMemory;
+ final Memory compressedMemory;
+ final byte markerType;
synchronized (lock) {
if (!canReadFrame()) {
throw new ISE("Frame of size [%,d] not yet ready to read", nextCompressedFrameLength);
}
- if (nextCompressedFrameLength > Integer.MAX_VALUE - FRAME_MARKER_BYTES - Frame.COMPRESSED_FRAME_ENVELOPE_SIZE) {
+ if (nextCompressedFrameLength > Integer.MAX_VALUE - FRAME_MARKER_BYTES - FrameCompression.COMPRESSED_DATA_ENVELOPE_SIZE) {
throw new ISE("Cannot read frame of size [%,d] bytes", nextCompressedFrameLength);
}
+ markerType = nextMarkerType;
final int numBytes = Ints.checkedCast(FRAME_MARKER_AND_COMPRESSED_ENVELOPE_BYTES + nextCompressedFrameLength);
- frameMemory = copyFromQueuedChunks(numBytes).region(
+ compressedMemory = copyFromQueuedChunks(numBytes).region(
FRAME_MARKER_BYTES,
FRAME_MARKER_AND_COMPRESSED_ENVELOPE_BYTES + nextCompressedFrameLength - FRAME_MARKER_BYTES
);
@@ -334,9 +370,26 @@ public class ReadableByteChunksFrameChannel implements ReadableFrameChannel
updateStreamState();
}
- final Frame frame = Frame.decompress(frameMemory, 0, frameMemory.getCapacity());
- log.debug("Read frame with [%,d] rows and [%,d] bytes.", frame.numRows(), frame.numBytes());
- return frame;
+ final byte[] decompressedBytes =
+ FrameCompression.decompress(compressedMemory, 0, compressedMemory.getCapacity());
+
+ switch (markerType) {
+ case FrameFileWriter.MARKER_FRAME:
+ final Frame frame = Frame.wrap(decompressedBytes);
+ log.debug("Read frame with rows[%,d] and bytes[%,d].", frame.numRows(), frame.numBytes());
+ return frame.asRAC();
+
+ case FrameFileWriter.MARKER_RAC:
+ if (wtContext == null) {
+ throw DruidException.defensive("Cannot read RAC, no WireTransferableContext");
+ }
+ final RowsAndColumns rac = wtContext.deserialize(ByteBufferUtils.wrapLE(decompressedBytes));
+ log.debug("Read RAC with rows[%,d] and bytes[%,d].", rac.numRows(), decompressedBytes.length);
+ return rac;
+
+ default:
+ throw new ISE("Unexpected marker type[%d]", markerType);
+ }
}
@GuardedBy("lock")
@@ -358,12 +411,15 @@ public class ReadableByteChunksFrameChannel implements ReadableFrameChannel
if (streamPart == StreamPart.FRAMES) {
if (bytesBuffered >= Byte.BYTES) {
final Memory memory = copyFromQueuedChunks(1);
+ final byte markerByte = memory.getByte(0);
- if (memory.getByte(0) == FrameFileWriter.MARKER_FRAME) {
+ if (markerByte == FrameFileWriter.MARKER_FRAME || markerByte == FrameFileWriter.MARKER_RAC) {
// Read nextFrameLength if needed; otherwise do nothing.
- final int bytesRequiredToReadLength = FRAME_MARKER_BYTES + Frame.COMPRESSED_FRAME_HEADER_SIZE;
+ // Both MARKER_FRAME and MARKER_RAC use the same compression envelope format.
+ final int bytesRequiredToReadLength = FRAME_MARKER_BYTES + FrameCompression.COMPRESSED_DATA_HEADER_SIZE;
if (nextCompressedFrameLength == UNKNOWN_LENGTH && bytesBuffered >= bytesRequiredToReadLength) {
+ nextMarkerType = markerByte;
nextCompressedFrameLength = copyFromQueuedChunks(bytesRequiredToReadLength)
.getLong(FRAME_MARKER_BYTES + Byte.BYTES /* Compression strategy byte */);
@@ -371,9 +427,10 @@ public class ReadableByteChunksFrameChannel implements ReadableFrameChannel
throw new ISE("Invalid frame size (size = %,d B)", nextCompressedFrameLength);
}
}
- } else if (memory.getByte(0) == FrameFileWriter.MARKER_NO_MORE_FRAMES) {
+ } else if (markerByte == FrameFileWriter.MARKER_NO_MORE_FRAMES) {
streamPart = StreamPart.FOOTER;
nextCompressedFrameLength = UNKNOWN_LENGTH;
+ nextMarkerType = NO_MARKER;
} else {
throw new ISE("Invalid midstream marker (id = %s, position = %d)", id, bytesAdded - bytesBuffered);
}
@@ -461,7 +518,8 @@ public class ReadableByteChunksFrameChannel implements ReadableFrameChannel
bytesBuffered -= numBytes;
- // Clear nextFrameLength; it won't be accurate anymore after deleting bytes.
+ // Clear nextFrameLength and nextMarkerType; they won't be accurate anymore after deleting bytes.
nextCompressedFrameLength = UNKNOWN_LENGTH;
+ nextMarkerType = NO_MARKER;
}
}
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/ReadableConcatFrameChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/ReadableConcatFrameChannel.java
index 2cdfdd96292..c1804f3b160 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/ReadableConcatFrameChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/ReadableConcatFrameChannel.java
@@ -21,7 +21,7 @@ package org.apache.druid.frame.channel;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
-import org.apache.druid.frame.Frame;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import javax.annotation.Nullable;
import java.util.Iterator;
@@ -72,7 +72,7 @@ public class ReadableConcatFrameChannel implements ReadableFrameChannel
}
@Override
- public Frame read()
+ public RowsAndColumns read()
{
if (!canRead()) {
throw new NoSuchElementException();
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/ReadableFileFrameChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/ReadableFileFrameChannel.java
index 7da6550ccca..07ab7c29088 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/ReadableFileFrameChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/ReadableFileFrameChannel.java
@@ -21,9 +21,10 @@ package org.apache.druid.frame.channel;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
-import org.apache.druid.frame.Frame;
import org.apache.druid.frame.file.FrameFile;
import org.apache.druid.java.util.common.IAE;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
import java.io.IOException;
import java.util.NoSuchElementException;
@@ -38,12 +39,19 @@ import java.util.NoSuchElementException;
public class ReadableFileFrameChannel implements ReadableFrameChannel
{
private final FrameFile frameFile;
+ private final WireTransferable.ConcreteDeserializer deserializer;
private final int endFrame;
private int currentFrame;
- public ReadableFileFrameChannel(final FrameFile frameFile, final int startFrame, final int endFrame)
+ public ReadableFileFrameChannel(
+ final FrameFile frameFile,
+ final int startFrame,
+ final int endFrame,
+ final WireTransferable.ConcreteDeserializer deserializer
+ )
{
this.frameFile = frameFile;
+ this.deserializer = deserializer;
this.currentFrame = startFrame;
this.endFrame = endFrame;
@@ -60,9 +68,12 @@ public class ReadableFileFrameChannel implements ReadableFrameChannel
}
}
- public ReadableFileFrameChannel(final FrameFile frameFile)
+ public ReadableFileFrameChannel(
+ final FrameFile frameFile,
+ final WireTransferable.ConcreteDeserializer deserializer
+ )
{
- this(frameFile, 0, frameFile.numFrames());
+ this(frameFile, 0, frameFile.numFrames(), deserializer);
}
@Override
@@ -78,13 +89,13 @@ public class ReadableFileFrameChannel implements ReadableFrameChannel
}
@Override
- public Frame read()
+ public RowsAndColumns read()
{
if (isFinished()) {
throw new NoSuchElementException();
}
- return frameFile.frame(currentFrame++);
+ return frameFile.rac(currentFrame++, deserializer);
}
@Override
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/ReadableFrameChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/ReadableFrameChannel.java
index 44c108f7ac1..dc8ea919662 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/ReadableFrameChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/ReadableFrameChannel.java
@@ -20,54 +20,80 @@
package org.apache.druid.frame.channel;
import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.error.DruidException;
import org.apache.druid.frame.Frame;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import java.io.Closeable;
+import java.util.NoSuchElementException;
/**
- * Interface for reading a sequence of frames. Supports nonblocking reads through the {@link #canRead()} and
+ * Interface for reading a sequence of batches of data. Supports nonblocking reads through the {@link #canRead()} and
* {@link #readabilityFuture()} methods.
*
* May be implemented using an in-memory queue, disk file, stream, etc.
*
* Channels implementing this interface are used by a single reader; they do not support concurrent reads.
+ *
+ * Despite its name, instances of this class can typically return any {@link RowsAndColumns} through the
+ * {@link #read()} method.
*/
public interface ReadableFrameChannel extends Closeable
{
/**
- * Returns whether this channel is finished. Finished channels will not generate any further frames or errors.
+ * Returns whether this channel is finished. Finished channels will not generate any further data batches or errors.
*
* Generally, once you discover that a channel is finished, you should call {@link #close()} and then
* discard it.
*
- * Note that it is possible for a channel to be unfinished and also have no available frames or errors. This happens
- * when it is not in a ready-for-reading state. See {@link #readabilityFuture()} for details.
+ * Note that it is possible for a channel to be unfinished and also have no available data batches or errors.
+ * This happens when it is not in a ready-for-reading state. See {@link #readabilityFuture()} for details.
*/
boolean isFinished();
/**
- * Returns whether this channel has a frame or error condition currently available. If this method returns true, then
- * you can call {@link #read()} to retrieve the frame or error.
+ * Returns whether this channel has a batch of data or error condition currently available. If this method returns
+ * true, then you can call {@link #readFrame()} or {@link #read()} to retrieve the batch or error.
*
- * Note that it is possible for a channel to be unfinished and also have no available frames or errors. This happens
+ * Note that it is possible for a channel to be unfinished and also have no available batches or errors. This happens
* when it is not in a ready-for-reading state. See {@link #readabilityFuture()} for details.
*/
boolean canRead();
/**
- * Returns the next available frame from this channel.
+ * Returns the next available batch of data from this channel as a {@link RowsAndColumns}.
+ *
+ * Before calling this method, you should check {@link #canRead()} to ensure there is a batch of data or
+ * error available.
+ *
+ * @throws NoSuchElementException if there is no batch currently available
+ */
+ RowsAndColumns read();
+
+ /**
+ * Returns the next available batch of data from this channel as a {@link Frame}.
*
- * Before calling this method, you should check {@link #canRead()} to ensure there is a frame or
+ * Before calling this method, you should check {@link #canRead()} to ensure there is a batch of data or
* error available.
*
- * @throws java.util.NoSuchElementException if there is no frame currently available
+ * @throws NoSuchElementException if there is no batch currently available
+ * @throws DruidException if the available batch was not available as a {@link Frame}
*/
- Frame read();
+ default Frame readFrame()
+ {
+ final RowsAndColumns rac = read();
+ final Frame frame = rac.as(Frame.class);
+ if (frame != null) {
+ return frame;
+ } else {
+ throw DruidException.defensive("Got RAC[%s] which is not a frame", rac);
+ }
+ }
/**
* Returns a future that will resolve when either {@link #isFinished()} or {@link #canRead()} would
* return true. The future will never resolve to an exception. If something exceptional has happened, the exception
- * can be retrieved from {@link #read()}.
+ * can be retrieved from {@link #readFrame()} or {@link #read()}.
*/
ListenableFuture<?> readabilityFuture();
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/ReadableInputStreamFrameChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/ReadableInputStreamFrameChannel.java
index 5f03a5e8efb..cccce0147dd 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/ReadableInputStreamFrameChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/ReadableInputStreamFrameChannel.java
@@ -22,10 +22,12 @@ package org.apache.druid.frame.channel;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.errorprone.annotations.concurrent.GuardedBy;
import org.apache.commons.io.IOUtils;
-import org.apache.druid.frame.Frame;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.concurrent.Execs;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
+import javax.annotation.Nullable;
import java.io.InputStream;
import java.util.Arrays;
import java.util.concurrent.ExecutorService;
@@ -90,17 +92,24 @@ public class ReadableInputStreamFrameChannel implements ReadableFrameChannel
/**
* Create an instance of this class and immediately start reading from the provided InputStream.
+ *
+ * @param inputStream the input stream to read from
+ * @param id identifier for this channel
+ * @param executorService executor service for reading from the stream
+ * @param framesOnly if true, expects stream to start at FRAMES part (no MAGIC header)
+ * @param wireTransferableContext context for wire transfer deserialization; may be null if all data uses legacy frames
*/
public static ReadableInputStreamFrameChannel open(
InputStream inputStream,
String id,
ExecutorService executorService,
- boolean framesOnly
+ boolean framesOnly,
+ @Nullable WireTransferableContext wireTransferableContext
)
{
return new ReadableInputStreamFrameChannel(
inputStream,
- ReadableByteChunksFrameChannel.create(id, framesOnly),
+ ReadableByteChunksFrameChannel.create(id, framesOnly, wireTransferableContext),
executorService
);
}
@@ -124,7 +133,7 @@ public class ReadableInputStreamFrameChannel implements ReadableFrameChannel
}
@Override
- public Frame read()
+ public RowsAndColumns read()
{
synchronized (lock) {
startReading();
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/ReadableNilFrameChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/ReadableNilFrameChannel.java
index 184629086d0..dc8f82add0e 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/ReadableNilFrameChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/ReadableNilFrameChannel.java
@@ -21,7 +21,7 @@ package org.apache.druid.frame.channel;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
-import org.apache.druid.frame.Frame;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import java.util.NoSuchElementException;
@@ -49,7 +49,7 @@ public class ReadableNilFrameChannel implements ReadableFrameChannel
}
@Override
- public Frame read()
+ public RowsAndColumns read()
{
throw new NoSuchElementException();
}
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/WritableFrameChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/WritableFrameChannel.java
index bb1527305fd..170ce4b6844 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/WritableFrameChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/WritableFrameChannel.java
@@ -21,36 +21,67 @@ package org.apache.druid.frame.channel;
import com.google.common.util.concurrent.ListenableFuture;
import org.apache.druid.frame.Frame;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import javax.annotation.Nullable;
import java.io.Closeable;
import java.io.IOException;
/**
- * Interface for writing a sequence of frames. Supports nonblocking writes through the {@link #writabilityFuture()}
- * method.
+ * Interface for writing a sequence of result batches. Supports nonblocking writes through the
+ * {@link #writabilityFuture()} method.
*
* May be implemented using an in-memory queue, disk file, stream, etc.
*
* Channels implementing this interface are used by a single writer; they do not support concurrent writes.
+ *
+ * Despite its name, instances of this class can typically accept any {@link RowsAndColumns} through the
+ * {@link #write(RowsAndColumns)} method.
*/
public interface WritableFrameChannel extends Closeable
{
/**
- * Writes a frame with an attached partition number.
+ * Partition number used when there is no meaningful partition to associate with a batch of data.
+ */
+ int NO_PARTITION = -1;
+
+ /**
+ * Writes a batch of data with an attached partition number.
+ *
+ * May throw an exception if {@link #writabilityFuture()} is unresolved.
+ */
+ void write(RowsAndColumns rac, int partitionNumber) throws IOException;
+
+ /**
+ * Writes a batch of data without an attached partition number.
*
* May throw an exception if {@link #writabilityFuture()} is unresolved.
*/
- void write(FrameWithPartition frameWithPartition) throws IOException;
+ default void write(RowsAndColumns rac) throws IOException
+ {
+ write(rac, NO_PARTITION);
+ }
+
+ /**
+ * Writes a frame with an attached partition number. Equivalent to calling {@link #write(RowsAndColumns, int)}
+ * with {@link Frame#asRAC()}.
+ *
+ * May throw an exception if {@link #writabilityFuture()} is unresolved.
+ */
+ default void write(Frame frame, int partitionNumber) throws IOException
+ {
+ write(frame.asRAC(), partitionNumber);
+ }
/**
- * Writes a frame without an attached partition number.
+ * Writes a frame without an attached partition number. Equivalent to calling {@link #write(RowsAndColumns)} with
+ * {@link Frame#asRAC()}.
*
* May throw an exception if {@link #writabilityFuture()} is unresolved.
*/
default void write(Frame frame) throws IOException
{
- write(new FrameWithPartition(frame, FrameWithPartition.NO_PARTITION));
+ write(frame.asRAC(), NO_PARTITION);
}
/**
@@ -79,8 +110,8 @@ public interface WritableFrameChannel extends Closeable
boolean isClosed();
/**
- * Returns a future that resolves when {@link #write} is able to receive a new frame without blocking or throwing
- * an exception. The future never resolves to an exception.
+ * Returns a future that resolves when {@link #write} is able to receive a new batch of data without blocking or
+ * throwing an exception. The future never resolves to an exception.
*/
ListenableFuture<?> writabilityFuture();
}
diff --git a/processing/src/main/java/org/apache/druid/frame/channel/WritableFrameFileChannel.java b/processing/src/main/java/org/apache/druid/frame/channel/WritableFrameFileChannel.java
index 6463abe974e..2541ef619b3 100644
--- a/processing/src/main/java/org/apache/druid/frame/channel/WritableFrameFileChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/channel/WritableFrameFileChannel.java
@@ -22,6 +22,7 @@ package org.apache.druid.frame.channel;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import org.apache.druid.frame.file.FrameFileWriter;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import javax.annotation.Nullable;
import java.io.IOException;
@@ -40,9 +41,9 @@ public class WritableFrameFileChannel implements WritableFrameChannel
}
@Override
- public void write(FrameWithPartition frame) throws IOException
+ public void write(RowsAndColumns rac, int partitionNumber) throws IOException
{
- writer.writeFrame(frame.frame(), frame.partition());
+ writer.write(rac, partitionNumber);
}
@Override
diff --git a/processing/src/main/java/org/apache/druid/frame/file/FrameFile.java b/processing/src/main/java/org/apache/druid/frame/file/FrameFile.java
index 358997efcdd..14f89c4b875 100644
--- a/processing/src/main/java/org/apache/druid/frame/file/FrameFile.java
+++ b/processing/src/main/java/org/apache/druid/frame/file/FrameFile.java
@@ -20,8 +20,12 @@
package org.apache.druid.frame.file;
import org.apache.datasketches.memory.Memory;
+import org.apache.druid.error.DruidException;
import org.apache.druid.frame.Frame;
import org.apache.druid.frame.channel.ByteTracker;
+import org.apache.druid.frame.channel.ReadableByteChunksFrameChannel;
+import org.apache.druid.frame.wire.FrameCompression;
+import org.apache.druid.java.util.common.ByteBufferUtils;
import org.apache.druid.java.util.common.FileUtils;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.IOE;
@@ -30,6 +34,8 @@ import org.apache.druid.java.util.common.MappedByteBufferHandler;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.io.Closer;
import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
import org.apache.druid.segment.ReferenceCountingCloseableObject;
import org.apache.druid.utils.CloseableUtils;
@@ -48,7 +54,7 @@ import java.util.EnumSet;
*
* Frame files are written by {@link FrameFileWriter}.
*
- * Frame files can optionally be partitioned, by providing partition numbers to the {@link FrameFileWriter#writeFrame}
+ * Frame files can optionally be partitioned, by providing partition numbers to the {@link FrameFileWriter#write}
* method when creating the file. Partitions are contiguous within the frame file.
*
* Frame files can contain up to {@link Integer#MAX_VALUE} frames. Generally, frames are on the order of 1 MB in size,
@@ -70,6 +76,9 @@ import java.util.EnumSet;
*
* Instances of this class are not thread-safe. For sharing across threads, use {@link #newReference()} to create
* an additional reference.
+ *
+ * @see FrameFileWriter writer for this file format
+ * @see ReadableByteChunksFrameChannel streaming reader for this file format
*/
public class FrameFile implements Closeable
{
@@ -309,36 +318,62 @@ public class FrameFile implements Closeable
}
/**
- * Reads a frame from the file.
+ * Reads a RowsAndColumns from the file.
+ *
+ * @param batchNumber the batch number to read. Starts at 0, ends at {@link #numFrames()} (exclusive)
+ * @param deserializer deserializer for WireTransferable format
+ *
+ * @return the RowsAndColumns
*/
- public Frame frame(final int frameNumber)
+ public RowsAndColumns rac(
+ final int batchNumber,
+ @Nullable final WireTransferable.ConcreteDeserializer deserializer
+ )
{
checkOpen();
- if (frameNumber < 0 || frameNumber >= numFrames()) {
- throw new IAE("Frame [%,d] out of bounds", frameNumber);
+ if (batchNumber < 0 || batchNumber >= numFrames()) {
+ throw new IAE("Batch[%,d] out of bounds", batchNumber);
}
- final long frameEnd = frameFileFooter.getFrameEndPosition(frameNumber);
- final long frameStart;
+ final long markerPosition = getMarkerPosition(batchNumber);
+ final long dataStart = markerPosition + Byte.BYTES;
+ final long dataEnd = frameFileFooter.getFrameEndPosition(batchNumber);
- if (frameNumber == 0) {
- frameStart = FrameFileWriter.MAGIC.length + Byte.BYTES /* MARKER_FRAME */;
- } else {
- frameStart = frameFileFooter.getFrameEndPosition(frameNumber - 1) + Byte.BYTES /* MARKER_FRAME */;
+ // Map buffer to include marker byte.
+ if (buffer == null || markerPosition < bufferOffset || dataEnd > bufferOffset + buffer.getCapacity()) {
+ remapBuffer(markerPosition);
}
- if (buffer == null || frameStart < bufferOffset || frameEnd > bufferOffset + buffer.getCapacity()) {
- remapBuffer(frameStart);
+ if (markerPosition < bufferOffset || dataEnd > bufferOffset + buffer.getCapacity()) {
+ // Still out of bounds after remapping successfully: must mean entry was too large to fit in maxMmapSize.
+ throw new ISE("Entry [%,d] too large (max size = %,d bytes)", batchNumber, maxMmapSize);
}
- if (frameStart < bufferOffset || frameEnd > bufferOffset + buffer.getCapacity()) {
- // Still out of bounds after remapping successfully: must mean frame was too large to fit in maxMmapSize.
- throw new ISE("Frame [%,d] too large (max size = %,d bytes)", frameNumber, maxMmapSize);
- }
+ // Read marker byte.
+ final byte marker = buffer.getByte(markerPosition - bufferOffset);
+
+ // Read the data.
+ final byte[] decompressedData = FrameCompression.decompress(
+ buffer,
+ dataStart - bufferOffset,
+ dataEnd - dataStart
+ );
- // Decompression is safe even on corrupt data: it validates position, length, checksum.
- return Frame.decompress(buffer, frameStart - bufferOffset, frameEnd - frameStart);
+ switch (marker) {
+ case FrameFileWriter.MARKER_FRAME:
+ final Frame frame = Frame.wrap(decompressedData);
+ return frame.asRAC();
+
+ case FrameFileWriter.MARKER_RAC:
+ if (deserializer == null) {
+ throw DruidException.defensive("Cannot read RAC, no WireTransferable deserializer");
+ }
+ return deserializer.deserialize(ByteBufferUtils.wrapLE(decompressedData));
+
+ default:
+ throw new ISE("Invalid marker byte[%d] for batch[%,d]", marker, batchNumber);
+ }
}
/**
@@ -434,4 +469,16 @@ public class FrameFile implements Closeable
bufferCloser = null;
}
}
+
+ /**
+ * Returns the position in the file of the marker byte for the given frame number.
+ */
+ private long getMarkerPosition(final int frameNumber)
+ {
+ if (frameNumber == 0) {
+ return FrameFileWriter.MAGIC.length;
+ } else {
+ return frameFileFooter.getFrameEndPosition(frameNumber - 1);
+ }
+ }
}
diff --git a/processing/src/main/java/org/apache/druid/frame/file/FrameFileWriter.java b/processing/src/main/java/org/apache/druid/frame/file/FrameFileWriter.java
index f855641571e..09570848aa2 100644
--- a/processing/src/main/java/org/apache/druid/frame/file/FrameFileWriter.java
+++ b/processing/src/main/java/org/apache/druid/frame/file/FrameFileWriter.java
@@ -21,13 +21,18 @@ package org.apache.druid.frame.file;
import com.google.common.primitives.Ints;
import org.apache.datasketches.memory.WritableMemory;
+import org.apache.druid.error.DruidException;
import org.apache.druid.frame.Frame;
import org.apache.druid.frame.allocation.AppendableMemory;
import org.apache.druid.frame.allocation.HeapMemoryAllocator;
import org.apache.druid.frame.allocation.MemoryRange;
import org.apache.druid.frame.channel.ByteTracker;
+import org.apache.druid.frame.wire.FrameCompression;
import org.apache.druid.io.Channels;
import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import javax.annotation.Nullable;
import java.io.Closeable;
@@ -44,6 +49,7 @@ public class FrameFileWriter implements Closeable
public static final byte[] MAGIC = {(byte) 0xff, 0x01};
public static final byte MARKER_FRAME = (byte) 0x01;
public static final byte MARKER_NO_MORE_FRAMES = (byte) 0x02;
+ public static final byte MARKER_RAC = (byte) 0x03;
public static final int TRAILER_LENGTH = Integer.BYTES * 4;
public static final int CHECKSUM_SEED = 0;
public static final int NO_PARTITION = -1;
@@ -52,6 +58,7 @@ public class FrameFileWriter implements Closeable
private final AppendableMemory tableOfContents;
private final AppendableMemory partitions;
private final ByteTracker byteTracker;
+ private final WireTransferableContext wtContext;
private long bytesWritten = 0;
private long trackedBytes = 0;
private int numFrames = 0;
@@ -64,7 +71,8 @@ public class FrameFileWriter implements Closeable
@Nullable final ByteBuffer compressionBuffer,
final AppendableMemory tableOfContents,
final AppendableMemory partitions,
- ByteTracker byteTracker
+ final ByteTracker byteTracker,
+ final WireTransferableContext wtContext
)
{
this.channel = channel;
@@ -72,20 +80,24 @@ public class FrameFileWriter implements Closeable
this.tableOfContents = tableOfContents;
this.partitions = partitions;
this.byteTracker = byteTracker;
+ this.wtContext = wtContext;
}
/**
* Opens a writer for a particular channel.
- * @param channel destination channel
- * @param compressionBuffer result of {@link Frame#compressionBufferSize} for the largest possible frame size that
- * will be written to this file, or null to allocate buffers dynamically.
- * Providing an explicit buffer here, if possible, improves performance.
- * @param byteTracker tracker to limit the number of bytes that can be written to the frame file
+ *
+ * @param channel destination channel
+ * @param compressionBuffer result of {@link Frame#compressionBufferSize} for the largest possible frame size
+ * that will be written to this file, or null to allocate buffers dynamically.
+ * Providing an explicit buffer here, if possible, improves performance.
+ * @param byteTracker tracker to limit the number of bytes that can be written to the frame file
+ * @param wireTransferableContext context for wire transfer serialization
*/
public static FrameFileWriter open(
final WritableByteChannel channel,
@Nullable final ByteBuffer compressionBuffer,
- ByteTracker byteTracker
+ final ByteTracker byteTracker,
+ final WireTransferableContext wireTransferableContext
)
{
// Unlimited allocator is for convenience. Only a few bytes per frame will be allocated.
@@ -95,18 +107,70 @@ public class FrameFileWriter implements Closeable
compressionBuffer,
AppendableMemory.create(allocator),
AppendableMemory.create(allocator),
- byteTracker
+ byteTracker,
+ wireTransferableContext
);
}
/**
- * Write a frame.
+ * Write a batch of data to the file. If legacy frame serialization is enabled and the RowsAndColumns can be
+ * converted to a Frame, it will be written as raw frame bytes with {@link #MARKER_FRAME}. Otherwise, it will be
+ * written using {@link WireTransferable} with {@link #MARKER_RAC}.
*
- * @param frame the frame
+ * @param rac the RowsAndColumns to write
* @param partition partition number for a partitioned frame file, or {@link #NO_PARTITION} for an unpartitioned file.
* Must be monotonically increasing.
*/
- public void writeFrame(final Frame frame, final int partition) throws IOException
+ public void write(final RowsAndColumns rac, final int partition) throws IOException
+ {
+ if (wtContext.useLegacyFrameSerialization()) {
+ // Check if we should write as a Frame using the legacy format.
+ final Frame frame = rac.as(Frame.class);
+ if (frame != null) {
+ writeFrameLegacy(frame, partition);
+ return;
+ }
+ }
+
+ // Write using WireTransferable.
+ final WireTransferable wireTransferable = rac.as(WireTransferable.class);
+ if (wireTransferable == null) {
+ throw DruidException.defensive("RAC[%s] is not WireTransferable", rac.getClass().getName());
+ }
+
+ final WireTransferable.ByteArrayOffsetAndLen serializedBytes = wtContext.serializedBytes(wireTransferable);
+
+ // Compress and write the serialized bytes.
+ final ByteBuffer compressedBuffer = FrameCompression.compress(
+ serializedBytes.asByteBuffer(),
+ getCompressionBuffer(serializedBytes.getLength())
+ );
+
+ writeCompressedEntry(MARKER_RAC, compressedBuffer, partition);
+ }
+
+ /**
+ * Internal method to write a frame using the legacy format with {@link #MARKER_FRAME}.
+ */
+ private void writeFrameLegacy(final Frame frame, final int partition) throws IOException
+ {
+ final ByteBuffer compressedBuffer = FrameCompression.compress(
+ frame.readableMemory(),
+ frame.numBytes(),
+ getCompressionBuffer(frame.numBytes())
+ );
+
+ writeCompressedEntry(MARKER_FRAME, compressedBuffer, partition);
+ }
+
+ /**
+ * Common helper method to write a compressed entry (either MARKER_FRAME or MARKER_RAC) to the file.
+ */
+ private void writeCompressedEntry(
+ final byte marker,
+ final ByteBuffer compressedBuffer,
+ final int partition
+ ) throws IOException
{
if (numFrames == Integer.MAX_VALUE) {
throw new ISE("Too many frames");
@@ -127,19 +191,35 @@ public class FrameFileWriter implements Closeable
writeMagicIfNeeded();
+ // Write marker.
byteTracker.reserve(1);
trackedBytes++;
- Channels.writeFully(channel, ByteBuffer.wrap(new byte[]{MARKER_FRAME}));
+ Channels.writeFully(channel, ByteBuffer.wrap(new byte[]{marker}));
bytesWritten++;
- long frameWrittenBytes = frame.writeTo(channel, true, getCompressionBuffer(frame.numBytes()), byteTracker);
- bytesWritten += frameWrittenBytes;
- trackedBytes += frameWrittenBytes;
+
+ // Write compressed data.
+ final int compressedBytes = compressedBuffer.remaining();
+ byteTracker.reserve(compressedBytes);
+ trackedBytes += compressedBytes;
+ Channels.writeFully(channel, compressedBuffer);
+ bytesWritten += compressedBytes;
// Write *end* of frame to tableOfContents.
final MemoryRange<WritableMemory> tocCursor = tableOfContents.cursor();
tocCursor.memory().putLong(tocCursor.start(), bytesWritten);
tableOfContents.advanceCursor(Long.BYTES);
+ // Update partitions if needed.
+ updatePartitionsIfNeeded(partition);
+
+ numFrames++;
+ }
+
+ /**
+ * Updates partition tracking after writing a frame.
+ */
+ private void updatePartitionsIfNeeded(final int partition)
+ {
if (usePartitions) {
// Write new partition if needed.
int highestPartitionWritten = Ints.checkedCast(partitions.size() / Integer.BYTES) - 1;
@@ -162,8 +242,6 @@ public class FrameFileWriter implements Closeable
partitions.advanceCursor(Integer.BYTES);
}
}
-
- numFrames++;
}
/**
diff --git a/processing/src/main/java/org/apache/druid/frame/processor/FileOutputChannelFactory.java b/processing/src/main/java/org/apache/druid/frame/processor/FileOutputChannelFactory.java
index 430f8929684..6d363504fa7 100644
--- a/processing/src/main/java/org/apache/druid/frame/processor/FileOutputChannelFactory.java
+++ b/processing/src/main/java/org/apache/druid/frame/processor/FileOutputChannelFactory.java
@@ -32,6 +32,8 @@ import org.apache.druid.frame.file.FrameFile;
import org.apache.druid.frame.file.FrameFileWriter;
import org.apache.druid.java.util.common.FileUtils;
import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import javax.annotation.Nullable;
import java.io.File;
@@ -51,16 +53,19 @@ public class FileOutputChannelFactory implements OutputChannelFactory
private final File fileChannelsDirectory;
private final int frameSize;
private final ByteTracker byteTracker;
+ private final WireTransferableContext wireTransferableContext;
public FileOutputChannelFactory(
final File fileChannelsDirectory,
final int frameSize,
- @Nullable final ByteTracker byteTracker
+ @Nullable final ByteTracker byteTracker,
+ final WireTransferableContext wireTransferableContext
)
{
this.fileChannelsDirectory = fileChannelsDirectory;
this.frameSize = frameSize;
this.byteTracker = byteTracker == null ? ByteTracker.unboundedTracker() : byteTracker;
+ this.wireTransferableContext = wireTransferableContext;
}
@Override
@@ -80,15 +85,17 @@ public class FileOutputChannelFactory implements OutputChannelFactory
StandardOpenOption.WRITE
),
ByteBuffer.allocate(Frame.compressionBufferSize(frameSize)),
- byteTracker
+ byteTracker,
+ wireTransferableContext
)
);
+ final WireTransferable.ConcreteDeserializer deserializer = wireTransferableContext.concreteDeserializer();
final Supplier<ReadableFrameChannel> readableChannelSupplier = Suppliers.memoize(
() -> {
try {
final FrameFile frameFile = FrameFile.open(file, byteTracker, FrameFile.Flag.DELETE_ON_CLOSE);
- return new ReadableFileFrameChannel(frameFile);
+ return new ReadableFileFrameChannel(frameFile, deserializer);
}
catch (IOException e) {
throw new RuntimeException(e);
@@ -118,9 +125,11 @@ public class FileOutputChannelFactory implements OutputChannelFactory
StandardOpenOption.WRITE
),
ByteBuffer.allocate(Frame.compressionBufferSize(frameSize)),
- byteTracker
+ byteTracker,
+ wireTransferableContext
)
);
+ final WireTransferable.ConcreteDeserializer deserializer = wireTransferableContext.concreteDeserializer();
Supplier<FrameFile> frameFileSupplier = Suppliers.memoize(
() -> {
try {
@@ -143,7 +152,8 @@ public class FileOutputChannelFactory implements OutputChannelFactory
return new ReadableFileFrameChannel(
fileHandle,
fileHandle.getPartitionStartFrame(partitionNumber),
- fileHandle.getPartitionStartFrame(partitionNumber + 1)
+ fileHandle.getPartitionStartFrame(partitionNumber + 1),
+ deserializer
);
}
diff --git a/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelBatcher.java b/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelBatcher.java
index f62dd78fd90..e09853c803f 100644
--- a/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelBatcher.java
+++ b/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelBatcher.java
@@ -97,7 +97,7 @@ public class FrameChannelBatcher implements FrameProcessor<Pair<List<Frame>, Int
if (readableInputs.contains(channelNumber) && channelsToRead.contains(channelNumber)) {
final ReadableFrameChannel channel = channels.get(channelNumber);
if (channel.canRead()) {
- out.add(channel.read());
+ out.add(channel.readFrame());
} else if (channel.isFinished()) {
channelsToRead.remove(channelNumber);
}
diff --git a/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelHashPartitioner.java b/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelHashPartitioner.java
index b3a47c5817e..033e8f921ea 100644
--- a/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelHashPartitioner.java
+++ b/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelHashPartitioner.java
@@ -195,7 +195,7 @@ public class FrameChannelHashPartitioner implements FrameProcessor<Long>
if (!channel.isFinished()) {
// Need row-based frame so we can hash memory directly.
- final Frame frame = channel.read().ensureRowBased();
+ final Frame frame = channel.readFrame().ensureRowBased();
final HashPartitionVirtualColumn hashPartitionVirtualColumn =
new HashPartitionVirtualColumn(
diff --git a/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelMerger.java b/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelMerger.java
index c879aa8ef7b..da52cbf2096 100644
--- a/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelMerger.java
+++ b/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelMerger.java
@@ -25,7 +25,6 @@ import it.unimi.dsi.fastutil.ints.IntSet;
import it.unimi.dsi.fastutil.ints.IntSets;
import org.apache.druid.frame.Frame;
import org.apache.druid.frame.FrameType;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.key.ClusterByPartitions;
@@ -87,7 +86,7 @@ public class FrameChannelMerger implements FrameProcessor<Long>
* @param frameWriterFactory writer for frames
* @param sortKey sort key for input and output frames
* @param partitions partitions for output frames. If non-null, output frames are written with
- * {@link FrameWithPartition#partition()} set according to this parameter
+ * partition numbers set according to this parameter
* @param rowLimit maximum number of rows to write to the output channel
*/
public FrameChannelMerger(
@@ -189,9 +188,9 @@ public class FrameChannelMerger implements FrameProcessor<Long>
}
// Generate one output frame and stop for now.
- outputChannel.write(nextFrame());
+ writeNextFrame();
- // Check finished() after nextFrame().
+ // Check finished() after writeNextFrame().
if (finished()) {
return ReturnOrAwait.returnObject(rowsOutput);
} else {
@@ -199,7 +198,7 @@ public class FrameChannelMerger implements FrameProcessor<Long>
}
}
- private FrameWithPartition nextFrame()
+ private void writeNextFrame() throws IOException
{
if (finished()) {
throw new NoSuchElementException();
@@ -227,7 +226,7 @@ public class FrameChannelMerger implements FrameProcessor<Long>
// Fall through: keep reading into the new partition.
mergedFramePartition = currentPartition;
} else {
- // Return current frame.
+ // Write current frame.
break;
}
}
@@ -240,7 +239,7 @@ public class FrameChannelMerger implements FrameProcessor<Long>
throw new FrameRowTooLargeException(frameWriterFactory.allocatorCapacity());
}
- // Frame is full. Return the current frame.
+ // Frame is full. Write the current frame.
break;
}
@@ -262,11 +261,11 @@ public class FrameChannelMerger implements FrameProcessor<Long>
if (channel.canRead()) {
// Read next frame from this channel.
- final Frame frame = channel.read();
+ final Frame frame = channel.readFrame();
final FramePlus framePlus = makeFramePlus(frame, frameReader);
if (framePlus.isDone()) {
// Nothing to read in this frame. Not finished; we can't continue.
- // Finish up the current frame and return it.
+ // Finish up the current frame and write it.
break;
} else {
currentFrames[currentChannel] = framePlus;
@@ -275,7 +274,7 @@ public class FrameChannelMerger implements FrameProcessor<Long>
// Done reading this channel. Fall through and continue with other channels.
remainingChannels.remove(currentChannel);
} else {
- // Nothing available, not finished; we can't continue. Finish up the current frame and return it.
+ // Nothing available, not finished; we can't continue. Finish up the current frame and write it.
break;
}
}
@@ -283,7 +282,7 @@ public class FrameChannelMerger implements FrameProcessor<Long>
}
final Frame nextFrame = Frame.wrap(mergedFrameWriter.toByteArray());
- return new FrameWithPartition(nextFrame, mergedFramePartition);
+ outputChannel.write(nextFrame, mergedFramePartition);
}
}
@@ -314,7 +313,7 @@ public class FrameChannelMerger implements FrameProcessor<Long>
final ReadableFrameChannel channel = inputChannels.get(i);
if (channel.canRead()) {
- final Frame frame = channel.read();
+ final Frame frame = channel.readFrame();
final FramePlus framePlus = makeFramePlus(frame, frameReader);
if (framePlus.isDone()) {
await.add(i);
diff --git a/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelMixer.java b/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelMixer.java
index 7f1c48f7ffa..d7db6d81f93 100644
--- a/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelMixer.java
+++ b/processing/src/main/java/org/apache/druid/frame/processor/FrameChannelMixer.java
@@ -21,9 +21,9 @@ package org.apache.druid.frame.processor;
import it.unimi.dsi.fastutil.ints.IntAVLTreeSet;
import it.unimi.dsi.fastutil.ints.IntSet;
-import org.apache.druid.frame.Frame;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import java.io.IOException;
import java.util.Collections;
@@ -85,9 +85,9 @@ public class FrameChannelMixer implements FrameProcessor<Long>
final ReadableFrameChannel channel = inputChannels.get(channelNumber);
if (!channel.isFinished()) {
- final Frame frame = channel.read();
- outputChannel.write(frame);
- rowsRead += frame.numRows();
+ final RowsAndColumns rac = channel.read();
+ outputChannel.write(rac);
+ rowsRead += rac.numRows();
}
}
diff --git a/processing/src/main/java/org/apache/druid/frame/processor/OutputChannel.java b/processing/src/main/java/org/apache/druid/frame/processor/OutputChannel.java
index 09fae34bb07..73fdc2c03bb 100644
--- a/processing/src/main/java/org/apache/druid/frame/processor/OutputChannel.java
+++ b/processing/src/main/java/org/apache/druid/frame/processor/OutputChannel.java
@@ -23,7 +23,6 @@ import com.google.common.base.Preconditions;
import com.google.common.base.Suppliers;
import org.apache.druid.error.DruidException;
import org.apache.druid.frame.allocation.MemoryAllocator;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.ReadableNilFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
@@ -63,7 +62,7 @@ public class OutputChannel
this.readableChannelSupplier = readableChannelSupplier;
this.partitionNumber = partitionNumber;
- if (partitionNumber < 0 && partitionNumber != FrameWithPartition.NO_PARTITION) {
+ if (partitionNumber < 0 && partitionNumber != WritableFrameChannel.NO_PARTITION) {
throw new IAE("Invalid partition number [%d]", partitionNumber);
}
}
@@ -75,7 +74,7 @@ public class OutputChannel
* @param frameMemoryAllocator memory allocator for producer to use while writing frames to the channel
* @param readableChannelSupplier readable channel for consumer. May be called multiple times, so you should wrap this
* in {@link Suppliers#memoize} if needed.
- * @param partitionNumber partition number, if any; may be {@link FrameWithPartition#NO_PARTITION} if unknown
+ * @param partitionNumber partition number, if any; may be {@link WritableFrameChannel#NO_PARTITION} if unknown
*/
public static OutputChannel pair(
final WritableFrameChannel writableChannel,
@@ -98,7 +97,7 @@ public class OutputChannel
* @param writableChannel writable channel for producer
* @param frameMemoryAllocator memory allocator for producer to use while writing frames to the channel
* @param readableChannel readable channel for consumer
- * @param partitionNumber partition number, if any; may be {@link FrameWithPartition#NO_PARTITION} if unknown
+ * @param partitionNumber partition number, if any; may be {@link WritableFrameChannel#NO_PARTITION} if unknown
*/
public static OutputChannel immediatelyReadablePair(
final WritableFrameChannel writableChannel,
@@ -119,7 +118,7 @@ public class OutputChannel
* Creates a read-only output channel.
*
* @param readableChannel readable channel for consumer.
- * @param partitionNumber partition number, if any; may be {@link FrameWithPartition#NO_PARTITION} if unknown
+ * @param partitionNumber partition number, if any; may be {@link WritableFrameChannel#NO_PARTITION} if unknown
*/
public static OutputChannel readOnly(
final ReadableFrameChannel readableChannel,
@@ -134,7 +133,7 @@ public class OutputChannel
*
* @param readableChannelSupplier readable channel for consumer. May be called multiple times, so you should wrap this
* in {@link Suppliers#memoize} if needed.
- * @param partitionNumber partition number, if any; may be {@link FrameWithPartition#NO_PARTITION} if unknown
+ * @param partitionNumber partition number, if any; may be {@link WritableFrameChannel#NO_PARTITION} if unknown
*/
public static OutputChannel readOnly(
final Supplier<ReadableFrameChannel> readableChannelSupplier,
diff --git a/processing/src/main/java/org/apache/druid/frame/processor/SuperSorter.java b/processing/src/main/java/org/apache/druid/frame/processor/SuperSorter.java
index bcf585fdda8..3c13b39c8b1 100644
--- a/processing/src/main/java/org/apache/druid/frame/processor/SuperSorter.java
+++ b/processing/src/main/java/org/apache/druid/frame/processor/SuperSorter.java
@@ -42,7 +42,6 @@ import org.apache.druid.frame.FrameType;
import org.apache.druid.frame.allocation.MemoryAllocatorFactory;
import org.apache.druid.frame.allocation.SingleMemoryAllocatorFactory;
import org.apache.druid.frame.channel.BlockingQueueFrameChannel;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.PartitionedReadableFrameChannel;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
@@ -499,7 +498,7 @@ public class SuperSorter
final List<ReadableFrameChannel> in = new ArrayList<>();
for (final Frame frame : inputBuffer) {
- in.add(singleReadableFrameChannel(new FrameWithPartition(frame, FrameWithPartition.NO_PARTITION)));
+ in.add(singleReadableFrameChannel(frame));
}
runMerger(0, ultimateMergersRunSoFar++, in, Collections.emptyList());
@@ -537,7 +536,7 @@ public class SuperSorter
break;
}
- in.add(singleReadableFrameChannel(new FrameWithPartition(frame, FrameWithPartition.NO_PARTITION)));
+ in.add(singleReadableFrameChannel(frame));
}
runMerger(0, levelZeroMergersRunSoFar++, in, ImmutableList.of());
@@ -1055,7 +1054,7 @@ public class SuperSorter
}
}
- private static ReadableFrameChannel singleReadableFrameChannel(final FrameWithPartition frame)
+ private static ReadableFrameChannel singleReadableFrameChannel(final Frame frame)
{
try {
final BlockingQueueFrameChannel channel = BlockingQueueFrameChannel.minimal();
diff --git a/processing/src/main/java/org/apache/druid/frame/wire/FrameCompression.java b/processing/src/main/java/org/apache/druid/frame/wire/FrameCompression.java
new file mode 100644
index 00000000000..0ef814ac5e9
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/frame/wire/FrameCompression.java
@@ -0,0 +1,221 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.frame.wire;
+
+import com.google.common.primitives.Ints;
+import net.jpountz.lz4.LZ4Compressor;
+import net.jpountz.lz4.LZ4Factory;
+import net.jpountz.lz4.LZ4SafeDecompressor;
+import org.apache.datasketches.memory.Memory;
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.segment.data.CompressionStrategy;
+
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+
+/**
+ * Utility class for compressing and decompressing data using LZ4 compression with a specific envelope format.
+ *
+ * Compressed format:
+ * - 1 byte: compression type: {@link CompressionStrategy#getId()}. Currently, only LZ4 is supported.
+ * - 8 bytes: compressed data length, little-endian long
+ * - 8 bytes: uncompressed data length, little-endian long
+ * - NNN bytes: LZ4-compressed data
+ * - 8 bytes: 64-bit xxhash checksum of prior content, including 17-byte header and compressed data, little-endian long
+ */
+public class FrameCompression
+{
+ public static final int COMPRESSED_DATA_HEADER_SIZE = Byte.BYTES + Long.BYTES * 2;
+ public static final int COMPRESSED_DATA_TRAILER_SIZE = Long.BYTES; // Checksum
+ public static final int COMPRESSED_DATA_ENVELOPE_SIZE = COMPRESSED_DATA_HEADER_SIZE + COMPRESSED_DATA_TRAILER_SIZE;
+
+ private static final LZ4Compressor LZ4_COMPRESSOR = LZ4Factory.fastestInstance().fastCompressor();
+ private static final LZ4SafeDecompressor LZ4_DECOMPRESSOR = LZ4Factory.fastestInstance().safeDecompressor();
+
+ private static final int CHECKSUM_SEED = 0;
+
+ private FrameCompression()
+ {
+ // No instantiation.
+ }
+
+ /**
+ * Minimum size of compression buffer required by {@link #compress}.
+ */
+ public static int compressionBufferSize(final long dataBytes)
+ {
+ return COMPRESSED_DATA_ENVELOPE_SIZE + LZ4_COMPRESSOR.maxCompressedLength(Ints.checkedCast(dataBytes));
+ }
+
+ /**
+ * Compresses data from Memory into the provided compression buffer.
+ *
+ * @param data the memory region to compress
+ * @param dataLength the number of bytes to compress from memory
+ * @param compressionBuffer buffer to hold compressed output; must be at least {@link #compressionBufferSize} bytes
+ * @return the compression buffer with position 0 and limit set to the number of bytes written
+ */
+ public static ByteBuffer compress(final Memory data, final long dataLength, final ByteBuffer compressionBuffer)
+ {
+ final ByteBuffer dataBuffer;
+
+ if (data.hasByteBuffer()) {
+ // Optimized path when Memory is backed by ByteBuffer.
+ final ByteBuffer underlying = data.getByteBuffer();
+ final int offset = Ints.checkedCast(data.getRegionOffset());
+ dataBuffer = underlying.duplicate();
+ dataBuffer.position(offset).limit(offset + Ints.checkedCast(dataLength));
+ } else {
+ // Copy to byte array first, then compress.
+ final byte[] dataBytes = new byte[Ints.checkedCast(dataLength)];
+ data.getByteArray(0, dataBytes, 0, Ints.checkedCast(dataLength));
+ dataBuffer = ByteBuffer.wrap(dataBytes);
+ }
+
+ return compress(dataBuffer, compressionBuffer);
+ }
+
+ /**
+ * Compresses data into the provided compression buffer.
+ *
+ * @param data the data to compress (position and limit are used)
+ * @param compressionBuffer buffer to hold compressed output; must be at least {@link #compressionBufferSize} bytes
+ * @return the compression buffer with position 0 and limit set to the number of bytes written
+ */
+ public static ByteBuffer compress(final ByteBuffer data, final ByteBuffer compressionBuffer)
+ {
+ final int dataLength = data.remaining();
+ final int requiredSize = compressionBufferSize(dataLength);
+
+ if (compressionBuffer.capacity() < requiredSize) {
+ throw new ISE(
+ "Compression buffer too small: expected [%,d] bytes but got [%,d] bytes",
+ requiredSize,
+ compressionBuffer.capacity()
+ );
+ }
+
+ final int compressedDataLength = LZ4_COMPRESSOR.compress(
+ data,
+ data.position(),
+ dataLength,
+ compressionBuffer,
+ COMPRESSED_DATA_HEADER_SIZE,
+ compressionBuffer.capacity() - COMPRESSED_DATA_ENVELOPE_SIZE
+ );
+
+ compressionBuffer.order(ByteOrder.LITTLE_ENDIAN)
+ .limit(COMPRESSED_DATA_ENVELOPE_SIZE + compressedDataLength)
+ .position(0);
+
+ compressionBuffer.put(0, CompressionStrategy.LZ4.getId())
+ .putLong(Byte.BYTES, compressedDataLength)
+ .putLong(Byte.BYTES + Long.BYTES, dataLength);
+
+ final long checksum = Memory.wrap(compressionBuffer)
+ .xxHash64(0, COMPRESSED_DATA_HEADER_SIZE + compressedDataLength, CHECKSUM_SEED);
+
+ compressionBuffer.putLong(COMPRESSED_DATA_HEADER_SIZE + compressedDataLength, checksum);
+
+ return compressionBuffer;
+ }
+
+ /**
+ * Decompresses the provided memory and returns a byte array with the decompressed data. This operation is
+ * safe even on corrupt data: it validates position, length, and checksum prior to decompressing.
+ *
+ * This operation allocates memory on-heap to store the decompressed data.
+ *
+ * @param memory the memory containing compressed data
+ * @param position the position in memory where compressed data starts
+ * @param length the length of compressed data
+ * @return a byte array containing the decompressed data
+ */
+ public static byte[] decompress(final Memory memory, final long position, final long length)
+ {
+ if (memory.getCapacity() < position + length) {
+ throw new ISE("Provided position, length is out of bounds");
+ }
+
+ if (length < COMPRESSED_DATA_ENVELOPE_SIZE) {
+ throw new ISE("Region too short");
+ }
+
+ // Verify checksum.
+ final long expectedChecksum = memory.getLong(position + length - COMPRESSED_DATA_TRAILER_SIZE);
+ final long actualChecksum = memory.xxHash64(position, length - COMPRESSED_DATA_TRAILER_SIZE, CHECKSUM_SEED);
+
+ if (expectedChecksum != actualChecksum) {
+ throw new ISE("Checksum mismatch");
+ }
+
+ final byte compressionTypeId = memory.getByte(position);
+ final CompressionStrategy compressionStrategy = CompressionStrategy.forId(compressionTypeId);
+
+ if (compressionStrategy != CompressionStrategy.LZ4) {
+ throw new ISE("Unsupported compression strategy [%s]", compressionStrategy);
+ }
+
+ final int compressedDataLength = Ints.checkedCast(memory.getLong(position + Byte.BYTES));
+ final int uncompressedDataLength = Ints.checkedCast(memory.getLong(position + Byte.BYTES + Long.BYTES));
+ final int compressedDataLengthFromRegionLength = Ints.checkedCast(length - COMPRESSED_DATA_ENVELOPE_SIZE);
+ final long dataStart = position + COMPRESSED_DATA_HEADER_SIZE;
+
+ // Verify length.
+ if (compressedDataLength != compressedDataLengthFromRegionLength) {
+ throw new ISE(
+ "Compressed sizes disagree: [%d] (embedded) vs [%d] (region length)",
+ compressedDataLength,
+ compressedDataLengthFromRegionLength
+ );
+ }
+
+ if (memory.hasByteBuffer()) {
+ // Decompress directly out of the ByteBuffer.
+ final ByteBuffer srcBuffer = memory.getByteBuffer();
+ final ByteBuffer dstBuffer = ByteBuffer.allocate(uncompressedDataLength);
+ final int numBytesDecompressed =
+ LZ4_DECOMPRESSOR.decompress(
+ srcBuffer,
+ Ints.checkedCast(memory.getRegionOffset() + dataStart),
+ compressedDataLength,
+ dstBuffer,
+ 0,
+ uncompressedDataLength
+ );
+
+ // Sanity check.
+ if (numBytesDecompressed != uncompressedDataLength) {
+ throw new ISE(
+ "Expected to decompress [%d] bytes but got [%d] bytes",
+ uncompressedDataLength,
+ numBytesDecompressed
+ );
+ }
+
+ return dstBuffer.array();
+ } else {
+ // Copy first, then decompress.
+ final byte[] compressedData = new byte[compressedDataLength];
+ memory.getByteArray(dataStart, compressedData, 0, compressedDataLength);
+ return LZ4_DECOMPRESSOR.decompress(compressedData, uncompressedDataLength);
+ }
+ }
+}
diff --git a/processing/src/main/java/org/apache/druid/frame/wire/FrameWireTransferable.java b/processing/src/main/java/org/apache/druid/frame/wire/FrameWireTransferable.java
new file mode 100644
index 00000000000..d2e968107fa
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/frame/wire/FrameWireTransferable.java
@@ -0,0 +1,157 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.frame.wire;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.primitives.Ints;
+import org.apache.datasketches.memory.Memory;
+import org.apache.druid.frame.Frame;
+import org.apache.druid.java.util.common.RE;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.ColumnBasedFrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.apache.druid.segment.column.RowSignature;
+
+import javax.annotation.Nullable;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+
+/**
+ * A {@link WireTransferable} implementation for frames, used to serialize
+ * {@link RowBasedFrameRowsAndColumns} and {@link ColumnBasedFrameRowsAndColumns}.
+ *
+ * Signature-included format is:
+ * 1) 0x05 (length of type code)
+ * 2) "frame" (type code)
+ * 3) flags byte
+ * 4) signature length, if signature is present (little-endian int)
+ * 5) signature bytes, encoded with object mapper, if signature is present
+ * 6) frame bytes
+ */
+public class FrameWireTransferable implements WireTransferable
+{
+ public static final String TYPE = "frame";
+ public static final byte[] TYPE_BYTES = StringUtils.toUtf8(TYPE);
+ private static final byte HAS_SIGNATURE = 0x01;
+
+ private final Frame frame;
+ @Nullable
+ private final RowSignature signature;
+
+ public FrameWireTransferable(Frame frame, @Nullable RowSignature signature)
+ {
+ this.frame = frame;
+ this.signature = signature;
+ }
+
+ @Override
+ public ByteArrayOffsetAndLen serializedBytes(ObjectMapper mapper) throws IOException
+ {
+ final byte[] signatureBytes;
+
+ if (signature != null) {
+ signatureBytes = mapper.writeValueAsBytes(signature);
+ } else {
+ signatureBytes = null;
+ }
+
+ final int serializedLength = Ints.checkedCast(
+ 1 // type code length
+ + TYPE_BYTES.length // type code
+ + 1 // flags
+ + (signatureBytes != null ? Integer.BYTES + signatureBytes.length : 0) // signature
+ + frame.numBytes() // frame data
+ );
+
+ final ByteBuffer serializedBytes = ByteBuffer.allocate(serializedLength).order(ByteOrder.LITTLE_ENDIAN);
+
+ // Write type code length (1 byte) + type code
+ serializedBytes.put((byte) TYPE_BYTES.length);
+ serializedBytes.put(TYPE_BYTES);
+
+ // Write flags
+ serializedBytes.put(signature != null ? HAS_SIGNATURE : 0x00);
+
+ // Write signature length (4 bytes, little-endian) + signature, if present
+ if (signature != null) {
+ serializedBytes.putInt(signatureBytes.length);
+ serializedBytes.put(signatureBytes);
+ }
+
+ // Write frame bytes
+ frame.readableMemory()
+ .getByteArray(0, serializedBytes.array(), serializedBytes.position(), Ints.checkedCast(frame.numBytes()));
+
+ return new ByteArrayOffsetAndLen(serializedBytes.array(), 0, serializedLength);
+ }
+
+ private static RowsAndColumns deserializeFrame(ObjectMapper mapper, ByteBuffer theBytes)
+ {
+ // Skip the type code.
+ final int typeCodeLen = 0xFF & theBytes.get();
+ theBytes.position(theBytes.position() + typeCodeLen);
+
+ // Read flags.
+ final byte flags = theBytes.get();
+ final boolean hasSignature = (flags & HAS_SIGNATURE) != 0;
+
+ final RowSignature signature;
+ if (hasSignature) {
+ // Read signature.
+ final int signatureLen = theBytes.getInt();
+ final byte[] signatureBytes = new byte[signatureLen];
+ theBytes.get(signatureBytes);
+
+ try {
+ signature = mapper.readValue(signatureBytes, RowSignature.class);
+ }
+ catch (IOException e) {
+ throw new RE(e, "Failed to deserialize RowSignature");
+ }
+ } else {
+ signature = null;
+ }
+
+ // Read remaining bytes as a Frame.
+ final Frame frame = Frame.wrap(Memory.wrap(theBytes, ByteOrder.LITTLE_ENDIAN)
+ .region(theBytes.position(), theBytes.remaining()));
+
+ if (frame.type().isRowBased()) {
+ return new RowBasedFrameRowsAndColumns(frame, signature);
+ } else {
+ return new ColumnBasedFrameRowsAndColumns(frame, signature);
+ }
+ }
+
+ /**
+ * Deserializer for frames.
+ */
+ public static class Deserializer implements WireTransferable.Deserializer
+ {
+ @Override
+ public RowsAndColumns deserialize(ObjectMapper mapper, ByteBuffer theBytes)
+ {
+ return deserializeFrame(mapper, theBytes);
+ }
+ }
+}
diff --git a/processing/src/main/java/org/apache/druid/guice/DruidSecondaryModule.java b/processing/src/main/java/org/apache/druid/guice/DruidSecondaryModule.java
index be0ea786f58..3c8cc228d98 100644
--- a/processing/src/main/java/org/apache/druid/guice/DruidSecondaryModule.java
+++ b/processing/src/main/java/org/apache/druid/guice/DruidSecondaryModule.java
@@ -22,22 +22,48 @@ package org.apache.druid.guice;
import com.fasterxml.jackson.databind.AnnotationIntrospector;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.introspect.AnnotationIntrospectorPair;
+import com.fasterxml.jackson.databind.module.SimpleModule;
import com.google.inject.Binder;
import com.google.inject.Inject;
import com.google.inject.Injector;
import com.google.inject.Key;
import com.google.inject.Module;
import com.google.inject.Provides;
+import com.google.inject.TypeLiteral;
+import com.google.inject.multibindings.MapBinder;
+import org.apache.druid.frame.wire.FrameWireTransferable;
import org.apache.druid.guice.annotations.Json;
import org.apache.druid.guice.annotations.JsonNonNull;
import org.apache.druid.guice.annotations.Smile;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.FrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.apache.druid.query.rowsandcols.serde.RowsAndColumnsDeserializer;
+import org.apache.druid.query.rowsandcols.serde.RowsAndColumnsSerializer;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import javax.validation.Validator;
+import java.nio.ByteBuffer;
+import java.util.Map;
import java.util.Properties;
@LazySingleton
public class DruidSecondaryModule implements Module
{
+ public static final String PROPERTY_RAC_LEGACY_FRAME_SERIALIZATION = "druid.serde.rac.useLegacyFrameSerialization";
+
+ public static MapBinder<ByteBuffer, WireTransferable.Deserializer> getWireTransferableDeserializerBinder(
+ Binder binder
+ )
+ {
+ return MapBinder.newMapBinder(
+ binder,
+ TypeLiteral.get(ByteBuffer.class),
+ TypeLiteral.get(WireTransferable.Deserializer.class)
+ );
+ }
+
private final Properties properties;
private final ObjectMapper jsonMapper;
private final ObjectMapper jsonMapperOnlyNonNullValueSerialization;
@@ -68,33 +94,108 @@ public class DruidSecondaryModule implements Module
binder.bind(ObjectMapper.class).to(Key.get(ObjectMapper.class, Json.class));
binder.bind(Validator.class).toInstance(validator);
binder.bind(JsonConfigurator.class);
+
+ // Register frame deserializers.
+ getWireTransferableDeserializerBinder(binder)
+ .addBinding(StringUtils.toUtf8ByteBuffer(FrameWireTransferable.TYPE))
+ .to(FrameWireTransferable.Deserializer.class);
}
@Provides @LazySingleton @Json
- public ObjectMapper getJsonMapper(final Injector injector)
+ public ObjectMapper getJsonMapper(
+ final Injector injector,
+ final Map<ByteBuffer, WireTransferable.Deserializer> wtDeserializers
+ )
{
- setupJackson(injector, jsonMapper);
+ setupJackson(injector, jsonMapper, wtDeserializers, isUseLegacyFrameSerialization());
return jsonMapper;
}
@Provides @LazySingleton @JsonNonNull
- public ObjectMapper getJsonMapperOnlyNonNullValueSerialization(final Injector injector)
+ public ObjectMapper getJsonMapperOnlyNonNullValueSerialization(
+ final Injector injector,
+ final Map<ByteBuffer, WireTransferable.Deserializer> wtDeserializers
+ )
{
- setupJackson(injector, jsonMapperOnlyNonNullValueSerialization);
+ setupJackson(injector, jsonMapperOnlyNonNullValueSerialization, wtDeserializers, isUseLegacyFrameSerialization());
return jsonMapperOnlyNonNullValueSerialization;
}
@Provides @LazySingleton @Smile
- public ObjectMapper getSmileMapper(Injector injector)
+ public ObjectMapper getSmileMapper(
+ final Injector injector,
+ final Map<ByteBuffer, WireTransferable.Deserializer> wtDeserializers
+ )
{
- setupJackson(injector, smileMapper);
+ setupJackson(injector, smileMapper, wtDeserializers, isUseLegacyFrameSerialization());
return smileMapper;
}
- public static void setupJackson(Injector injector, final ObjectMapper mapper)
+ /**
+ * Whether to use the legacy serialization format (true) or {@link FrameWireTransferable} (false) when serializing
+ * {@link FrameRowsAndColumns} with Jackson. Deserialization supports both formats regardless of the setting of
+ * this property.
+ */
+ private boolean isUseLegacyFrameSerialization()
+ {
+ return Boolean.parseBoolean(
+ properties.getProperty(
+ PROPERTY_RAC_LEGACY_FRAME_SERIALIZATION,
+ String.valueOf(WireTransferableContext.DEFAULT_LEGACY_FRAME_SERIALIZATION)
+ )
+ );
+ }
+
+ @Provides @LazySingleton @Smile
+ public WireTransferable.ConcreteDeserializer getConcreteDeserializer(
+ @Smile final ObjectMapper smileMapper,
+ final Map<ByteBuffer, WireTransferable.Deserializer> wtDeserializers
+ )
+ {
+ return new WireTransferable.ConcreteDeserializer(smileMapper, wtDeserializers);
+ }
+
+ @Provides @LazySingleton
+ public WireTransferableContext getWireTransferableContext(
+ @Smile final ObjectMapper smileMapper,
+ @Smile final WireTransferable.ConcreteDeserializer concreteDeserializer
+ )
+ {
+ return new WireTransferableContext(smileMapper, concreteDeserializer, isUseLegacyFrameSerialization());
+ }
+
+ public static void setupJackson(
+ final Injector injector,
+ final ObjectMapper mapper,
+ final Map<ByteBuffer, WireTransferable.Deserializer> wtDeserializers,
+ final boolean useLegacyFrameSerialization
+ )
{
mapper.setInjectableValues(new GuiceInjectableValues(injector));
setupAnnotationIntrospector(mapper, new GuiceAnnotationIntrospector());
+
+ attachWireTransferables(mapper, wtDeserializers, useLegacyFrameSerialization);
+ }
+
+ /**
+ * Attach {@link RowsAndColumnsSerializer} and {@link RowsAndColumnsDeserializer} to an {@link ObjectMapper}.
+ */
+ public static void attachWireTransferables(
+ ObjectMapper mapper,
+ Map<ByteBuffer, WireTransferable.Deserializer> wtDeserializers,
+ boolean useLegacyFrameSerialization
+ )
+ {
+ final WireTransferableContext wtContext = new WireTransferableContext(
+ mapper,
+ new WireTransferable.ConcreteDeserializer(mapper, wtDeserializers),
+ useLegacyFrameSerialization
+ );
+ mapper.registerModule(
+ new SimpleModule("wireTransferableSerde")
+ .addDeserializer(RowsAndColumns.class, new RowsAndColumnsDeserializer(wtContext))
+ .addSerializer(RowsAndColumns.class, new RowsAndColumnsSerializer(wtContext))
+ );
}
public static void setupAnnotationIntrospector(
diff --git a/processing/src/main/java/org/apache/druid/jackson/DruidDefaultSerializersModule.java b/processing/src/main/java/org/apache/druid/jackson/DruidDefaultSerializersModule.java
index daa7ea81317..da78a0a7461 100644
--- a/processing/src/main/java/org/apache/druid/jackson/DruidDefaultSerializersModule.java
+++ b/processing/src/main/java/org/apache/druid/jackson/DruidDefaultSerializersModule.java
@@ -36,7 +36,6 @@ import org.apache.druid.query.FrameBasedInlineDataSource;
import org.apache.druid.query.FrameBasedInlineDataSourceSerializer;
import org.apache.druid.query.context.ResponseContext;
import org.apache.druid.query.context.ResponseContextDeserializer;
-import org.apache.druid.query.rowsandcols.RowsAndColumns;
import org.joda.time.DateTimeZone;
import java.io.IOException;
@@ -187,8 +186,5 @@ public class DruidDefaultSerializersModule extends SimpleModule
}
);
addDeserializer(ResponseContext.class, new ResponseContextDeserializer());
-
- addSerializer(RowsAndColumns.class, new RowsAndColumns.RowsAndColumnsSerializer());
- addDeserializer(RowsAndColumns.class, new RowsAndColumns.RowsAndColumnsDeserializer());
}
}
diff --git a/processing/src/main/java/org/apache/druid/java/util/common/ByteBufferUtils.java b/processing/src/main/java/org/apache/druid/java/util/common/ByteBufferUtils.java
index 696f0b629c9..c81c445e258 100644
--- a/processing/src/main/java/org/apache/druid/java/util/common/ByteBufferUtils.java
+++ b/processing/src/main/java/org/apache/druid/java/util/common/ByteBufferUtils.java
@@ -26,6 +26,7 @@ import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
import java.nio.MappedByteBuffer;
import java.util.Comparator;
import java.util.concurrent.atomic.AtomicBoolean;
@@ -212,4 +213,28 @@ public class ByteBufferUtils
return compareUtf8ByteBuffers(o1, o2);
}
}
+
+ /**
+ * Wraps a byte array in a little-endian ByteBuffer.
+ *
+ * @param array byte array to wrap
+ * @return ByteBuffer wrapping the array with little-endian byte order
+ */
+ public static ByteBuffer wrapLE(byte[] array)
+ {
+ return ByteBuffer.wrap(array).order(ByteOrder.LITTLE_ENDIAN);
+ }
+
+ /**
+ * Wraps a portion of a byte array in a little-endian ByteBuffer.
+ *
+ * @param array byte array to wrap
+ * @param offset offset within the array
+ * @param length length of the region to wrap
+ * @return ByteBuffer wrapping the specified portion with little-endian byte order
+ */
+ public static ByteBuffer wrapLE(byte[] array, int offset, int length)
+ {
+ return ByteBuffer.wrap(array, offset, length).order(ByteOrder.LITTLE_ENDIAN);
+ }
}
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/AppendableMapOfColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/AppendableMapOfColumns.java
index d83f56c7ba5..61f6855cd01 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/AppendableMapOfColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/AppendableMapOfColumns.java
@@ -79,4 +79,14 @@ public class AppendableMapOfColumns implements AppendableRowsAndColumns
}
return retVal;
}
+
+ @Override
+ @SuppressWarnings("unchecked")
+ public <T> T as(Class<T> clazz)
+ {
+ if (AppendableRowsAndColumns.class.equals(clazz)) {
+ return (T) this;
+ }
+ return null;
+ }
}
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumns.java
index 3f70f82a253..c6ced60849d 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/ConcatRowsAndColumns.java
@@ -141,6 +141,13 @@ public class ConcatRowsAndColumns implements RowsAndColumns
}
}
+ @Nullable
+ @Override
+ public <T> T as(Class<T> clazz)
+ {
+ return null;
+ }
+
private class ConcatedidColumn implements Column
{
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/CursorFactoryRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/CursorFactoryRowsAndColumns.java
index e636921baf7..8331676ea7d 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/CursorFactoryRowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/CursorFactoryRowsAndColumns.java
@@ -61,7 +61,7 @@ public class CursorFactoryRowsAndColumns implements CloseableShapeshifter, RowsA
if (CursorFactory.class == clazz) {
return (T) cursorFactory;
}
- return RowsAndColumns.super.as(clazz);
+ return null;
}
@Override
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/EmptyRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/EmptyRowsAndColumns.java
index 56647e0f568..dd0c7dab1cd 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/EmptyRowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/EmptyRowsAndColumns.java
@@ -21,6 +21,7 @@ package org.apache.druid.query.rowsandcols;
import org.apache.druid.query.rowsandcols.column.Column;
+import javax.annotation.Nullable;
import java.util.ArrayList;
import java.util.Collection;
@@ -43,4 +44,11 @@ public class EmptyRowsAndColumns implements RowsAndColumns
{
return null;
}
+
+ @Nullable
+ @Override
+ public <T> T as(Class<T> clazz)
+ {
+ return null;
+ }
}
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/LimitedRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/LimitedRowsAndColumns.java
index 8cfadecb4dd..30d1c4d717a 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/LimitedRowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/LimitedRowsAndColumns.java
@@ -23,6 +23,7 @@ import org.apache.druid.java.util.common.ISE;
import org.apache.druid.query.rowsandcols.column.Column;
import org.apache.druid.query.rowsandcols.column.LimitedColumn;
+import javax.annotation.Nullable;
import java.util.Collection;
public class LimitedRowsAndColumns implements RowsAndColumns
@@ -65,4 +66,11 @@ public class LimitedRowsAndColumns implements RowsAndColumns
return new LimitedColumn(column, start, end);
}
+
+ @Nullable
+ @Override
+ public <T> T as(Class<T> clazz)
+ {
+ return null;
+ }
}
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/MapOfColumnsRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/MapOfColumnsRowsAndColumns.java
index d6bc1026a98..29f092f6744 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/MapOfColumnsRowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/MapOfColumnsRowsAndColumns.java
@@ -164,7 +164,7 @@ public class MapOfColumnsRowsAndColumns implements RowsAndColumns
if (AppendableRowsAndColumns.class.equals(clazz)) {
return (T) new AppendableMapOfColumns(this);
}
- return RowsAndColumns.super.as(clazz);
+ return null;
}
public static class Builder
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/RearrangedRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/RearrangedRowsAndColumns.java
index e64f086edd7..f1793f8fd0e 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/RearrangedRowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/RearrangedRowsAndColumns.java
@@ -164,4 +164,11 @@ public class RearrangedRowsAndColumns implements RowsAndColumns
);
}
}
+
+ @Nullable
+ @Override
+ public <T> T as(Class<T> clazz)
+ {
+ return null;
+ }
}
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/RowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/RowsAndColumns.java
index 5fe3d49bcfb..da5a57f1e48 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/RowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/RowsAndColumns.java
@@ -19,29 +19,12 @@
package org.apache.druid.query.rowsandcols;
-import com.fasterxml.jackson.core.JsonGenerator;
-import com.fasterxml.jackson.core.JsonParser;
-import com.fasterxml.jackson.databind.DeserializationContext;
-import com.fasterxml.jackson.databind.SerializerProvider;
-import com.fasterxml.jackson.databind.deser.std.StdDeserializer;
-import com.fasterxml.jackson.databind.ser.std.StdSerializer;
-import org.apache.druid.error.DruidException;
-import org.apache.druid.frame.Frame;
-import org.apache.druid.frame.channel.ByteTracker;
-import org.apache.druid.java.util.common.jackson.JacksonUtils;
import org.apache.druid.query.rowsandcols.column.Column;
-import org.apache.druid.query.rowsandcols.concrete.ColumnBasedFrameRowsAndColumns;
-import org.apache.druid.query.rowsandcols.concrete.FrameRowsAndColumns;
-import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowsAndColumns;
import org.apache.druid.query.rowsandcols.semantic.AppendableRowsAndColumns;
import org.apache.druid.query.rowsandcols.semantic.FramedOnHeapAggregatable;
-import org.apache.druid.segment.column.RowSignature;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
-import java.nio.channels.Channels;
import java.util.Collection;
/**
@@ -127,72 +110,7 @@ public interface RowsAndColumns
* @return A concrete implementation of the interface, or null if there is no meaningful optimization to be had
* through a local implementation of the interface.
*/
- @SuppressWarnings("unchecked")
@Nullable
- default <T> T as(Class<T> clazz)
- {
- if (clazz.isInstance(this)) {
- return (T) this;
- }
- return null;
- }
+ <T> T as(Class<T> clazz);
- /**
- * Serializer for {@link RowsAndColumns} by converting the instance to {@link FrameRowsAndColumns}
- */
- class RowsAndColumnsSerializer extends StdSerializer<RowsAndColumns>
- {
- public RowsAndColumnsSerializer()
- {
- super(RowsAndColumns.class);
- }
-
- @Override
- public void serialize(
- RowsAndColumns rac,
- JsonGenerator jsonGenerator,
- SerializerProvider serializerProvider
- ) throws IOException
- {
- FrameRowsAndColumns frameRAC = rac.as(FrameRowsAndColumns.class);
- if (frameRAC == null) {
- throw DruidException.defensive("Unable to serialize RAC");
- }
- JacksonUtils.writeObjectUsingSerializerProvider(jsonGenerator, serializerProvider, frameRAC.getSignature());
-
- Frame frame = frameRAC.getFrame();
- final ByteArrayOutputStream baos = new ByteArrayOutputStream();
- frame.writeTo(Channels.newChannel(baos), false, null, ByteTracker.unboundedTracker());
-
- jsonGenerator.writeBinary(baos.toByteArray());
- }
- }
-
- /**
- * Deserializer for {@link RowsAndColumns} returning as an instance of {@link FrameRowsAndColumns}
- */
- class RowsAndColumnsDeserializer extends StdDeserializer<RowsAndColumns>
- {
- public RowsAndColumnsDeserializer()
- {
- super(RowsAndColumns.class);
- }
-
- @Override
- public FrameRowsAndColumns deserialize(JsonParser jsonParser, DeserializationContext deserializationContext)
- throws IOException
- {
- RowSignature sig = jsonParser.readValueAs(RowSignature.class);
- jsonParser.nextValue();
-
- ByteArrayOutputStream baos = new ByteArrayOutputStream();
- jsonParser.readBinaryValue(baos);
- Frame frame = Frame.wrap(baos.toByteArray());
- if (frame.type().isColumnar()) {
- return new ColumnBasedFrameRowsAndColumns(frame, sig);
- } else {
- return new RowBasedFrameRowsAndColumns(frame, sig);
- }
- }
- }
}
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/AbstractFrameRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/AbstractFrameRowsAndColumns.java
index 5295326c862..a6960ea6ebc 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/AbstractFrameRowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/AbstractFrameRowsAndColumns.java
@@ -19,10 +19,12 @@
package org.apache.druid.query.rowsandcols.concrete;
-import com.google.common.base.Objects;
+import org.apache.druid.error.DruidException;
import org.apache.druid.frame.Frame;
import org.apache.druid.frame.read.FrameReader;
+import org.apache.druid.frame.wire.FrameWireTransferable;
import org.apache.druid.query.rowsandcols.column.Column;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
import org.apache.druid.segment.CloseableShapeshifter;
import org.apache.druid.segment.CursorFactory;
import org.apache.druid.segment.column.RowSignature;
@@ -30,14 +32,16 @@ import org.apache.druid.segment.column.RowSignature;
import javax.annotation.Nullable;
import java.util.Collection;
import java.util.LinkedHashMap;
+import java.util.Objects;
public abstract class AbstractFrameRowsAndColumns implements FrameRowsAndColumns, AutoCloseable, CloseableShapeshifter
{
final Frame frame;
+ @Nullable
final RowSignature signature;
final LinkedHashMap<String, Column> colCache = new LinkedHashMap<>();
- public AbstractFrameRowsAndColumns(Frame frame, RowSignature signature)
+ public AbstractFrameRowsAndColumns(Frame frame, @Nullable RowSignature signature)
{
this.frame = frame;
this.signature = signature;
@@ -49,15 +53,26 @@ public abstract class AbstractFrameRowsAndColumns implements FrameRowsAndColumns
return frame;
}
+ public boolean hasSignature()
+ {
+ return signature != null;
+ }
+
@Override
public RowSignature getSignature()
{
+ if (signature == null) {
+ throw DruidException.defensive("No signature present, cannot call getSignature()");
+ }
return signature;
}
@Override
public Collection<String> getColumnNames()
{
+ if (signature == null) {
+ throw DruidException.defensive("No signature present, cannot call getColumnNames()");
+ }
return signature.getColumnNames();
}
@@ -72,10 +87,19 @@ public abstract class AbstractFrameRowsAndColumns implements FrameRowsAndColumns
@Override
public <T> T as(Class<T> clazz)
{
- if (CursorFactory.class.equals(clazz)) {
+ if (CursorFactory.class.equals(clazz) && signature != null) {
return (T) FrameReader.create(signature).makeCursorFactory(frame);
}
- return FrameRowsAndColumns.super.as(clazz);
+ if (WireTransferable.class.equals(clazz)) {
+ return (T) new FrameWireTransferable(frame, signature);
+ }
+ if (Frame.class.equals(clazz)) {
+ return (T) frame;
+ }
+ if (clazz.isInstance(this)) {
+ return (T) this;
+ }
+ return null;
}
@Override
@@ -84,23 +108,19 @@ public abstract class AbstractFrameRowsAndColumns implements FrameRowsAndColumns
// nothing to close
}
- @Override
- public int hashCode()
- {
- return Objects.hashCode(frame, signature);
- }
-
@Override
public boolean equals(Object o)
{
- if (this == o) {
- return true;
- }
- if (!(o instanceof AbstractFrameRowsAndColumns)) {
+ if (o == null || getClass() != o.getClass()) {
return false;
}
- AbstractFrameRowsAndColumns otherFrame = (AbstractFrameRowsAndColumns) o;
+ AbstractFrameRowsAndColumns that = (AbstractFrameRowsAndColumns) o;
+ return Objects.equals(frame, that.frame) && Objects.equals(signature, that.signature);
+ }
- return frame.writableMemory().equals(otherFrame.frame.writableMemory()) && signature.equals(otherFrame.signature);
+ @Override
+ public int hashCode()
+ {
+ return Objects.hash(frame, signature);
}
}
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/ColumnBasedFrameRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/ColumnBasedFrameRowsAndColumns.java
index d85838a6366..9eab954cd26 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/ColumnBasedFrameRowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/ColumnBasedFrameRowsAndColumns.java
@@ -30,7 +30,7 @@ import javax.annotation.Nullable;
public class ColumnBasedFrameRowsAndColumns extends AbstractFrameRowsAndColumns
{
- public ColumnBasedFrameRowsAndColumns(Frame frame, RowSignature signature)
+ public ColumnBasedFrameRowsAndColumns(Frame frame, @Nullable RowSignature signature)
{
super(frame.ensureColumnar(), signature);
}
@@ -39,6 +39,10 @@ public class ColumnBasedFrameRowsAndColumns extends AbstractFrameRowsAndColumns
@Override
public Column findColumn(String name)
{
+ if (signature == null) {
+ throw DruidException.defensive("No signature present, cannot call findColumn()");
+ }
+
// Use contains so that we can negative cache.
if (!colCache.containsKey(name)) {
final int columnIndex = signature.indexOf(name);
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java
index 022a0f91ac1..d05ac98eef8 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/FrameRowsAndColumns.java
@@ -19,13 +19,31 @@
package org.apache.druid.query.rowsandcols.concrete;
+import org.apache.druid.error.DruidException;
import org.apache.druid.frame.Frame;
import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.segment.CursorFactory;
import org.apache.druid.segment.column.RowSignature;
+/**
+ * Wrapper around {@link Frame} with an optional signature (see {@link #getSignature()}).
+ *
+ * If the signature is available, this object has working {@link #getColumnNames()} and {@link #findColumn(String)}
+ * methods, and can become a {@link CursorFactory} through {@link #as(Class)}. If the signature is not available,
+ * this object can only really be used by getting the underlying {@link Frame} via {@link #getFrame()} or
+ * {@link #as(Class)}.
+ */
public interface FrameRowsAndColumns extends RowsAndColumns
{
+ /**
+ * Returns the frame.
+ */
Frame getFrame();
+ /**
+ * Returns a signature, which may or may not be available.
+ *
+ * @throws DruidException if signature is not available
+ */
RowSignature getSignature();
}
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowsAndColumns.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowsAndColumns.java
index 32f453f0626..5cf27fce0c3 100644
--- a/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowsAndColumns.java
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/concrete/RowBasedFrameRowsAndColumns.java
@@ -31,7 +31,7 @@ import javax.annotation.Nullable;
public class RowBasedFrameRowsAndColumns extends AbstractFrameRowsAndColumns
{
- public RowBasedFrameRowsAndColumns(Frame frame, RowSignature signature)
+ public RowBasedFrameRowsAndColumns(Frame frame, @Nullable RowSignature signature)
{
super(frame.ensureRowBased(), signature);
}
@@ -40,6 +40,10 @@ public class RowBasedFrameRowsAndColumns extends AbstractFrameRowsAndColumns
@Override
public Column findColumn(String name)
{
+ if (signature == null) {
+ throw DruidException.defensive("No signature present, cannot call findColumn()");
+ }
+
// Use contains so that we can negative cache.
if (!colCache.containsKey(name)) {
final int columnIndex = signature.indexOf(name);
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/semantic/WireTransferable.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/semantic/WireTransferable.java
new file mode 100644
index 00000000000..fb73ee1e727
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/semantic/WireTransferable.java
@@ -0,0 +1,157 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.rowsandcols.semantic;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.guice.DruidSecondaryModule;
+import org.apache.druid.java.util.common.ByteBufferUtils;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Map;
+
+/**
+ * A Semantic interface that enables serializing a {@link RowsAndColumns} over the wire.
+ * <p>
+ * Serialization and Deserialization are a weird beast. It is very easy to follow an Object-Oriented mechanism for
+ * serializing out an object, but when you want to deserialize, it is unclear exactly which class should be
+ * deserialized. This requires extra metadata that helps point the deserializing code to how to interpret the
+ * bytes.
+ * <p>
+ * The {@link WireTransferable} works by serializing the object as a singular byte array, the first byte of the byte array
+ * is interpretted as an unsigned integer (N) and represents the length of the "extra metadata" to identify the correct
+ * deserializer. Those N bytes are read as a utf8 String and used to look up the deserialization logic. This
+ * identifier must be registered with {@link DruidSecondaryModule#getWireTransferableDeserializerBinder}, which will
+ * use the identifier to look up the logic and deserialize the RowsAndColumns object.
+ * <p>
+ * There is, intentionally, no default implementation of {@link WireTransferable} that just wraps a RowsAndColumns,
+ * this is because the point at which data is transfered over the wire tends to be a very sensitive point in terms of
+ * what the best format for transmission is and what optimizations are available. As such, we don't want a concrete
+ * class that makes its way to the wire to accidentally get a least-common-denominator serialization and instead want
+ * it to be explicitly chosen.
+ */
+public interface WireTransferable
+{
+ /**
+ * Serializes the object into bytes. This method is passed a jackson ObjectMapper that can be used to serialize
+ * objects if needed. In general, it's better to generate binary serializations and not rely on the mapper
+ * but it's there as an expediency layer to support some compatibility stuff. Hopefully it goes away at some point.
+ *
+ * @param mapper ObjectMapper
+ * @return bytes that represent the wire-format of this object
+ * @throws IOException a problem with IO
+ */
+ ByteArrayOffsetAndLen serializedBytes(ObjectMapper mapper) throws IOException;
+
+ interface Deserializer
+ {
+ /**
+ * Deserializes a ByteBuffer into a RowsAndColumns. The deserialized object owns the passed-in buffer as well
+ * as the data it points to. It is safe for the deserializer to mutate the buffer, retain a reference to it, etc.
+ *
+ * @param mapper ObjectMapper
+ * @param theBytes ByteBuffer to deserialize, in little-endian order
+ * @return RowsAndColumns object deserialized from the ByteBuffer
+ */
+ RowsAndColumns deserialize(ObjectMapper mapper, ByteBuffer theBytes);
+ }
+
+ /**
+ * A holder class for a byte array, offset and length, this exists to minimize copies. It's pretty similar to a
+ * ByteBuffer, but is separate because jackson cannot write a ByteBuffer anyway so might as well wrap in this.
+ */
+ class ByteArrayOffsetAndLen
+ {
+ private final byte[] array;
+ private final int offset;
+ private final int length;
+
+ public ByteArrayOffsetAndLen(byte[] array, int offset, int length)
+ {
+ this.array = array;
+ this.offset = offset;
+ this.length = length;
+ }
+
+ public byte[] getArray()
+ {
+ return array;
+ }
+
+ public int getOffset()
+ {
+ return offset;
+ }
+
+ public int getLength()
+ {
+ return length;
+ }
+
+ public ByteBuffer asByteBuffer()
+ {
+ return ByteBufferUtils.wrapLE(array, offset, length);
+ }
+ }
+
+ class ConcreteDeserializer
+ {
+ private final ObjectMapper mapper;
+ private final Map<ByteBuffer, Deserializer> delegates;
+
+ public ConcreteDeserializer(ObjectMapper mapper, Map<ByteBuffer, Deserializer> delegates)
+ {
+ this.mapper = mapper;
+ this.delegates = delegates;
+ }
+
+ /**
+ * Deserializes a ByteBuffer into a RowsAndColumns. The deserialized object owns the passed-in buffer as well
+ * as the data it points to. It is safe for the deserializer to mutate the buffer, retain a reference to it, etc.
+ *
+ * @param theBytes ByteBuffer to be deserialized, in little-endian order
+ * @return RowsAndColumns object deserialized from the ByteBuffer
+ */
+ public RowsAndColumns deserialize(ByteBuffer theBytes)
+ {
+ int originalPos = theBytes.position();
+ int originalLimit = theBytes.limit();
+
+ int identifierLen = 0xff & theBytes.get();
+ theBytes.limit(theBytes.position() + identifierLen);
+ final Deserializer deserializer = delegates.get(theBytes);
+ if (deserializer == null) {
+ throw DruidException
+ .forPersona(DruidException.Persona.OPERATOR)
+ .ofCategory(DruidException.Category.RUNTIME_FAILURE)
+ .build(
+ "Got unhandled RAC blob type [%s], either an extension is missing or code is not wired correctly.",
+ StringUtils.fromUtf8Nullable(theBytes)
+ );
+ }
+ theBytes.position(originalPos);
+ theBytes.limit(originalLimit);
+ return deserializer.deserialize(mapper, theBytes);
+ }
+ }
+}
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/serde/RowsAndColumnsDeserializer.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/serde/RowsAndColumnsDeserializer.java
new file mode 100644
index 00000000000..e6305f15792
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/serde/RowsAndColumnsDeserializer.java
@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.rowsandcols.serde;
+
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.DeserializationContext;
+import com.fasterxml.jackson.databind.deser.std.StdDeserializer;
+import com.fasterxml.jackson.databind.exc.MismatchedInputException;
+import org.apache.druid.error.ErrorResponse;
+import org.apache.druid.frame.Frame;
+import org.apache.druid.java.util.common.ByteBufferUtils;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.ColumnBasedFrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.FrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.apache.druid.segment.column.RowSignature;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+
+/**
+ * Deserializer for {@link RowsAndColumns} following the contract set out in {@link WireTransferable}.
+ *
+ * There is some legacy code in {@link #deserializeLegacyFrame(JsonParser)} that handles deserializing
+ * {@link FrameRowsAndColumns} specifically as a signature object followed by frame bytes. This code is
+ * for compatibility with the format written by {@link RowsAndColumnsSerializer} when
+ * {@link WireTransferableContext#useLegacyFrameSerialization()} is true.
+ */
+public class RowsAndColumnsDeserializer extends StdDeserializer<RowsAndColumns>
+{
+ private final WireTransferableContext wtContext;
+
+ public RowsAndColumnsDeserializer(WireTransferableContext wtContext)
+ {
+ super(RowsAndColumns.class);
+ this.wtContext = wtContext;
+ }
+
+ @Override
+ public RowsAndColumns deserialize(JsonParser jsonParser, DeserializationContext deserializationContext)
+ throws IOException
+ {
+ if (jsonParser.currentToken().isScalarValue()) {
+ final byte[] bytes = jsonParser.readValueAs(byte[].class);
+ return wtContext.deserialize(ByteBufferUtils.wrapLE(bytes));
+ } else {
+ return deserializeLegacyFrame(jsonParser);
+ }
+ }
+
+ /**
+ * Deserializes the format written by {@link RowsAndColumnsSerializer} when
+ * {@link WireTransferableContext#useLegacyFrameSerialization()} is true.
+ */
+ private RowsAndColumns deserializeLegacyFrame(JsonParser jsonParser) throws IOException
+ {
+ final RowSignature sig;
+ try {
+ sig = jsonParser.readValueAs(RowSignature.class);
+ }
+ catch (MismatchedInputException e) {
+ throw jsonParser.readValueAs(ErrorResponse.class).getUnderlyingException();
+ }
+ jsonParser.nextValue();
+
+ final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ jsonParser.readBinaryValue(baos);
+ final Frame frame = Frame.wrap(baos.toByteArray());
+ if (frame.type().isColumnar()) {
+ return new ColumnBasedFrameRowsAndColumns(frame, sig);
+ } else {
+ return new RowBasedFrameRowsAndColumns(frame, sig);
+ }
+ }
+}
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/serde/RowsAndColumnsSerializer.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/serde/RowsAndColumnsSerializer.java
new file mode 100644
index 00000000000..0676305248d
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/serde/RowsAndColumnsSerializer.java
@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.rowsandcols.serde;
+
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.databind.SerializerProvider;
+import com.fasterxml.jackson.databind.ser.std.StdSerializer;
+import org.apache.druid.error.NotYetImplemented;
+import org.apache.druid.frame.Frame;
+import org.apache.druid.java.util.common.jackson.JacksonUtils;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.FrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.channels.Channels;
+
+/**
+ * Serializer for {@link RowsAndColumns} by converting the instance to {@link WireTransferable}. See the javadoc
+ * on {@link WireTransferable} for information on the serde contract.
+ *
+ * There is a legacy serialization format for {@link FrameRowsAndColumns} that serializes the signature object
+ * followed by frame bytes. This legacy format can be enabled via
+ * {@link WireTransferableContext#useLegacyFrameSerialization()} for compatibility during rolling updates.
+ * Once all servers have been upgraded, this can be set to false.
+ */
+public class RowsAndColumnsSerializer extends StdSerializer<RowsAndColumns>
+{
+ private final WireTransferableContext wtContext;
+
+ public RowsAndColumnsSerializer(WireTransferableContext wtContext)
+ {
+ super(RowsAndColumns.class);
+ this.wtContext = wtContext;
+ }
+
+ @Override
+ public void serialize(
+ RowsAndColumns rac,
+ JsonGenerator jsonGenerator,
+ SerializerProvider serializerProvider
+ ) throws IOException
+ {
+ // Check if we should use the legacy Frame serialization path
+ if (wtContext.useLegacyFrameSerialization()) {
+ final FrameRowsAndColumns frameRAC = rac.as(FrameRowsAndColumns.class);
+ if (frameRAC != null) {
+ serializeFrameLegacy(frameRAC, jsonGenerator, serializerProvider);
+ return;
+ }
+ }
+
+ // Use WireTransferable path
+ final WireTransferable as = rac.as(WireTransferable.class);
+ if (as != null) {
+ final WireTransferable.ByteArrayOffsetAndLen serMe = wtContext.serializedBytes(as);
+ jsonGenerator.writeBinary(serMe.getArray(), serMe.getOffset(), serMe.getLength());
+ jsonGenerator.flush();
+ return;
+ }
+
+ throw NotYetImplemented.ex(null, "RAC [%s] cannot be serialized, implement WireTransferable", rac.getClass());
+ }
+
+ /**
+ * Legacy serialization format for FrameRowsAndColumns: signature object followed by frame bytes.
+ * This format exists for compatibility during rolling updates and should be removed in a future release.
+ */
+ private static void serializeFrameLegacy(
+ FrameRowsAndColumns frameRAC,
+ JsonGenerator jsonGenerator,
+ SerializerProvider serializerProvider
+ ) throws IOException
+ {
+ JacksonUtils.writeObjectUsingSerializerProvider(jsonGenerator, serializerProvider, frameRAC.getSignature());
+
+ final Frame frame = frameRAC.getFrame();
+ final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ frame.writeTo(Channels.newChannel(baos), false, null);
+
+ jsonGenerator.writeBinary(baos.toByteArray());
+ }
+}
diff --git a/processing/src/main/java/org/apache/druid/query/rowsandcols/serde/WireTransferableContext.java b/processing/src/main/java/org/apache/druid/query/rowsandcols/serde/WireTransferableContext.java
new file mode 100644
index 00000000000..920ffd4c08b
--- /dev/null
+++ b/processing/src/main/java/org/apache/druid/query/rowsandcols/serde/WireTransferableContext.java
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.rowsandcols.serde;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.frame.file.FrameFileWriter;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+/**
+ * Object that holds the components and configuration needed for serde of {@link RowsAndColumns}
+ * via {@link WireTransferable}.
+ */
+public class WireTransferableContext
+{
+ public static final boolean DEFAULT_LEGACY_FRAME_SERIALIZATION = true;
+
+ private final ObjectMapper smileMapper;
+ private final WireTransferable.ConcreteDeserializer concreteDeserializer;
+ private final boolean useLegacyFrameSerialization;
+
+ public WireTransferableContext(
+ final ObjectMapper smileMapper,
+ final WireTransferable.ConcreteDeserializer concreteDeserializer,
+ final boolean useLegacyFrameSerialization
+ )
+ {
+ this.smileMapper = smileMapper;
+ this.concreteDeserializer = concreteDeserializer;
+ this.useLegacyFrameSerialization = useLegacyFrameSerialization;
+ }
+
+ /**
+ * Serializes an object into bytes.
+ *
+ * @return bytes that represent the wire-format of this object
+ *
+ * @throws IOException a problem with IO
+ */
+ public WireTransferable.ByteArrayOffsetAndLen serializedBytes(WireTransferable wireTransferable)
+ throws IOException
+ {
+ return wireTransferable.serializedBytes(smileMapper);
+ }
+
+ /**
+ * Deserializes a ByteBuffer into a RowsAndColumns. The deserialized object owns the passed-in buffer as well
+ * as the data it points to. It is safe for the deserializer to mutate the buffer, retain a reference to it, etc.
+ *
+ * @param theBytes ByteBuffer to be deserialized, in little-endian order
+ *
+ * @return RowsAndColumns object deserialized from the ByteBuffer
+ */
+ public RowsAndColumns deserialize(ByteBuffer theBytes)
+ {
+ if (concreteDeserializer == null) {
+ throw DruidException.defensive("Cannot deserialize, no concreteDeserializer");
+ }
+ return concreteDeserializer.deserialize(theBytes);
+ }
+
+ /**
+ * Returns the Smile ObjectMapper for serialization, or null if not available.
+ */
+ public ObjectMapper smileMapper()
+ {
+ return smileMapper;
+ }
+
+ /**
+ * Returns the deserializer for RowsAndColumns wire format, or null if not available.
+ */
+ public WireTransferable.ConcreteDeserializer concreteDeserializer()
+ {
+ return concreteDeserializer;
+ }
+
+ /**
+ * Returns true if legacy frame serialization formats should be used in {@link RowsAndColumnsSerializer} and
+ * {@link FrameFileWriter}. Returns false if the {@link WireTransferable} mechanism should be used.
+ */
+ public boolean useLegacyFrameSerialization()
+ {
+ return useLegacyFrameSerialization;
+ }
+}
diff --git a/processing/src/test/java/org/apache/druid/frame/FrameTest.java b/processing/src/test/java/org/apache/druid/frame/FrameTest.java
index 45080ee262d..bba3249076d 100644
--- a/processing/src/test/java/org/apache/druid/frame/FrameTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/FrameTest.java
@@ -25,7 +25,6 @@ import com.google.common.io.Files;
import com.google.common.primitives.Ints;
import org.apache.datasketches.memory.Memory;
import org.apache.datasketches.memory.WritableMemory;
-import org.apache.druid.frame.channel.ByteTracker;
import org.apache.druid.frame.key.KeyColumn;
import org.apache.druid.frame.key.KeyOrder;
import org.apache.druid.frame.testutil.FrameSequenceBuilder;
@@ -351,8 +350,7 @@ public class FrameTest
frame.writeTo(
Channels.newChannel(baos),
compressed,
- ByteBuffer.allocate(Frame.compressionBufferSize((int) frame.numBytes())),
- ByteTracker.unboundedTracker()
+ ByteBuffer.allocate(Frame.compressionBufferSize((int) frame.numBytes()))
);
if (!compressed) {
@@ -415,8 +413,7 @@ public class FrameTest
frame.writeTo(
Channels.newChannel(baos),
compressed,
- ByteBuffer.allocate(Frame.compressionBufferSize((int) frame.numBytes())),
- ByteTracker.unboundedTracker()
+ ByteBuffer.allocate(Frame.compressionBufferSize((int) frame.numBytes()))
);
return baos.toByteArray();
}
diff --git a/processing/src/test/java/org/apache/druid/frame/channel/ComposingWritableFrameChannelTest.java b/processing/src/test/java/org/apache/druid/frame/channel/ComposingWritableFrameChannelTest.java
index abcb9f03cfb..7a6dfef2c6b 100644
--- a/processing/src/test/java/org/apache/druid/frame/channel/ComposingWritableFrameChannelTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/channel/ComposingWritableFrameChannelTest.java
@@ -27,6 +27,7 @@ import org.apache.druid.frame.Frame;
import org.apache.druid.frame.allocation.ArenaMemoryAllocator;
import org.apache.druid.frame.processor.OutputChannel;
import org.apache.druid.query.ResourceLimitExceededException;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import org.hamcrest.CoreMatchers;
import org.hamcrest.MatcherAssert;
import org.junit.Assert;
@@ -84,9 +85,9 @@ public class ComposingWritableFrameChannelTest
partitionToChannelMap
);
- composingWritableFrameChannel.write(new FrameWithPartition(Mockito.mock(Frame.class), 1));
- composingWritableFrameChannel.write(new FrameWithPartition(Mockito.mock(Frame.class), 2));
- composingWritableFrameChannel.write(new FrameWithPartition(Mockito.mock(Frame.class), 3));
+ composingWritableFrameChannel.write(Mockito.mock(Frame.class), 1);
+ composingWritableFrameChannel.write(Mockito.mock(Frame.class), 2);
+ composingWritableFrameChannel.write(Mockito.mock(Frame.class), 3);
// Assert the location of the channels where the frames have been written to
Assert.assertEquals(ImmutableSet.of(0), partitionToChannelMap.get(1));
@@ -124,7 +125,7 @@ public class ComposingWritableFrameChannelTest
}
@Override
- public void write(FrameWithPartition frameWithPartition)
+ public void write(RowsAndColumns rac, int partitionNumber)
{
if (curFrame >= maxFrames) {
throw new ResourceLimitExceededException("Cannot write more frames to the channel");
@@ -132,11 +133,6 @@ public class ComposingWritableFrameChannelTest
++curFrame;
}
- @Override
- public void write(Frame frame)
- {
- }
-
@Override
public void fail(@Nullable Throwable cause)
{
diff --git a/processing/src/test/java/org/apache/druid/frame/channel/ReadableByteChunksFrameChannelTest.java b/processing/src/test/java/org/apache/druid/frame/channel/ReadableByteChunksFrameChannelTest.java
index 7ff4cd31a20..bf8d8558a34 100644
--- a/processing/src/test/java/org/apache/druid/frame/channel/ReadableByteChunksFrameChannelTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/channel/ReadableByteChunksFrameChannelTest.java
@@ -19,17 +19,28 @@
package org.apache.druid.frame.channel;
+import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.io.ByteStreams;
import com.google.common.io.Files;
import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.frame.Frame;
import org.apache.druid.frame.FrameType;
import org.apache.druid.frame.allocation.ArenaMemoryAllocator;
+import org.apache.druid.frame.file.FrameFileWriter;
import org.apache.druid.frame.read.FrameReader;
import org.apache.druid.frame.testutil.FrameSequenceBuilder;
import org.apache.druid.frame.testutil.FrameTestUtil;
+import org.apache.druid.frame.wire.FrameWireTransferable;
import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.common.guava.Sequence;
import org.apache.druid.java.util.common.guava.Sequences;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
+import org.apache.druid.segment.CursorFactory;
import org.apache.druid.segment.TestIndex;
+import org.apache.druid.segment.column.RowSignature;
import org.apache.druid.segment.incremental.IncrementalIndexCursorFactory;
import org.apache.druid.testing.InitializedNullHandlingTest;
import org.hamcrest.CoreMatchers;
@@ -47,9 +58,11 @@ import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
+import java.util.Map;
public class ReadableByteChunksFrameChannelTest
{
@@ -67,7 +80,7 @@ public class ReadableByteChunksFrameChannelTest
@Test
public void testZeroBytes()
{
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false);
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false, null);
channel.doneWriting();
Assert.assertTrue(channel.canRead());
@@ -77,13 +90,13 @@ public class ReadableByteChunksFrameChannelTest
expectedException.expect(IllegalStateException.class);
expectedException.expectMessage("Incomplete or missing frame at end of stream (id = test, position = 0)");
- channel.read();
+ channel.readFrame();
}
@Test
public void testZeroBytesWithSpecialError()
{
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false);
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false, null);
channel.setError(new IllegalArgumentException("test error"));
channel.doneWriting();
@@ -94,7 +107,7 @@ public class ReadableByteChunksFrameChannelTest
expectedException.expect(IllegalArgumentException.class);
expectedException.expectMessage("test error");
- channel.read();
+ channel.readFrame();
}
@Test
@@ -103,7 +116,7 @@ public class ReadableByteChunksFrameChannelTest
// File with no frames (but still well-formed).
final File file = FrameTestUtil.writeFrameFile(Sequences.empty(), temporaryFolder.newFile());
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false);
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false, null);
channel.addChunk(Files.toByteArray(file));
channel.doneWriting();
Assert.assertEquals(file.length(), channel.getBytesAdded());
@@ -111,7 +124,7 @@ public class ReadableByteChunksFrameChannelTest
while (channel.canRead()) {
Assert.assertFalse(channel.isFinished());
Assert.assertFalse(channel.isErrorOrFinished());
- channel.read();
+ channel.readFrame();
}
Assert.assertTrue(channel.isFinished());
@@ -121,7 +134,7 @@ public class ReadableByteChunksFrameChannelTest
@Test
public void testAddChunkAfterDoneWriting()
{
- try (final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false)) {
+ try (final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false, null)) {
channel.doneWriting();
Assert.assertThrows(
@@ -154,7 +167,7 @@ public class ReadableByteChunksFrameChannelTest
ByteStreams.readFully(in, truncatedFile);
}
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false);
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false, null);
channel.addChunk(truncatedFile);
channel.doneWriting();
Assert.assertEquals(truncatedFile.length, channel.getBytesAdded());
@@ -162,12 +175,12 @@ public class ReadableByteChunksFrameChannelTest
Assert.assertTrue(channel.canRead());
Assert.assertFalse(channel.isFinished());
Assert.assertFalse(channel.isErrorOrFinished());
- channel.read(); // Throw away value.
+ channel.readFrame(); // Throw away value.
Assert.assertTrue(channel.canRead());
Assert.assertFalse(channel.isFinished());
Assert.assertFalse(channel.isErrorOrFinished());
- channel.read(); // Throw away value.
+ channel.readFrame(); // Throw away value.
Assert.assertTrue(channel.canRead());
Assert.assertFalse(channel.isFinished());
@@ -175,7 +188,7 @@ public class ReadableByteChunksFrameChannelTest
expectedException.expect(IllegalStateException.class);
expectedException.expectMessage(CoreMatchers.startsWith("Incomplete or missing frame at end of stream"));
- channel.read();
+ channel.readFrame();
}
@Test
@@ -195,7 +208,7 @@ public class ReadableByteChunksFrameChannelTest
temporaryFolder.newFile()
);
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false);
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false, null);
final byte[] fileBytes = Files.toByteArray(file);
final byte[] chunk1 = new byte[errorAtBytePosition];
System.arraycopy(fileBytes, 0, chunk1, 0, chunk1.length);
@@ -208,7 +221,7 @@ public class ReadableByteChunksFrameChannelTest
expectedException.expect(IllegalStateException.class);
expectedException.expectMessage("Test error!");
- channel.read();
+ channel.readFrame();
}
}
@@ -224,15 +237,22 @@ public class ReadableByteChunksFrameChannelTest
private final FrameType frameType;
private final int maxRowsPerFrame;
private final int chunkSize;
-
- public ParameterizedWithTestIndexTests(final FrameType frameType, final int maxRowsPerFrame, final int chunkSize)
+ private final boolean useLegacyFrameSerialization;
+
+ public ParameterizedWithTestIndexTests(
+ final FrameType frameType,
+ final int maxRowsPerFrame,
+ final int chunkSize,
+ final boolean useLegacyFrameSerialization
+ )
{
this.frameType = frameType;
this.maxRowsPerFrame = maxRowsPerFrame;
this.chunkSize = chunkSize;
+ this.useLegacyFrameSerialization = useLegacyFrameSerialization;
}
- @Parameterized.Parameters(name = "frameType = {0}, maxRowsPerFrame = {1}, chunkSize = {2}")
+ @Parameterized.Parameters(name = "frameType = {0}, maxRowsPerFrame = {1}, chunkSize = {2}, useLegacyFrameSerialization = {3}")
public static Iterable<Object[]> constructorFeeder()
{
final List<Object[]> constructors = new ArrayList<>();
@@ -240,7 +260,9 @@ public class ReadableByteChunksFrameChannelTest
for (FrameType frameType : FrameType.values()) {
for (int maxRowsPerFrame : new int[]{1, 50, Integer.MAX_VALUE}) {
for (int chunkSize : new int[]{1, 10, 1_000, 5_000, 50_000, 1_000_000}) {
- constructors.add(new Object[]{frameType, maxRowsPerFrame, chunkSize});
+ for (boolean useLegacyFrameSerialization : new boolean[]{true, false}) {
+ constructors.add(new Object[]{frameType, maxRowsPerFrame, chunkSize, useLegacyFrameSerialization});
+ }
}
}
}
@@ -248,6 +270,51 @@ public class ReadableByteChunksFrameChannelTest
return constructors;
}
+ private WireTransferableContext makeWireTransferableContext()
+ {
+ if (useLegacyFrameSerialization) {
+ return FrameTestUtil.WT_CONTEXT_LEGACY;
+ } else {
+ final ObjectMapper objectMapper = new ObjectMapper();
+ final WireTransferable.ConcreteDeserializer deserializer = new WireTransferable.ConcreteDeserializer(
+ objectMapper,
+ Map.of(
+ ByteBuffer.wrap(StringUtils.toUtf8(FrameWireTransferable.TYPE)),
+ new FrameWireTransferable.Deserializer()
+ )
+ );
+ return new WireTransferableContext(objectMapper, deserializer, false);
+ }
+ }
+
+ private File writeFrameFile(final CursorFactory cursorFactory) throws IOException
+ {
+ final WireTransferableContext wireTransferableContext = makeWireTransferableContext();
+ final Sequence<Frame> frames = FrameSequenceBuilder.fromCursorFactory(cursorFactory)
+ .maxRowsPerFrame(maxRowsPerFrame)
+ .frameType(frameType)
+ .frames();
+
+ final File file = temporaryFolder.newFile();
+ try (final FrameFileWriter writer = FrameFileWriter.open(
+ Channels.newChannel(java.nio.file.Files.newOutputStream(file.toPath())),
+ null,
+ ByteTracker.unboundedTracker(),
+ wireTransferableContext
+ )) {
+ frames.forEach(frame -> {
+ try {
+ writer.write(frame.asRAC(), FrameFileWriter.NO_PARTITION);
+ }
+ catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ });
+ }
+
+ return file;
+ }
+
@Test
public void testWriteFullyThenRead() throws IOException
{
@@ -255,15 +322,11 @@ public class ReadableByteChunksFrameChannelTest
final IncrementalIndexCursorFactory cursorFactory =
new IncrementalIndexCursorFactory(TestIndex.getIncrementalTestIndex());
- final File file = FrameTestUtil.writeFrameFile(
- FrameSequenceBuilder.fromCursorFactory(cursorFactory)
- .maxRowsPerFrame(maxRowsPerFrame)
- .frameType(frameType)
- .frames(),
- temporaryFolder.newFile()
- );
+ final File file = writeFrameFile(cursorFactory);
+ final WireTransferableContext wireTransferableContext = makeWireTransferableContext();
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false);
+ final ReadableByteChunksFrameChannel channel =
+ ReadableByteChunksFrameChannel.create("test", false, wireTransferableContext);
ListenableFuture<?> firstBackpressureFuture = null;
long totalSize = 0;
@@ -310,15 +373,11 @@ public class ReadableByteChunksFrameChannelTest
final IncrementalIndexCursorFactory cursorFactory =
new IncrementalIndexCursorFactory(TestIndex.getIncrementalTestIndex());
- final File file = FrameTestUtil.writeFrameFile(
- FrameSequenceBuilder.fromCursorFactory(cursorFactory)
- .maxRowsPerFrame(maxRowsPerFrame)
- .frameType(frameType)
- .frames(),
- temporaryFolder.newFile()
- );
+ final File file = writeFrameFile(cursorFactory);
+ final WireTransferableContext wireTransferableContext = makeWireTransferableContext();
- final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create("test", false);
+ final ReadableByteChunksFrameChannel channel =
+ ReadableByteChunksFrameChannel.create("test", false, wireTransferableContext);
final BlockingQueueFrameChannel outChannel = new BlockingQueueFrameChannel(10_000); // Enough to hold all frames
ListenableFuture<?> backpressureFuture = null;
@@ -332,14 +391,14 @@ public class ReadableByteChunksFrameChannelTest
// Read one frame every 3 iterations. Read everything every 11 iterations. Otherwise, read nothing.
if (iteration % 3 == 0) {
while (channel.canRead()) {
- outChannel.writable().write(channel.read());
+ outChannel.writable().write(channel.readFrame());
}
// After reading everything, backpressure should be off.
Assert.assertTrue(backpressureFuture == null || backpressureFuture.isDone());
} else if (iteration % 11 == 0) {
if (channel.canRead()) {
- outChannel.writable().write(channel.read());
+ outChannel.writable().write(channel.readFrame());
}
}
@@ -370,7 +429,7 @@ public class ReadableByteChunksFrameChannelTest
// Get all the remaining frames.
while (channel.canRead()) {
- outChannel.writable().write(channel.read());
+ outChannel.writable().write(channel.readFrame());
}
outChannel.writable().close();
@@ -382,6 +441,124 @@ public class ReadableByteChunksFrameChannelTest
);
}
+ @Test
+ public void testReadRac() throws IOException
+ {
+ final CursorFactory cursorFactory = new IncrementalIndexCursorFactory(
+ TestIndex.getNoRollupIncrementalTestIndex()
+ );
+ final RowSignature signature = cursorFactory.getRowSignature();
+ final FrameReader frameReader = FrameReader.create(signature);
+
+ final File file = writeFrameFile(cursorFactory);
+ final WireTransferableContext wireTransferableContext = makeWireTransferableContext();
+
+ // Create channel with wireTransferableContext
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create(
+ "test-rac",
+ false,
+ wireTransferableContext
+ );
+
+ // Add file contents in chunks
+ try (final FileInputStream in = new FileInputStream(file)) {
+ final byte[] buf = new byte[chunkSize];
+ int bytesRead;
+ while ((bytesRead = in.read(buf)) >= 0) {
+ if (bytesRead > 0) {
+ channel.addChunk(Arrays.copyOf(buf, bytesRead));
+ }
+ }
+ }
+ channel.doneWriting();
+
+ // Read and verify using readRAC()
+ final List<List<Object>> readRows = new ArrayList<>();
+ while (channel.canRead()) {
+ final RowsAndColumns rac = channel.read();
+ final Frame frame = rac.as(Frame.class);
+ Assert.assertNotNull("Frame should be retrievable from RAC", frame);
+
+ // Read rows from this frame
+ final Sequence<List<Object>> frameRows = FrameTestUtil.readRowsFromCursorFactory(
+ frameReader.makeCursorFactory(frame)
+ );
+ readRows.addAll(frameRows.toList());
+ }
+
+ Assert.assertTrue(channel.isFinished());
+
+ // Verify data integrity
+ final List<List<Object>> originalRows =
+ FrameTestUtil.readRowsFromCursorFactory(cursorFactory).toList();
+ Assert.assertEquals(originalRows, readRows);
+ }
+
+ @Test
+ public void testReadRacInterleaved() throws IOException
+ {
+ final CursorFactory cursorFactory = new IncrementalIndexCursorFactory(
+ TestIndex.getNoRollupIncrementalTestIndex()
+ );
+ final RowSignature signature = cursorFactory.getRowSignature();
+ final FrameReader frameReader = FrameReader.create(signature);
+
+ final File file = writeFrameFile(cursorFactory);
+ final WireTransferableContext wireTransferableContext = makeWireTransferableContext();
+
+ // Create channel with wireTransferableContext
+ final ReadableByteChunksFrameChannel channel = ReadableByteChunksFrameChannel.create(
+ "test-rac-interleaved",
+ false,
+ wireTransferableContext
+ );
+
+ // Add file contents in chunks, interleaving reads
+ final List<List<Object>> readRows = new ArrayList<>();
+ int iteration = 0;
+
+ try (final FileInputStream in = new FileInputStream(file)) {
+ final byte[] buf = new byte[chunkSize];
+ int bytesRead;
+ while ((bytesRead = in.read(buf)) >= 0) {
+ // Read every 3 iterations
+ if (iteration % 3 == 0) {
+ while (channel.canRead()) {
+ final RowsAndColumns rac = channel.read();
+ final Frame frame = rac.as(Frame.class);
+ final Sequence<List<Object>> frameRows = FrameTestUtil.readRowsFromCursorFactory(
+ frameReader.makeCursorFactory(frame)
+ );
+ readRows.addAll(frameRows.toList());
+ }
+ }
+
+ if (bytesRead > 0) {
+ channel.addChunk(Arrays.copyOf(buf, bytesRead));
+ }
+ iteration++;
+ }
+ }
+ channel.doneWriting();
+
+ // Read remaining frames
+ while (channel.canRead()) {
+ final RowsAndColumns rac = channel.read();
+ final Frame frame = rac.as(Frame.class);
+ final Sequence<List<Object>> frameRows = FrameTestUtil.readRowsFromCursorFactory(
+ frameReader.makeCursorFactory(frame)
+ );
+ readRows.addAll(frameRows.toList());
+ }
+
+ Assert.assertTrue(channel.isFinished());
+
+ // Verify data integrity
+ final List<List<Object>> originalRows =
+ FrameTestUtil.readRowsFromCursorFactory(cursorFactory).toList();
+ Assert.assertEquals(originalRows, readRows);
+ }
+
private static class Chunker implements Closeable
{
private final FileInputStream in;
diff --git a/processing/src/test/java/org/apache/druid/frame/channel/ReadableInputStreamFrameChannelTest.java b/processing/src/test/java/org/apache/druid/frame/channel/ReadableInputStreamFrameChannelTest.java
index cc1d67dc595..3cfc15c74ce 100644
--- a/processing/src/test/java/org/apache/druid/frame/channel/ReadableInputStreamFrameChannelTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/channel/ReadableInputStreamFrameChannelTest.java
@@ -72,7 +72,8 @@ public class ReadableInputStreamFrameChannelTest extends InitializedNullHandling
inputStream,
"readSimpleFrameFile",
executorService,
- false
+ false,
+ null
);
FrameTestUtil.assertRowsEqual(
@@ -96,7 +97,8 @@ public class ReadableInputStreamFrameChannelTest extends InitializedNullHandling
Files.newInputStream(file.toPath()),
"readEmptyFrameFile",
executorService,
- false
+ false,
+ null
);
Assert.assertEquals(FrameTestUtil.readRowsFromFrameChannel(
@@ -119,7 +121,8 @@ public class ReadableInputStreamFrameChannelTest extends InitializedNullHandling
Files.newInputStream(file.toPath()),
"testZeroBytesFrameFile",
executorService,
- false
+ false,
+ null
);
final IllegalStateException e = Assert.assertThrows(
@@ -165,7 +168,8 @@ public class ReadableInputStreamFrameChannelTest extends InitializedNullHandling
new ByteArrayInputStream(truncatedFile),
"readTruncatedFrameFile",
executorService,
- false
+ false,
+ null
);
expectedException.expect(ISE.class);
@@ -191,7 +195,8 @@ public class ReadableInputStreamFrameChannelTest extends InitializedNullHandling
Files.newInputStream(file.toPath()),
"readIncorrectFrameFile",
executorService,
- false
+ false,
+ null
);
expectedException.expect(ISE.class);
@@ -215,7 +220,8 @@ public class ReadableInputStreamFrameChannelTest extends InitializedNullHandling
inputStream,
"closeInputStreamWhileReading",
executorService,
- false
+ false,
+ null
);
inputStream.close();
@@ -240,7 +246,8 @@ public class ReadableInputStreamFrameChannelTest extends InitializedNullHandling
inputStream,
"closeInputStreamWhileReadingCheckError",
executorService,
- false
+ false,
+ null
);
inputStream.close();
@@ -251,7 +258,7 @@ public class ReadableInputStreamFrameChannelTest extends InitializedNullHandling
while (!readableInputStreamFrameChannel.canRead()) {
Thread.sleep(10);
}
- readableInputStreamFrameChannel.read();
+ readableInputStreamFrameChannel.readFrame();
Assert.assertTrue(readableInputStreamFrameChannel.isFinished());
readableInputStreamFrameChannel.close();
}
diff --git a/processing/src/test/java/org/apache/druid/frame/file/FrameFileHttpResponseHandlerTest.java b/processing/src/test/java/org/apache/druid/frame/file/FrameFileHttpResponseHandlerTest.java
index 45e45f04867..958ecf88efd 100644
--- a/processing/src/test/java/org/apache/druid/frame/file/FrameFileHttpResponseHandlerTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/file/FrameFileHttpResponseHandlerTest.java
@@ -104,7 +104,7 @@ public class FrameFileHttpResponseHandlerTest extends InitializedNullHandlingTes
temporaryFolder.newFile()
);
- channel = ReadableByteChunksFrameChannel.create("test", false);
+ channel = ReadableByteChunksFrameChannel.create("test", false, null);
handler = new FrameFileHttpResponseHandler(channel);
}
diff --git a/processing/src/test/java/org/apache/druid/frame/file/FrameFileTest.java b/processing/src/test/java/org/apache/druid/frame/file/FrameFileTest.java
index 3853cf697cd..8ded1a5b7b0 100644
--- a/processing/src/test/java/org/apache/druid/frame/file/FrameFileTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/file/FrameFileTest.java
@@ -19,6 +19,7 @@
package org.apache.druid.frame.file;
+import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.math.IntMath;
import it.unimi.dsi.fastutil.ints.IntObjectPair;
import org.apache.druid.frame.Frame;
@@ -27,9 +28,12 @@ import org.apache.druid.frame.TestArrayCursorFactory;
import org.apache.druid.frame.read.FrameReader;
import org.apache.druid.frame.testutil.FrameSequenceBuilder;
import org.apache.druid.frame.testutil.FrameTestUtil;
+import org.apache.druid.frame.wire.FrameWireTransferable;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.guava.Sequence;
import org.apache.druid.java.util.common.guava.Sequences;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.segment.CursorFactory;
import org.apache.druid.segment.QueryableIndexCursorFactory;
import org.apache.druid.segment.RowAdapters;
@@ -50,11 +54,13 @@ import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
+import javax.annotation.Nullable;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.math.RoundingMode;
+import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.HashMap;
@@ -149,6 +155,7 @@ public class FrameFileTest extends InitializedNullHandlingTest
private final boolean partitioned;
private final AdapterType adapterType;
private final int maxMmapSize;
+ private final boolean useLegacyFrameSerialization;
private CursorFactory cursorFactory;
private int rowCount;
@@ -159,7 +166,8 @@ public class FrameFileTest extends InitializedNullHandlingTest
final int maxRowsPerFrame,
final boolean partitioned,
final AdapterType adapterType,
- final int maxMmapSize
+ final int maxMmapSize,
+ final boolean useLegacyFrameSerialization
)
{
this.frameType = frameType;
@@ -167,6 +175,7 @@ public class FrameFileTest extends InitializedNullHandlingTest
this.partitioned = partitioned;
this.adapterType = adapterType;
this.maxMmapSize = maxMmapSize;
+ this.useLegacyFrameSerialization = useLegacyFrameSerialization;
}
@Parameterized.Parameters(
@@ -174,7 +183,8 @@ public class FrameFileTest extends InitializedNullHandlingTest
+ "maxRowsPerFrame = {1}, "
+ "partitioned = {2}, "
+ "adapter = {3}, "
- + "maxMmapSize = {4}"
+ + "maxMmapSize = {4}, "
+ + "useLegacyFrameSerialization = {5}"
)
public static Iterable<Object[]> constructorFeeder()
{
@@ -193,7 +203,9 @@ public class FrameFileTest extends InitializedNullHandlingTest
}
for (int maxMmapSize : maxMmapSizes) {
- constructors.add(new Object[]{frameType, maxRowsPerFrame, partitioned, adapterType, maxMmapSize});
+ for (boolean useLegacyFrameSerialization : new boolean[]{true, false}) {
+ constructors.add(new Object[]{frameType, maxRowsPerFrame, partitioned, adapterType, maxMmapSize, useLegacyFrameSerialization});
+ }
}
}
}
@@ -203,6 +215,23 @@ public class FrameFileTest extends InitializedNullHandlingTest
return constructors;
}
+ @Nullable
+ private WireTransferable.ConcreteDeserializer makeConcreteDeserializer()
+ {
+ if (useLegacyFrameSerialization) {
+ return null;
+ } else {
+ final ObjectMapper objectMapper = new ObjectMapper();
+ return new WireTransferable.ConcreteDeserializer(
+ objectMapper,
+ Map.of(
+ ByteBuffer.wrap(StringUtils.toUtf8(FrameWireTransferable.TYPE)),
+ new FrameWireTransferable.Deserializer()
+ )
+ );
+ }
+ }
+
@Before
public void setUp() throws IOException
{
@@ -211,7 +240,7 @@ public class FrameFileTest extends InitializedNullHandlingTest
file = temporaryFolder.newFile();
try (final OutputStream out = Files.newOutputStream(file.toPath())) {
- final FrameFileKey frameFileKey = new FrameFileKey(adapterType, frameType, maxRowsPerFrame, partitioned);
+ final FrameFileKey frameFileKey = new FrameFileKey(adapterType, frameType, maxRowsPerFrame, partitioned, useLegacyFrameSerialization);
final byte[] frameFileBytes = FRAME_FILES.computeIfAbsent(frameFileKey, FrameFileTest::computeFrameFile);
out.write(frameFileBytes);
}
@@ -240,25 +269,25 @@ public class FrameFileTest extends InitializedNullHandlingTest
}
@Test
- public void test_frame_first() throws IOException
+ public void test_rac_first() throws IOException
{
try (final FrameFile frameFile = FrameFile.open(file, maxMmapSize, null)) {
// Skip test for empty files.
Assume.assumeThat(frameFile.numFrames(), Matchers.greaterThan(0));
- final Frame firstFrame = frameFile.frame(0);
+ final Frame firstFrame = frameFile.rac(0, makeConcreteDeserializer()).as(Frame.class);
Assert.assertEquals(Math.min(rowCount, maxRowsPerFrame), firstFrame.numRows());
}
}
@Test
- public void test_frame_last() throws IOException
+ public void test_rac_last() throws IOException
{
try (final FrameFile frameFile = FrameFile.open(file, maxMmapSize, null)) {
// Skip test for empty files.
Assume.assumeThat(frameFile.numFrames(), Matchers.greaterThan(0));
- final Frame lastFrame = frameFile.frame(frameFile.numFrames() - 1);
+ final Frame lastFrame = frameFile.rac(frameFile.numFrames() - 1, makeConcreteDeserializer()).as(Frame.class);
Assert.assertEquals(
rowCount % maxRowsPerFrame != 0
? rowCount % maxRowsPerFrame
@@ -269,34 +298,34 @@ public class FrameFileTest extends InitializedNullHandlingTest
}
@Test
- public void test_frame_outOfBoundsNegative() throws IOException
+ public void test_rac_outOfBoundsNegative() throws IOException
{
try (final FrameFile frameFile = FrameFile.open(file, maxMmapSize, null)) {
expectedException.expect(IllegalArgumentException.class);
- expectedException.expectMessage("Frame [-1] out of bounds");
- frameFile.frame(-1);
+ expectedException.expectMessage("Batch[-1] out of bounds");
+ frameFile.rac(-1, null);
}
}
@Test
- public void test_frame_outOfBoundsTooLarge() throws IOException
+ public void test_rac_outOfBoundsTooLarge() throws IOException
{
try (final FrameFile frameFile = FrameFile.open(file, maxMmapSize, null)) {
expectedException.expect(IllegalArgumentException.class);
- expectedException.expectMessage(StringUtils.format("Frame [%,d] out of bounds", frameFile.numFrames()));
- frameFile.frame(frameFile.numFrames());
+ expectedException.expectMessage(StringUtils.format("Batch[%,d] out of bounds", frameFile.numFrames()));
+ frameFile.rac(frameFile.numFrames(), null);
}
}
@Test
- public void test_frame_readAllDataViaCursorFactory() throws IOException
+ public void test_rac_readAllDataViaCursorFactory() throws IOException
{
final FrameReader frameReader = FrameReader.create(cursorFactory.getRowSignature());
try (final FrameFile frameFile = FrameFile.open(file, maxMmapSize, null)) {
final Sequence<List<Object>> frameFileRows = Sequences.concat(
() -> IntStream.range(0, frameFile.numFrames())
- .mapToObj(frameFile::frame)
+ .mapToObj(i -> frameFile.rac(i, makeConcreteDeserializer()).as(Frame.class))
.map(frameReader::makeCursorFactory)
.map(FrameTestUtil::readRowsFromCursorFactoryWithRowNumber)
.iterator()
@@ -417,6 +446,7 @@ public class FrameFileTest extends InitializedNullHandlingTest
private static byte[] computeFrameFile(final FrameFileKey frameFileKey)
{
final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ final WireTransferableContext wireTransferableContext = frameFileKey.makeWireTransferableContext();
try {
if (frameFileKey.partitioned) {
@@ -443,7 +473,8 @@ public class FrameFileTest extends InitializedNullHandlingTest
}
}
),
- baos
+ baos,
+ wireTransferableContext
);
} else {
FrameTestUtil.writeFrameFile(
@@ -451,7 +482,8 @@ public class FrameFileTest extends InitializedNullHandlingTest
.frameType(frameFileKey.frameType)
.maxRowsPerFrame(frameFileKey.maxRowsPerFrame)
.frames(),
- baos
+ baos,
+ wireTransferableContext
);
}
}
@@ -471,13 +503,39 @@ public class FrameFileTest extends InitializedNullHandlingTest
final FrameType frameType;
final int maxRowsPerFrame;
final boolean partitioned;
-
- public FrameFileKey(AdapterType adapterType, FrameType frameType, int maxRowsPerFrame, boolean partitioned)
+ final boolean useLegacyFrameSerialization;
+
+ public FrameFileKey(
+ AdapterType adapterType,
+ FrameType frameType,
+ int maxRowsPerFrame,
+ boolean partitioned,
+ boolean useLegacyFrameSerialization
+ )
{
this.adapterType = adapterType;
this.frameType = frameType;
this.maxRowsPerFrame = maxRowsPerFrame;
this.partitioned = partitioned;
+ this.useLegacyFrameSerialization = useLegacyFrameSerialization;
+ }
+
+ @Nullable
+ WireTransferableContext makeWireTransferableContext()
+ {
+ if (useLegacyFrameSerialization) {
+ return null;
+ } else {
+ final ObjectMapper objectMapper = new ObjectMapper();
+ final WireTransferable.ConcreteDeserializer deserializer = new WireTransferable.ConcreteDeserializer(
+ objectMapper,
+ Map.of(
+ ByteBuffer.wrap(StringUtils.toUtf8(FrameWireTransferable.TYPE)),
+ new FrameWireTransferable.Deserializer()
+ )
+ );
+ return new WireTransferableContext(objectMapper, deserializer, false);
+ }
}
@Override
@@ -492,6 +550,7 @@ public class FrameFileTest extends InitializedNullHandlingTest
FrameFileKey that = (FrameFileKey) o;
return maxRowsPerFrame == that.maxRowsPerFrame
&& partitioned == that.partitioned
+ && useLegacyFrameSerialization == that.useLegacyFrameSerialization
&& adapterType == that.adapterType
&& frameType == that.frameType;
}
@@ -499,7 +558,7 @@ public class FrameFileTest extends InitializedNullHandlingTest
@Override
public int hashCode()
{
- return Objects.hash(adapterType, frameType, maxRowsPerFrame, partitioned);
+ return Objects.hash(adapterType, frameType, maxRowsPerFrame, partitioned, useLegacyFrameSerialization);
}
@Override
@@ -510,6 +569,7 @@ public class FrameFileTest extends InitializedNullHandlingTest
", frameType=" + frameType +
", maxRowsPerFrame=" + maxRowsPerFrame +
", partitioned=" + partitioned +
+ ", useLegacyFrameSerialization=" + useLegacyFrameSerialization +
'}';
}
}
diff --git a/processing/src/test/java/org/apache/druid/frame/file/FrameFileWireTransferableTest.java b/processing/src/test/java/org/apache/druid/frame/file/FrameFileWireTransferableTest.java
new file mode 100644
index 00000000000..7966e4df68f
--- /dev/null
+++ b/processing/src/test/java/org/apache/druid/frame/file/FrameFileWireTransferableTest.java
@@ -0,0 +1,114 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.frame.file;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.frame.Frame;
+import org.apache.druid.frame.FrameType;
+import org.apache.druid.frame.channel.ByteTracker;
+import org.apache.druid.frame.testutil.FrameSequenceBuilder;
+import org.apache.druid.frame.wire.FrameWireTransferable;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.common.guava.Sequence;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
+import org.apache.druid.segment.CursorFactory;
+import org.apache.druid.segment.TestIndex;
+import org.apache.druid.segment.incremental.IncrementalIndexCursorFactory;
+import org.apache.druid.testing.InitializedNullHandlingTest;
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
+import java.nio.file.Files;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Tests error cases for FrameFile with WireTransferable format.
+ * Basic read/write functionality is covered by {@link FrameFileTest} with useLegacyFrameSerialization parameter.
+ */
+public class FrameFileWireTransferableTest extends InitializedNullHandlingTest
+{
+ @Rule
+ public TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ @Test
+ public void testRacMethodThrowsWithoutDeserializerForRacEntry() throws IOException
+ {
+ // Test that rac(int, null) throws when encountering a RAC entry
+ final CursorFactory cursorFactory = new IncrementalIndexCursorFactory(
+ TestIndex.getNoRollupIncrementalTestIndex()
+ );
+
+ // Create wireTransferableContext for RAC serialization
+ final ObjectMapper smileMapper = new ObjectMapper();
+ final WireTransferable.ConcreteDeserializer deserializer = new WireTransferable.ConcreteDeserializer(
+ smileMapper,
+ Map.of(
+ ByteBuffer.wrap(StringUtils.toUtf8(FrameWireTransferable.TYPE)),
+ new FrameWireTransferable.Deserializer()
+ )
+ );
+ final WireTransferableContext wireTransferableContext = new WireTransferableContext(
+ smileMapper,
+ deserializer,
+ false // useLegacyFrameSerialization = false to enable RAC serialization
+ );
+
+ // Build frames
+ final Sequence<Frame> frames = FrameSequenceBuilder.fromCursorFactory(cursorFactory)
+ .frameType(FrameType.latestColumnar())
+ .maxRowsPerFrame(50)
+ .frames();
+ final List<Frame> frameList = frames.toList();
+
+ // Write frame file with useWireTransferableForFrames = true (RAC format)
+ final File file = temporaryFolder.newFile();
+ try (final FrameFileWriter writer = FrameFileWriter.open(
+ Channels.newChannel(Files.newOutputStream(file.toPath())),
+ null,
+ ByteTracker.unboundedTracker(),
+ wireTransferableContext
+ )) {
+ for (Frame frame : frameList) {
+ writer.write(frame.asRAC(), FrameFileWriter.NO_PARTITION);
+ }
+ }
+
+ // Try to read using rac(int, null) without deserializer - should throw
+ try (final FrameFile frameFile = FrameFile.open(file, null)) {
+ Assert.assertEquals(frameList.size(), frameFile.numFrames());
+
+ // This should throw DruidException because no deserializer was provided
+ Assert.assertThrows(
+ "Should throw DruidException when reading RAC entry without deserializer",
+ DruidException.class,
+ () -> frameFile.rac(0, null)
+ );
+ }
+ }
+}
diff --git a/processing/src/test/java/org/apache/druid/frame/file/FrameFileWriterTest.java b/processing/src/test/java/org/apache/druid/frame/file/FrameFileWriterTest.java
index e78705ea56f..6a88c1e5818 100644
--- a/processing/src/test/java/org/apache/druid/frame/file/FrameFileWriterTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/file/FrameFileWriterTest.java
@@ -24,6 +24,7 @@ import org.apache.druid.frame.FrameType;
import org.apache.druid.frame.allocation.ArenaMemoryAllocator;
import org.apache.druid.frame.channel.ByteTracker;
import org.apache.druid.frame.testutil.FrameSequenceBuilder;
+import org.apache.druid.frame.testutil.FrameTestUtil;
import org.apache.druid.java.util.common.guava.Sequence;
import org.apache.druid.segment.TestIndex;
import org.apache.druid.segment.incremental.IncrementalIndexCursorFactory;
@@ -55,14 +56,19 @@ public class FrameFileWriterTest extends InitializedNullHandlingTest
.frames();
final File file = temporaryFolder.newFile();
- final FrameFileWriter fileWriter = FrameFileWriter.open(Files.newByteChannel(
- file.toPath(),
- StandardOpenOption.WRITE
- ), null, ByteTracker.unboundedTracker());
+ final FrameFileWriter fileWriter = FrameFileWriter.open(
+ Files.newByteChannel(
+ file.toPath(),
+ StandardOpenOption.WRITE
+ ),
+ null,
+ ByteTracker.unboundedTracker(),
+ FrameTestUtil.WT_CONTEXT_LEGACY
+ );
frames.forEach(frame -> {
try {
- fileWriter.writeFrame(frame, FrameFileWriter.NO_PARTITION);
+ fileWriter.write(frame.asRAC(), FrameFileWriter.NO_PARTITION);
}
catch (IOException e) {
throw new RuntimeException(e);
diff --git a/processing/src/test/java/org/apache/druid/frame/processor/ComposingOutputChannelFactoryTest.java b/processing/src/test/java/org/apache/druid/frame/processor/ComposingOutputChannelFactoryTest.java
index 235ba06691e..90cf8f428ad 100644
--- a/processing/src/test/java/org/apache/druid/frame/processor/ComposingOutputChannelFactoryTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/processor/ComposingOutputChannelFactoryTest.java
@@ -21,8 +21,8 @@ package org.apache.druid.frame.processor;
import com.google.common.collect.ImmutableList;
import org.apache.druid.frame.channel.ByteTracker;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.WritableFrameChannel;
+import org.apache.druid.frame.testutil.FrameTestUtil;
import org.junit.Assert;
import org.junit.ClassRule;
import org.junit.Test;
@@ -43,8 +43,8 @@ public class ComposingOutputChannelFactoryTest extends OutputChannelFactoryTest
ImmutableList.of(
// TODO : currently hardcoded 256k since it allows one frame to be written to each factory
// nicer to do that automatically
- new FileOutputChannelFactory(folder.newFolder(), 100, new ByteTracker(256_000)),
- new FileOutputChannelFactory(folder.newFolder(), 100, new ByteTracker(256_000))
+ new FileOutputChannelFactory(folder.newFolder(), 100, new ByteTracker(256_000), FrameTestUtil.WT_CONTEXT_LEGACY),
+ new FileOutputChannelFactory(folder.newFolder(), 100, new ByteTracker(256_000), FrameTestUtil.WT_CONTEXT_LEGACY)
),
100
),
@@ -57,7 +57,7 @@ public class ComposingOutputChannelFactoryTest extends OutputChannelFactoryTest
{
ComposingOutputChannelFactory outputChannelFactory = new ComposingOutputChannelFactory(
ImmutableList.of(
- new FileOutputChannelFactory(folder.newFolder(), 100, new ByteTracker(1)),
+ new FileOutputChannelFactory(folder.newFolder(), 100, new ByteTracker(1), FrameTestUtil.WT_CONTEXT_LEGACY),
new ThrowingOutputChannelFactory() // adding this to check if it gets called
),
100
@@ -69,7 +69,7 @@ public class ComposingOutputChannelFactoryTest extends OutputChannelFactoryTest
writableFrameChannel.writabilityFuture().get();
Assert.assertThrows(
UnsupportedOperationException.class,
- () -> writableFrameChannel.write(new FrameWithPartition(frame, 1))
+ () -> writableFrameChannel.write(frame, 1)
);
}
@@ -80,7 +80,7 @@ public class ComposingOutputChannelFactoryTest extends OutputChannelFactoryTest
// can handle the test data frames
ComposingOutputChannelFactory outputChannelFactory = new ComposingOutputChannelFactory(
ImmutableList.of(
- new FileOutputChannelFactory(folder.newFolder(), 100, new ByteTracker(1_000_000)),
+ new FileOutputChannelFactory(folder.newFolder(), 100, new ByteTracker(1_000_000), FrameTestUtil.WT_CONTEXT_LEGACY),
new ThrowingOutputChannelFactory()
),
100
@@ -90,7 +90,7 @@ public class ComposingOutputChannelFactoryTest extends OutputChannelFactoryTest
Assert.assertEquals(1, channel.getPartitionNumber());
WritableFrameChannel writableFrameChannel = channel.getWritableChannel();
writableFrameChannel.writabilityFuture().get();
- writableFrameChannel.write(new FrameWithPartition(frame, 1));
+ writableFrameChannel.write(frame, 1);
writableFrameChannel.close();
verifySingleFrameReadableChannel(channel.getReadableChannel(), sourceCursorFactory);
diff --git a/processing/src/test/java/org/apache/druid/frame/processor/FileOutputChannelFactoryTest.java b/processing/src/test/java/org/apache/druid/frame/processor/FileOutputChannelFactoryTest.java
index 1a0038479d3..8b26641e57c 100644
--- a/processing/src/test/java/org/apache/druid/frame/processor/FileOutputChannelFactoryTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/processor/FileOutputChannelFactoryTest.java
@@ -19,6 +19,7 @@
package org.apache.druid.frame.processor;
+import org.apache.druid.frame.testutil.FrameTestUtil;
import org.junit.ClassRule;
import org.junit.rules.TemporaryFolder;
@@ -31,6 +32,6 @@ public class FileOutputChannelFactoryTest extends OutputChannelFactoryTest
public FileOutputChannelFactoryTest() throws IOException
{
- super(new FileOutputChannelFactory(folder.newFolder(), 100, null), 100);
+ super(new FileOutputChannelFactory(folder.newFolder(), 100, null, FrameTestUtil.WT_CONTEXT_LEGACY), 100);
}
}
diff --git a/processing/src/test/java/org/apache/druid/frame/processor/FrameProcessorExecutorTest.java b/processing/src/test/java/org/apache/druid/frame/processor/FrameProcessorExecutorTest.java
index 60c2e1135b8..0d5029e56fd 100644
--- a/processing/src/test/java/org/apache/druid/frame/processor/FrameProcessorExecutorTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/processor/FrameProcessorExecutorTest.java
@@ -138,7 +138,8 @@ public class FrameProcessorExecutorTest
FrameFileWriter.open(
Channels.newChannel(Files.newOutputStream(outFile.toPath())),
null,
- ByteTracker.unboundedTracker()
+ ByteTracker.unboundedTracker(),
+ FrameTestUtil.WT_CONTEXT_LEGACY
)
)
);
@@ -152,7 +153,7 @@ public class FrameProcessorExecutorTest
Assert.assertEquals(
index.numRows() * 2,
FrameTestUtil.readRowsFromFrameChannel(
- new ReadableFileFrameChannel(FrameFile.open(outFile, null)),
+ new ReadableFileFrameChannel(FrameFile.open(outFile, null), null),
FrameReader.create(cursorFactory.getRowSignature())
).toList().size()
);
@@ -206,7 +207,7 @@ public class FrameProcessorExecutorTest
final RuntimeException readException = Assert.assertThrows(
RuntimeException.class,
- outReadableChannel::read
+ outReadableChannel::readFrame
);
MatcherAssert.assertThat(
@@ -529,7 +530,8 @@ public class FrameProcessorExecutorTest
FrameFileWriter.open(
Channels.newChannel(Files.newOutputStream(files.get(i).toPath())),
null,
- ByteTracker.unboundedTracker()
+ ByteTracker.unboundedTracker(),
+ FrameTestUtil.WT_CONTEXT_LEGACY
)
);
}
@@ -543,7 +545,7 @@ public class FrameProcessorExecutorTest
public void accept(Frame frame)
{
try {
- writers.get(j % writers.size()).writeFrame(frame, FrameFileWriter.NO_PARTITION);
+ writers.get(j % writers.size()).write(frame.asRAC(), FrameFileWriter.NO_PARTITION);
}
catch (IOException e) {
throw new RuntimeException(e);
@@ -572,7 +574,7 @@ public class FrameProcessorExecutorTest
private static ReadableFrameChannel openFileChannel(final File file)
{
try {
- return new ReadableFileFrameChannel(FrameFile.open(file, null));
+ return new ReadableFileFrameChannel(FrameFile.open(file, null), null);
}
catch (IOException e) {
throw new RuntimeException(e);
diff --git a/processing/src/test/java/org/apache/druid/frame/processor/OutputChannelFactoryTest.java b/processing/src/test/java/org/apache/druid/frame/processor/OutputChannelFactoryTest.java
index 0e472cd29fc..80ab2519daf 100644
--- a/processing/src/test/java/org/apache/druid/frame/processor/OutputChannelFactoryTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/processor/OutputChannelFactoryTest.java
@@ -23,7 +23,6 @@ import com.google.common.collect.Iterables;
import org.apache.druid.error.DruidException;
import org.apache.druid.frame.Frame;
import org.apache.druid.frame.FrameType;
-import org.apache.druid.frame.channel.FrameWithPartition;
import org.apache.druid.frame.channel.PartitionedReadableFrameChannel;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
@@ -72,7 +71,7 @@ public abstract class OutputChannelFactoryTest extends InitializedNullHandlingTe
// write data to the channel
WritableFrameChannel writableFrameChannel = channel.getWritableChannel();
writableFrameChannel.writabilityFuture().get();
- writableFrameChannel.write(new FrameWithPartition(frame, 1));
+ writableFrameChannel.write(frame, 1);
writableFrameChannel.close();
// read back data from the channel
@@ -114,7 +113,7 @@ public abstract class OutputChannelFactoryTest extends InitializedNullHandlingTe
WritableFrameChannel writableFrameChannel = channel.getWritableChannel();
writableFrameChannel.writabilityFuture().get();
for (int partition : partitions) {
- writableFrameChannel.write(new FrameWithPartition(frame, partition));
+ writableFrameChannel.write(frame, partition);
}
writableFrameChannel.close();
@@ -142,7 +141,7 @@ public abstract class OutputChannelFactoryTest extends InitializedNullHandlingTe
break;
}
}
- Frame readbackFrame = readableFrameChannel.read();
+ Frame readbackFrame = readableFrameChannel.readFrame();
readableFrameChannel.readabilityFuture().get();
Assert.assertFalse(readableFrameChannel.canRead());
Assert.assertTrue(readableFrameChannel.isFinished());
diff --git a/processing/src/test/java/org/apache/druid/frame/processor/ReadableFileFrameChannelTest.java b/processing/src/test/java/org/apache/druid/frame/processor/ReadableFileFrameChannelTest.java
index f6bb3fa1953..54e09de615d 100644
--- a/processing/src/test/java/org/apache/druid/frame/processor/ReadableFileFrameChannelTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/processor/ReadableFileFrameChannelTest.java
@@ -77,7 +77,7 @@ public class ReadableFileFrameChannelTest extends InitializedNullHandlingTest
@Test
public void test_fullFile()
{
- final ReadableFileFrameChannel channel = new ReadableFileFrameChannel(frameFile);
+ final ReadableFileFrameChannel channel = new ReadableFileFrameChannel(frameFile, null);
Assert.assertTrue(channel.isEntireFile());
FrameTestUtil.assertRowsEqual(
@@ -91,7 +91,7 @@ public class ReadableFileFrameChannelTest extends InitializedNullHandlingTest
@Test
public void test_partialFile()
{
- final ReadableFileFrameChannel channel = new ReadableFileFrameChannel(frameFile, 1, 2);
+ final ReadableFileFrameChannel channel = new ReadableFileFrameChannel(frameFile, 1, 2, null);
Assert.assertFalse(channel.isEntireFile());
FrameTestUtil.assertRowsEqual(
diff --git a/processing/src/test/java/org/apache/druid/frame/processor/SuperSorterTest.java b/processing/src/test/java/org/apache/druid/frame/processor/SuperSorterTest.java
index 44c696e508b..680df7a3eb4 100644
--- a/processing/src/test/java/org/apache/druid/frame/processor/SuperSorterTest.java
+++ b/processing/src/test/java/org/apache/druid/frame/processor/SuperSorterTest.java
@@ -130,8 +130,8 @@ public class SuperSorterTest
outputPartitionsFuture,
exec,
FrameProcessorDecorator.NONE,
- new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null),
- new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null),
+ new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null, FrameTestUtil.WT_CONTEXT_LEGACY),
+ new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null, FrameTestUtil.WT_CONTEXT_LEGACY),
FrameType.latestRowBased(),
2,
2,
@@ -167,8 +167,8 @@ public class SuperSorterTest
Futures.immediateFuture(ClusterByPartitions.oneUniversalPartition()),
exec,
FrameProcessorDecorator.NONE,
- new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null),
- new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null),
+ new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null, FrameTestUtil.WT_CONTEXT_LEGACY),
+ new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null, FrameTestUtil.WT_CONTEXT_LEGACY),
FrameType.latestRowBased(),
2,
2,
@@ -203,8 +203,8 @@ public class SuperSorterTest
Futures.immediateFuture(ClusterByPartitions.oneUniversalPartition()),
exec,
FrameProcessorDecorator.NONE,
- new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null),
- new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null),
+ new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null, FrameTestUtil.WT_CONTEXT_LEGACY),
+ new FileOutputChannelFactory(tempFolder, FRAME_SIZE, null, FrameTestUtil.WT_CONTEXT_LEGACY),
FrameType.latestRowBased(),
2,
2,
@@ -430,11 +430,11 @@ public class SuperSorterTest
final File tempFolder = temporaryFolder.newFolder();
final OutputChannelFactory outputChannelFactory = isComposedStorage ? new ComposingOutputChannelFactory(
ImmutableList.of(
- new FileOutputChannelFactory(new File(tempFolder, "1"), maxBytesPerFrame, null),
- new FileOutputChannelFactory(new File(tempFolder, "2"), maxBytesPerFrame, null)
+ new FileOutputChannelFactory(new File(tempFolder, "1"), maxBytesPerFrame, null, FrameTestUtil.WT_CONTEXT_LEGACY),
+ new FileOutputChannelFactory(new File(tempFolder, "2"), maxBytesPerFrame, null, FrameTestUtil.WT_CONTEXT_LEGACY)
),
maxBytesPerFrame
- ) : new FileOutputChannelFactory(tempFolder, maxBytesPerFrame, null);
+ ) : new FileOutputChannelFactory(tempFolder, maxBytesPerFrame, null, FrameTestUtil.WT_CONTEXT_LEGACY);
final RowKeyReader keyReader = clusterBy.keyReader(signature, outputFrameType);
final Comparator<RowKey> keyComparator = clusterBy.keyComparator(signature);
final SettableFuture<ClusterByPartitions> clusterByPartitionsFuture = SettableFuture.create();
@@ -451,7 +451,7 @@ public class SuperSorterTest
clusterByPartitionsFuture,
exec,
FrameProcessorDecorator.NONE,
- makeOutputChannelFactory(new FileOutputChannelFactory(tempFolder, maxBytesPerFrame, null)),
+ makeOutputChannelFactory(new FileOutputChannelFactory(tempFolder, maxBytesPerFrame, null, FrameTestUtil.WT_CONTEXT_LEGACY)),
makeOutputChannelFactory(outputChannelFactory),
outputFrameType,
maxActiveProcessors,
diff --git a/processing/src/test/java/org/apache/druid/frame/processor/test/AlwaysAsyncReadableFrameChannel.java b/processing/src/test/java/org/apache/druid/frame/processor/test/AlwaysAsyncReadableFrameChannel.java
index 8ff10aeb7b0..b602180ee85 100644
--- a/processing/src/test/java/org/apache/druid/frame/processor/test/AlwaysAsyncReadableFrameChannel.java
+++ b/processing/src/test/java/org/apache/druid/frame/processor/test/AlwaysAsyncReadableFrameChannel.java
@@ -20,12 +20,12 @@
package org.apache.druid.frame.processor.test;
import com.google.common.util.concurrent.ListenableFuture;
-import org.apache.druid.frame.Frame;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
/**
- * Wraps an underlying channel and forces an async style of reading. After each call to {@link #read()}, the
+ * Wraps an underlying channel and forces an async style of reading. After each call to {@link #readFrame()}, the
* {@link #canRead()} and {@link #isFinished()} methods return false until {@link #readabilityFuture()} is called.
*/
public class AlwaysAsyncReadableFrameChannel implements ReadableFrameChannel
@@ -59,7 +59,7 @@ public class AlwaysAsyncReadableFrameChannel implements ReadableFrameChannel
}
@Override
- public Frame read()
+ public RowsAndColumns read()
{
if (defer) {
throw new ISE("Cannot call read() while deferred");
diff --git a/processing/src/test/java/org/apache/druid/frame/processor/test/ChompingFrameProcessor.java b/processing/src/test/java/org/apache/druid/frame/processor/test/ChompingFrameProcessor.java
index 67ec46c3776..5a44f2db703 100644
--- a/processing/src/test/java/org/apache/druid/frame/processor/test/ChompingFrameProcessor.java
+++ b/processing/src/test/java/org/apache/druid/frame/processor/test/ChompingFrameProcessor.java
@@ -76,7 +76,7 @@ public class ChompingFrameProcessor implements FrameProcessor<Long>
if (channel.isFinished()) {
awaitSet.remove(channelNumber);
} else {
- // Call valueOrThrow() for side effects: we want to throw any exceptions that come in on the channel.
+ // Call readRAC() for side effects: we want to throw any exceptions that come in on the channel.
channel.read();
didReadFrame.countDown();
numFrames++;
diff --git a/processing/src/test/java/org/apache/druid/frame/processor/test/SuperBlasterFrameProcessor.java b/processing/src/test/java/org/apache/druid/frame/processor/test/SuperBlasterFrameProcessor.java
index bc8790d4671..8a00c488795 100644
--- a/processing/src/test/java/org/apache/druid/frame/processor/test/SuperBlasterFrameProcessor.java
+++ b/processing/src/test/java/org/apache/druid/frame/processor/test/SuperBlasterFrameProcessor.java
@@ -21,12 +21,12 @@ package org.apache.druid.frame.processor.test;
import it.unimi.dsi.fastutil.ints.IntSet;
import it.unimi.dsi.fastutil.ints.IntSets;
-import org.apache.druid.frame.Frame;
import org.apache.druid.frame.channel.ReadableFrameChannel;
import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.processor.FrameProcessor;
import org.apache.druid.frame.processor.FrameProcessors;
import org.apache.druid.frame.processor.ReturnOrAwait;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
import java.io.IOException;
import java.util.List;
@@ -93,11 +93,11 @@ public class SuperBlasterFrameProcessor implements FrameProcessor<Long>
final ReadableFrameChannel inChannel = inChannels.get(channelNumber);
if (!inChannel.isFinished()) {
- final Frame frame = inChannel.read();
- rowsRead += frame.numRows();
+ final RowsAndColumns rac = inChannel.read();
+ rowsRead += rac.numRows();
for (WritableFrameChannel outChannel : outChannels) {
- outChannel.write(frame);
+ outChannel.write(rac);
}
// Only permitted to write one frame to each output channel.
diff --git a/processing/src/test/java/org/apache/druid/frame/testutil/FrameTestUtil.java b/processing/src/test/java/org/apache/druid/frame/testutil/FrameTestUtil.java
index 50b4e80dba0..9ec348d7ea6 100644
--- a/processing/src/test/java/org/apache/druid/frame/testutil/FrameTestUtil.java
+++ b/processing/src/test/java/org/apache/druid/frame/testutil/FrameTestUtil.java
@@ -35,6 +35,8 @@ import org.apache.druid.java.util.common.guava.Sequence;
import org.apache.druid.java.util.common.guava.Sequences;
import org.apache.druid.query.OrderBy;
import org.apache.druid.query.dimension.DefaultDimensionSpec;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.apache.druid.query.rowsandcols.serde.WireTransferableContext;
import org.apache.druid.segment.ColumnProcessors;
import org.apache.druid.segment.ColumnSelectorFactory;
import org.apache.druid.segment.ColumnValueSelector;
@@ -71,6 +73,13 @@ public class FrameTestUtil
{
public static final String ROW_NUMBER_COLUMN = "__row_number";
+ /**
+ * A context that uses legacy frame serde, and does not support {@link WireTransferable}.
+ * Use this for tests that only need to exercise the legacy frame serde path.
+ */
+ public static final WireTransferableContext WT_CONTEXT_LEGACY =
+ new WireTransferableContext(null, null, true);
+
private FrameTestUtil()
{
// No instantiation.
@@ -78,18 +87,33 @@ public class FrameTestUtil
public static File writeFrameFile(final Sequence<Frame> frames, final File file) throws IOException
{
- writeFrameFile(frames, Files.newOutputStream(file.toPath()));
+ return writeFrameFile(frames, file, null);
+ }
+
+ public static File writeFrameFile(
+ final Sequence<Frame> frames,
+ final File file,
+ @Nullable final WireTransferableContext wireTransferableContext
+ ) throws IOException
+ {
+ writeFrameFile(frames, Files.newOutputStream(file.toPath()), wireTransferableContext);
return file;
}
- public static void writeFrameFile(final Sequence<Frame> frames, final OutputStream out) throws IOException
+ public static void writeFrameFile(
+ final Sequence<Frame> frames,
+ final OutputStream out,
+ @Nullable final WireTransferableContext wireTransferableContext
+ ) throws IOException
{
+ final WireTransferableContext context =
+ wireTransferableContext != null ? wireTransferableContext : WT_CONTEXT_LEGACY;
try (final FrameFileWriter writer =
- FrameFileWriter.open(Channels.newChannel(out), null, ByteTracker.unboundedTracker())) {
+ FrameFileWriter.open(Channels.newChannel(out), null, ByteTracker.unboundedTracker(), context)) {
frames.forEach(
frame -> {
try {
- writer.writeFrame(frame, FrameFileWriter.NO_PARTITION);
+ writer.write(frame.asRAC(), FrameFileWriter.NO_PARTITION);
}
catch (IOException e) {
throw new RuntimeException(e);
@@ -101,18 +125,22 @@ public class FrameTestUtil
public static void writeFrameFileWithPartitions(
final Sequence<IntObjectPair<Frame>> framesWithPartitions,
- final OutputStream out
+ final OutputStream out,
+ @Nullable final WireTransferableContext wireTransferableContext
) throws IOException
{
+ final WireTransferableContext context =
+ wireTransferableContext != null ? wireTransferableContext : WT_CONTEXT_LEGACY;
try (final FrameFileWriter writer = FrameFileWriter.open(
Channels.newChannel(out),
null,
- ByteTracker.unboundedTracker()
+ ByteTracker.unboundedTracker(),
+ context
)) {
framesWithPartitions.forEach(
frameWithPartition -> {
try {
- writer.writeFrame(frameWithPartition.right(), frameWithPartition.leftInt());
+ writer.write(frameWithPartition.right().asRAC(), frameWithPartition.leftInt());
}
catch (IOException e) {
throw new RuntimeException(e);
@@ -214,8 +242,8 @@ public class FrameTestUtil
{
return new FrameChannelSequence(channel)
.flatMap(
- frame -> {
- final CursorHolder cursorHolder = frameReader.makeCursorFactory(frame)
+ rac -> {
+ final CursorHolder cursorHolder = frameReader.makeCursorFactory(rac.as(Frame.class))
.makeCursorHolder(CursorBuildSpec.FULL_SCAN);
final Cursor cursor = cursorHolder.asCursor();
if (cursor == null) {
diff --git a/processing/src/test/java/org/apache/druid/frame/wire/FrameCompressionTest.java b/processing/src/test/java/org/apache/druid/frame/wire/FrameCompressionTest.java
new file mode 100644
index 00000000000..d2c7f0c5830
--- /dev/null
+++ b/processing/src/test/java/org/apache/druid/frame/wire/FrameCompressionTest.java
@@ -0,0 +1,302 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.frame.wire;
+
+import org.apache.datasketches.memory.Memory;
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.StringUtils;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Random;
+
+public class FrameCompressionTest
+{
+ private static final int SEED = 12345;
+
+ @Test
+ public void test_compressionBufferSize()
+ {
+ // Verify buffer size is at least envelope size plus some room for data.
+ final int size = FrameCompression.compressionBufferSize(100);
+ Assert.assertTrue(size >= FrameCompression.COMPRESSED_DATA_ENVELOPE_SIZE + 100);
+ }
+
+ @Test
+ public void test_compressionBufferSize_zero()
+ {
+ final int size = FrameCompression.compressionBufferSize(0);
+ Assert.assertTrue(size >= FrameCompression.COMPRESSED_DATA_ENVELOPE_SIZE);
+ }
+
+ @Test
+ public void test_compressAndDecompress_byteBuffer_smallData()
+ {
+ final byte[] data = StringUtils.toUtf8("Hello, World!");
+ final ByteBuffer dataBuffer = ByteBuffer.wrap(data);
+ final ByteBuffer compressionBuffer = ByteBuffer.allocate(FrameCompression.compressionBufferSize(data.length));
+
+ final ByteBuffer compressed = FrameCompression.compress(dataBuffer, compressionBuffer);
+ Assert.assertEquals(0, compressed.position());
+ Assert.assertTrue(compressed.limit() > 0);
+
+ final Memory compressedMemory = Memory.wrap(compressed);
+ final byte[] decompressed = FrameCompression.decompress(compressedMemory, 0, compressed.limit());
+
+ Assert.assertArrayEquals(data, decompressed);
+ }
+
+ @Test
+ public void test_compressAndDecompress_byteBuffer_emptyData()
+ {
+ final byte[] data = new byte[0];
+ final ByteBuffer dataBuffer = ByteBuffer.wrap(data);
+ final ByteBuffer compressionBuffer = ByteBuffer.allocate(FrameCompression.compressionBufferSize(data.length));
+
+ final ByteBuffer compressed = FrameCompression.compress(dataBuffer, compressionBuffer);
+ Assert.assertEquals(0, compressed.position());
+ Assert.assertTrue(compressed.limit() > 0);
+
+ final Memory compressedMemory = Memory.wrap(compressed);
+ final byte[] decompressed = FrameCompression.decompress(compressedMemory, 0, compressed.limit());
+
+ Assert.assertArrayEquals(data, decompressed);
+ }
+
+ @Test
+ public void test_compressAndDecompress_byteBuffer_largeData()
+ {
+ final Random random = new Random(SEED);
+ final byte[] data = new byte[100_000];
+ random.nextBytes(data);
+
+ final ByteBuffer dataBuffer = ByteBuffer.wrap(data);
+ final ByteBuffer compressionBuffer = ByteBuffer.allocate(FrameCompression.compressionBufferSize(data.length));
+
+ final ByteBuffer compressed = FrameCompression.compress(dataBuffer, compressionBuffer);
+ Assert.assertEquals(0, compressed.position());
+ Assert.assertTrue(compressed.limit() > 0);
+
+ final Memory compressedMemory = Memory.wrap(compressed);
+ final byte[] decompressed = FrameCompression.decompress(compressedMemory, 0, compressed.limit());
+
+ Assert.assertArrayEquals(data, decompressed);
+ }
+
+ @Test
+ public void test_compressAndDecompress_byteBuffer_highlyCompressibleData()
+ {
+ // Highly compressible data (all zeros).
+ final byte[] data = new byte[10_000];
+ Arrays.fill(data, (byte) 0);
+
+ final ByteBuffer dataBuffer = ByteBuffer.wrap(data);
+ final ByteBuffer compressionBuffer = ByteBuffer.allocate(FrameCompression.compressionBufferSize(data.length));
+
+ final ByteBuffer compressed = FrameCompression.compress(dataBuffer, compressionBuffer);
+ Assert.assertEquals(0, compressed.position());
+ Assert.assertTrue(compressed.limit() > 0);
+ // Compressed data should be much smaller than original for highly compressible data.
+ Assert.assertTrue(compressed.limit() < data.length);
+
+ final Memory compressedMemory = Memory.wrap(compressed);
+ final byte[] decompressed = FrameCompression.decompress(compressedMemory, 0, compressed.limit());
+
+ Assert.assertArrayEquals(data, decompressed);
+ }
+
+ @Test
+ public void test_compressAndDecompress_memory()
+ {
+ final byte[] data = StringUtils.toUtf8("Test data for Memory compression");
+ final Memory dataMemory = Memory.wrap(data);
+ final ByteBuffer compressionBuffer = ByteBuffer.allocate(FrameCompression.compressionBufferSize(data.length));
+
+ final ByteBuffer compressed = FrameCompression.compress(dataMemory, data.length, compressionBuffer);
+ Assert.assertEquals(0, compressed.position());
+ Assert.assertTrue(compressed.limit() > 0);
+
+ final Memory compressedMemory = Memory.wrap(compressed);
+ final byte[] decompressed = FrameCompression.decompress(compressedMemory, 0, compressed.limit());
+
+ Assert.assertArrayEquals(data, decompressed);
+ }
+
+ @Test
+ public void test_compressAndDecompress_memory_partialRegion()
+ {
+ // Create a buffer with padding at the start.
+ final byte[] expected = StringUtils.toUtf8("Test data");
+ final byte[] paddedData = new byte[100 + expected.length + 100];
+ System.arraycopy(expected, 0, paddedData, 100, expected.length);
+
+ // Create a Memory region that points only to the actual data.
+ final Memory fullMemory = Memory.wrap(paddedData);
+ final Memory dataRegion = fullMemory.region(100, expected.length);
+
+ final ByteBuffer compressionBuffer = ByteBuffer.allocate(FrameCompression.compressionBufferSize(expected.length));
+ final ByteBuffer compressed = FrameCompression.compress(dataRegion, expected.length, compressionBuffer);
+
+ final Memory compressedMemory = Memory.wrap(compressed);
+ final byte[] decompressed = FrameCompression.decompress(compressedMemory, 0, compressed.limit());
+
+ Assert.assertArrayEquals(expected, decompressed);
+ }
+
+ @Test
+ public void test_compress_bufferTooSmall()
+ {
+ final byte[] data = StringUtils.toUtf8("Test data");
+ final ByteBuffer dataBuffer = ByteBuffer.wrap(data);
+ final ByteBuffer tooSmallBuffer = ByteBuffer.allocate(10);
+
+ final ISE exception = Assert.assertThrows(ISE.class, () ->
+ FrameCompression.compress(dataBuffer, tooSmallBuffer)
+ );
+ Assert.assertTrue(exception.getMessage().contains("Compression buffer too small"));
+ }
+
+ @Test
+ public void test_decompress_regionTooShort()
+ {
+ final ByteBuffer tooShort = ByteBuffer.allocate(FrameCompression.COMPRESSED_DATA_ENVELOPE_SIZE - 1);
+ final Memory memory = Memory.wrap(tooShort);
+
+ final ISE exception = Assert.assertThrows(ISE.class, () ->
+ FrameCompression.decompress(memory, 0, tooShort.capacity())
+ );
+ Assert.assertTrue(exception.getMessage().contains("Region too short"));
+ }
+
+ @Test
+ public void test_decompress_outOfBounds()
+ {
+ final ByteBuffer buffer = ByteBuffer.allocate(100);
+ final Memory memory = Memory.wrap(buffer);
+
+ final ISE exception = Assert.assertThrows(ISE.class, () ->
+ FrameCompression.decompress(memory, 50, 100)
+ );
+ Assert.assertTrue(exception.getMessage().contains("out of bounds"));
+ }
+
+ @Test
+ public void test_decompress_checksumMismatch()
+ {
+ // First, create valid compressed data.
+ final byte[] data = StringUtils.toUtf8("Test data for checksum test");
+ final ByteBuffer dataBuffer = ByteBuffer.wrap(data);
+ final ByteBuffer compressionBuffer = ByteBuffer.allocate(FrameCompression.compressionBufferSize(data.length));
+
+ final ByteBuffer compressed = FrameCompression.compress(dataBuffer, compressionBuffer);
+ final int compressedLength = compressed.limit();
+
+ // Corrupt the checksum (last 8 bytes).
+ final byte[] corruptedData = new byte[compressedLength];
+ compressed.get(corruptedData);
+ corruptedData[compressedLength - 1] ^= 0xFF;
+
+ final Memory corruptedMemory = Memory.wrap(corruptedData);
+
+ final ISE exception = Assert.assertThrows(ISE.class, () ->
+ FrameCompression.decompress(corruptedMemory, 0, compressedLength)
+ );
+ Assert.assertTrue(exception.getMessage().contains("Checksum mismatch"));
+ }
+
+ @Test
+ public void test_decompress_atOffset()
+ {
+ // Create compressed data.
+ final byte[] data = StringUtils.toUtf8("Test data at offset");
+ final ByteBuffer dataBuffer = ByteBuffer.wrap(data);
+ final ByteBuffer compressionBuffer = ByteBuffer.allocate(FrameCompression.compressionBufferSize(data.length));
+
+ final ByteBuffer compressed = FrameCompression.compress(dataBuffer, compressionBuffer);
+ final int compressedLength = compressed.limit();
+
+ // Put the compressed data at an offset in a larger buffer.
+ final int offset = 50;
+ final byte[] paddedBuffer = new byte[offset + compressedLength + 50];
+ compressed.get(paddedBuffer, offset, compressedLength);
+
+ final Memory memory = Memory.wrap(paddedBuffer);
+ final byte[] decompressed = FrameCompression.decompress(memory, offset, compressedLength);
+
+ Assert.assertArrayEquals(data, decompressed);
+ }
+
+ @Test
+ public void test_compressAndDecompress_byteBuffer_withSlice()
+ {
+ // Test with a sliced ByteBuffer to verify position handling.
+ final byte[] fullData = new byte[200];
+ final Random random = new Random(SEED);
+ random.nextBytes(fullData);
+
+ final ByteBuffer fullBuffer = ByteBuffer.wrap(fullData);
+ fullBuffer.position(50);
+ fullBuffer.limit(150);
+ final ByteBuffer slicedBuffer = fullBuffer.slice();
+
+ final ByteBuffer compressionBuffer = ByteBuffer.allocate(FrameCompression.compressionBufferSize(100));
+ final ByteBuffer compressed = FrameCompression.compress(slicedBuffer, compressionBuffer);
+
+ final Memory compressedMemory = Memory.wrap(compressed);
+ final byte[] decompressed = FrameCompression.decompress(compressedMemory, 0, compressed.limit());
+
+ final byte[] expected = Arrays.copyOfRange(fullData, 50, 150);
+ Assert.assertArrayEquals(expected, decompressed);
+ }
+
+ @Test
+ public void test_constants()
+ {
+ // Verify constants are as expected per the javadoc.
+ // Header: 1 byte compression type + 8 bytes compressed length + 8 bytes uncompressed length = 17 bytes
+ Assert.assertEquals(17, FrameCompression.COMPRESSED_DATA_HEADER_SIZE);
+ // Trailer: 8 bytes checksum
+ Assert.assertEquals(8, FrameCompression.COMPRESSED_DATA_TRAILER_SIZE);
+ // Envelope: header + trailer
+ Assert.assertEquals(25, FrameCompression.COMPRESSED_DATA_ENVELOPE_SIZE);
+ }
+
+ @Test
+ public void test_compressAndDecompress_directByteBuffer()
+ {
+ final byte[] data = StringUtils.toUtf8("Direct buffer test");
+ final ByteBuffer directDataBuffer = ByteBuffer.allocateDirect(data.length);
+ directDataBuffer.put(data);
+ directDataBuffer.flip();
+
+ final ByteBuffer compressionBuffer = ByteBuffer.allocateDirect(FrameCompression.compressionBufferSize(data.length));
+
+ final ByteBuffer compressed = FrameCompression.compress(directDataBuffer, compressionBuffer);
+ Assert.assertEquals(0, compressed.position());
+ Assert.assertTrue(compressed.limit() > 0);
+
+ final Memory compressedMemory = Memory.wrap(compressed);
+ final byte[] decompressed = FrameCompression.decompress(compressedMemory, 0, compressed.limit());
+
+ Assert.assertArrayEquals(data, decompressed);
+ }
+}
diff --git a/processing/src/test/java/org/apache/druid/guice/DruidSecondaryModuleTest.java b/processing/src/test/java/org/apache/druid/guice/DruidSecondaryModuleTest.java
index 9bbd2de9df6..5ea346d88b6 100644
--- a/processing/src/test/java/org/apache/druid/guice/DruidSecondaryModuleTest.java
+++ b/processing/src/test/java/org/apache/druid/guice/DruidSecondaryModuleTest.java
@@ -30,12 +30,24 @@ import com.google.common.collect.ImmutableMap;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.Module;
+import org.apache.druid.frame.wire.FrameWireTransferable;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.query.rowsandcols.MapOfColumnsRowsAndColumns;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.column.IntArrayColumn;
+import org.apache.druid.query.rowsandcols.concrete.ColumnBasedFrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.ColumnBasedFrameRowsAndColumnsTest;
+import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowsAndColumnsTest;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
import org.junit.Assert;
import org.junit.Test;
import javax.annotation.Nullable;
import javax.validation.Validation;
import javax.validation.Validator;
+import java.nio.ByteBuffer;
+import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;
@@ -499,7 +511,138 @@ public class DruidSecondaryModuleTest
private static ObjectMapper makeObjectMapper(Injector injector)
{
final ObjectMapper mapper = new ObjectMapper();
- DruidSecondaryModule.setupJackson(injector, mapper);
+ DruidSecondaryModule.setupJackson(
+ injector,
+ mapper,
+ Collections.emptyMap(),
+ true
+ );
return mapper;
}
+
+ /**
+ * Tests for WireTransferable serialization and deserialization.
+ */
+ public static class WireTransferableTests
+ {
+ private static final Map<ByteBuffer, WireTransferable.Deserializer> FRAME_DESERIALIZERS =
+ Collections.singletonMap(
+ StringUtils.toUtf8ByteBuffer(FrameWireTransferable.TYPE),
+ new FrameWireTransferable.Deserializer()
+ );
+
+ @Test
+ public void testColumnBasedFrameRowsAndColumns() throws Exception
+ {
+ final ObjectMapper om = new ObjectMapper();
+ DruidSecondaryModule.attachWireTransferables(om, FRAME_DESERIALIZERS, false);
+
+ final MapOfColumnsRowsAndColumns input = MapOfColumnsRowsAndColumns.fromMap(
+ ImmutableMap.of(
+ "colA", new IntArrayColumn(new int[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}),
+ "colB", new IntArrayColumn(new int[]{4, -4, 3, -3, 4, 82, -90, 4, 0, 0})
+ ));
+
+ final ColumnBasedFrameRowsAndColumns frc = ColumnBasedFrameRowsAndColumnsTest.buildFrame(input);
+ final byte[] bytes = om.writeValueAsBytes(frc);
+
+ final ColumnBasedFrameRowsAndColumns frc2 = (ColumnBasedFrameRowsAndColumns) om.readValue(bytes, RowsAndColumns.class);
+ Assert.assertEquals(frc, frc2);
+ }
+
+ @Test
+ public void testRowBasedFrameRowsAndColumns() throws Exception
+ {
+ final ObjectMapper om = new ObjectMapper();
+ DruidSecondaryModule.attachWireTransferables(om, FRAME_DESERIALIZERS, false);
+
+ final MapOfColumnsRowsAndColumns input = MapOfColumnsRowsAndColumns.fromMap(
+ ImmutableMap.of(
+ "colA", new IntArrayColumn(new int[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}),
+ "colB", new IntArrayColumn(new int[]{4, -4, 3, -3, 4, 82, -90, 4, 0, 0})
+ ));
+
+ final RowBasedFrameRowsAndColumns frc = RowBasedFrameRowsAndColumnsTest.MAKER.apply(input);
+ final byte[] bytes = om.writeValueAsBytes(frc);
+
+ final RowBasedFrameRowsAndColumns frc2 = (RowBasedFrameRowsAndColumns) om.readValue(bytes, RowsAndColumns.class);
+ Assert.assertEquals(frc, frc2);
+ }
+
+ @Test
+ public void testColumnBasedFrameRowsAndColumnsLegacySerialization() throws Exception
+ {
+ // Test legacy serialization format (for rolling update compatibility)
+ final ObjectMapper legacySerializer = new ObjectMapper();
+ DruidSecondaryModule.attachWireTransferables(
+ legacySerializer,
+ FRAME_DESERIALIZERS,
+ true // useLegacyFrameSerialization
+ );
+
+ // New deserializer should be able to read legacy format
+ final ObjectMapper newDeserializer = new ObjectMapper();
+ DruidSecondaryModule.attachWireTransferables(
+ newDeserializer,
+ FRAME_DESERIALIZERS,
+ false // useLegacyFrameSerialization
+ );
+
+ final MapOfColumnsRowsAndColumns input = MapOfColumnsRowsAndColumns.fromMap(
+ ImmutableMap.of(
+ "colA", new IntArrayColumn(new int[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}),
+ "colB", new IntArrayColumn(new int[]{4, -4, 3, -3, 4, 82, -90, 4, 0, 0})
+ ));
+
+ final ColumnBasedFrameRowsAndColumns frc = ColumnBasedFrameRowsAndColumnsTest.buildFrame(input);
+
+ // Serialize with legacy format
+ final byte[] bytes = legacySerializer.writeValueAsBytes(frc);
+
+ // Deserialize with new deserializer (should handle legacy format)
+ final ColumnBasedFrameRowsAndColumns frc2 = (ColumnBasedFrameRowsAndColumns) newDeserializer.readValue(
+ bytes,
+ RowsAndColumns.class
+ );
+ Assert.assertEquals(frc, frc2);
+ }
+
+ @Test
+ public void testRowBasedFrameRowsAndColumnsLegacySerialization() throws Exception
+ {
+ // Test legacy serialization format (for rolling update compatibility)
+ final ObjectMapper legacySerializer = new ObjectMapper();
+ DruidSecondaryModule.attachWireTransferables(
+ legacySerializer,
+ FRAME_DESERIALIZERS,
+ true // useLegacyFrameSerialization
+ );
+
+ // New deserializer should be able to read legacy format
+ final ObjectMapper newDeserializer = new ObjectMapper();
+ DruidSecondaryModule.attachWireTransferables(
+ newDeserializer,
+ FRAME_DESERIALIZERS,
+ false // useLegacyFrameSerialization
+ );
+
+ final MapOfColumnsRowsAndColumns input = MapOfColumnsRowsAndColumns.fromMap(
+ ImmutableMap.of(
+ "colA", new IntArrayColumn(new int[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}),
+ "colB", new IntArrayColumn(new int[]{4, -4, 3, -3, 4, 82, -90, 4, 0, 0})
+ ));
+
+ final RowBasedFrameRowsAndColumns frc = RowBasedFrameRowsAndColumnsTest.MAKER.apply(input);
+
+ // Serialize with legacy format
+ final byte[] bytes = legacySerializer.writeValueAsBytes(frc);
+
+ // Deserialize with new deserializer (should handle legacy format)
+ final RowBasedFrameRowsAndColumns frc2 = (RowBasedFrameRowsAndColumns) newDeserializer.readValue(
+ bytes,
+ RowsAndColumns.class
+ );
+ Assert.assertEquals(frc, frc2);
+ }
+ }
}
diff --git a/processing/src/test/java/org/apache/druid/jackson/DefaultObjectMapperTest.java b/processing/src/test/java/org/apache/druid/jackson/DefaultObjectMapperTest.java
index ae461d26c2a..74181ebf93f 100644
--- a/processing/src/test/java/org/apache/druid/jackson/DefaultObjectMapperTest.java
+++ b/processing/src/test/java/org/apache/druid/jackson/DefaultObjectMapperTest.java
@@ -22,18 +22,12 @@ package org.apache.druid.jackson;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.exc.InvalidTypeIdException;
-import com.google.common.collect.ImmutableMap;
import org.apache.druid.java.util.common.DateTimes;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.guava.Sequence;
import org.apache.druid.java.util.common.guava.Sequences;
import org.apache.druid.java.util.common.guava.Yielders;
import org.apache.druid.query.Query;
-import org.apache.druid.query.rowsandcols.MapOfColumnsRowsAndColumns;
-import org.apache.druid.query.rowsandcols.RowsAndColumns;
-import org.apache.druid.query.rowsandcols.column.IntArrayColumn;
-import org.apache.druid.query.rowsandcols.concrete.ColumnBasedFrameRowsAndColumns;
-import org.apache.druid.query.rowsandcols.concrete.ColumnBasedFrameRowsAndColumnsTest;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
import org.junit.Assert;
@@ -41,8 +35,6 @@ import org.junit.Test;
import java.util.Arrays;
-import static org.junit.Assert.assertEquals;
-
/**
*
*/
@@ -124,22 +116,5 @@ public class DefaultObjectMapperTest
}
Assert.fail("We expect InvalidTypeIdException to be thrown");
}
-
- @Test
- public void testColumnBasedFrameRowsAndColumns() throws Exception
- {
- DefaultObjectMapper om = new DefaultObjectMapper("test");
-
- MapOfColumnsRowsAndColumns input = (MapOfColumnsRowsAndColumns.fromMap(
- ImmutableMap.of(
- "colA", new IntArrayColumn(new int[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}),
- "colB", new IntArrayColumn(new int[]{4, -4, 3, -3, 4, 82, -90, 4, 0, 0})
- )));
-
- ColumnBasedFrameRowsAndColumns frc = ColumnBasedFrameRowsAndColumnsTest.buildFrame(input);
- byte[] bytes = om.writeValueAsBytes(frc);
-
- ColumnBasedFrameRowsAndColumns frc2 = (ColumnBasedFrameRowsAndColumns) om.readValue(bytes, RowsAndColumns.class);
- assertEquals(frc, frc2);
- }
}
+
diff --git a/processing/src/test/java/org/apache/druid/query/rowsandcols/NoAsRowsAndColumns.java b/processing/src/test/java/org/apache/druid/query/rowsandcols/NoAsRowsAndColumns.java
index 16cd44e870b..78b2e992b3c 100644
--- a/processing/src/test/java/org/apache/druid/query/rowsandcols/NoAsRowsAndColumns.java
+++ b/processing/src/test/java/org/apache/druid/query/rowsandcols/NoAsRowsAndColumns.java
@@ -21,8 +21,14 @@ package org.apache.druid.query.rowsandcols;
import org.apache.druid.query.rowsandcols.column.Column;
+import javax.annotation.Nullable;
import java.util.Collection;
+/**
+ * This class exists as a test class to force validate that the default implementations of the semantic interfaces
+ * exist. As it is named "NoAs" it is intentionally blocking access to the as method. It should not be adjusted
+ * to remove said override.
+ */
public class NoAsRowsAndColumns implements RowsAndColumns
{
private final RowsAndColumns rac;
@@ -49,4 +55,13 @@ public class NoAsRowsAndColumns implements RowsAndColumns
{
return rac.findColumn(name);
}
+
+ @Nullable
+ @Override
+ public <T> T as(Class<T> clazz)
+ {
+ // Pretend that it doesn't implement the clazz. This is the whole point of this class, so if you are considering
+ // changing this method, please stop.
+ return null;
+ }
}
diff --git a/processing/src/test/java/org/apache/druid/query/rowsandcols/serde/FrameWireTransferableTest.java b/processing/src/test/java/org/apache/druid/query/rowsandcols/serde/FrameWireTransferableTest.java
new file mode 100644
index 00000000000..27b5af9d4c1
--- /dev/null
+++ b/processing/src/test/java/org/apache/druid/query/rowsandcols/serde/FrameWireTransferableTest.java
@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.rowsandcols.serde;
+
+import org.apache.druid.frame.wire.FrameWireTransferable;
+import org.apache.druid.jackson.DefaultObjectMapper;
+import org.apache.druid.java.util.common.ByteBufferUtils;
+import org.apache.druid.query.rowsandcols.MapOfColumnsRowsAndColumns;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.column.IntArrayColumn;
+import org.apache.druid.query.rowsandcols.column.ObjectArrayColumn;
+import org.apache.druid.query.rowsandcols.concrete.ColumnBasedFrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.ColumnBasedFrameRowsAndColumnsTest;
+import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowsAndColumns;
+import org.apache.druid.query.rowsandcols.concrete.RowBasedFrameRowsAndColumnsTest;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.apache.druid.segment.column.ColumnType;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class FrameWireTransferableTest
+{
+ private static final MapOfColumnsRowsAndColumns INPUT_RAC = MapOfColumnsRowsAndColumns
+ .builder()
+ .add("colA", new IntArrayColumn(new int[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}))
+ .add("colB", new IntArrayColumn(new int[]{4, -4, 3, -3, 4, 82, -90, 4, 0, 0}))
+ .add(
+ "colC",
+ new ObjectArrayColumn(
+ new Object[]{"foo", "bar", "baz", "spot", "automotive", "business", "a", "b", "c", "d"},
+ ColumnType.STRING
+ )
+ )
+ .build();
+
+ private final DefaultObjectMapper objectMapper = new DefaultObjectMapper();
+
+ @Test
+ public void testColumnBasedFrame() throws Exception
+ {
+ final ColumnBasedFrameRowsAndColumns frameRac = ColumnBasedFrameRowsAndColumnsTest.buildFrame(INPUT_RAC);
+ final FrameWireTransferable wt = new FrameWireTransferable(frameRac.getFrame(), frameRac.getSignature());
+
+ // Serialize
+ final WireTransferable.ByteArrayOffsetAndLen serialized = wt.serializedBytes(objectMapper);
+ final byte[] bytes = new byte[serialized.getLength()];
+ System.arraycopy(serialized.getArray(), serialized.getOffset(), bytes, 0, serialized.getLength());
+
+ // Deserialize
+ final FrameWireTransferable.Deserializer deserializer = new FrameWireTransferable.Deserializer();
+ final RowsAndColumns deserialized = deserializer.deserialize(objectMapper, ByteBufferUtils.wrapLE(bytes));
+
+ Assert.assertTrue(deserialized instanceof ColumnBasedFrameRowsAndColumns);
+ Assert.assertEquals(frameRac, deserialized);
+ }
+
+ @Test
+ public void testColumnBasedFrame_withoutSignature() throws Exception
+ {
+ final ColumnBasedFrameRowsAndColumns frameRac = ColumnBasedFrameRowsAndColumnsTest.buildFrame(INPUT_RAC);
+ final FrameWireTransferable wt = new FrameWireTransferable(frameRac.getFrame(), null);
+
+ // Serialize
+ final WireTransferable.ByteArrayOffsetAndLen serialized = wt.serializedBytes(objectMapper);
+ final byte[] bytes = new byte[serialized.getLength()];
+ System.arraycopy(serialized.getArray(), serialized.getOffset(), bytes, 0, serialized.getLength());
+
+ // Deserialize
+ final FrameWireTransferable.Deserializer deserializer = new FrameWireTransferable.Deserializer();
+ final RowsAndColumns deserialized = deserializer.deserialize(objectMapper, ByteBufferUtils.wrapLE(bytes));
+
+ Assert.assertTrue(deserialized instanceof ColumnBasedFrameRowsAndColumns);
+ Assert.assertEquals(frameRac.getFrame(), ((ColumnBasedFrameRowsAndColumns) deserialized).getFrame());
+ Assert.assertFalse(((ColumnBasedFrameRowsAndColumns) deserialized).hasSignature());
+ }
+
+ @Test
+ public void testRowBasedFrame() throws Exception
+ {
+ final RowBasedFrameRowsAndColumns frameRac = RowBasedFrameRowsAndColumnsTest.MAKER.apply(INPUT_RAC);
+ final FrameWireTransferable wt = new FrameWireTransferable(frameRac.getFrame(), frameRac.getSignature());
+
+ // Serialize
+ final WireTransferable.ByteArrayOffsetAndLen serialized = wt.serializedBytes(objectMapper);
+ final byte[] bytes = new byte[serialized.getLength()];
+ System.arraycopy(serialized.getArray(), serialized.getOffset(), bytes, 0, serialized.getLength());
+
+ // Deserialize
+ final FrameWireTransferable.Deserializer deserializer = new FrameWireTransferable.Deserializer();
+ final RowsAndColumns deserialized = deserializer.deserialize(objectMapper, ByteBufferUtils.wrapLE(bytes));
+
+ Assert.assertTrue(deserialized instanceof RowBasedFrameRowsAndColumns);
+ Assert.assertEquals(frameRac, deserialized);
+ }
+
+ @Test
+ public void testColumnBasedFrameViaAs() throws Exception
+ {
+ final ColumnBasedFrameRowsAndColumns frameRac = ColumnBasedFrameRowsAndColumnsTest.buildFrame(INPUT_RAC);
+
+ // Get WireTransferable via as()
+ final WireTransferable wt = frameRac.as(WireTransferable.class);
+ Assert.assertNotNull(wt);
+
+ // Serialize
+ final WireTransferable.ByteArrayOffsetAndLen serialized = wt.serializedBytes(objectMapper);
+ final byte[] bytes = new byte[serialized.getLength()];
+ System.arraycopy(serialized.getArray(), serialized.getOffset(), bytes, 0, serialized.getLength());
+
+ // Deserialize
+ final FrameWireTransferable.Deserializer deserializer = new FrameWireTransferable.Deserializer();
+ final RowsAndColumns deserialized = deserializer.deserialize(objectMapper, ByteBufferUtils.wrapLE(bytes));
+
+ Assert.assertTrue(deserialized instanceof ColumnBasedFrameRowsAndColumns);
+ Assert.assertEquals(frameRac, deserialized);
+ }
+
+ @Test
+ public void testRowBasedFrameViaAs() throws Exception
+ {
+ final RowBasedFrameRowsAndColumns frameRac = RowBasedFrameRowsAndColumnsTest.MAKER.apply(INPUT_RAC);
+
+ // Get WireTransferable via as()
+ final WireTransferable wt = frameRac.as(WireTransferable.class);
+ Assert.assertNotNull(wt);
+
+ // Serialize
+ final WireTransferable.ByteArrayOffsetAndLen serialized = wt.serializedBytes(objectMapper);
+ final byte[] bytes = new byte[serialized.getLength()];
+ System.arraycopy(serialized.getArray(), serialized.getOffset(), bytes, 0, serialized.getLength());
+
+ // Deserialize
+ final FrameWireTransferable.Deserializer deserializer = new FrameWireTransferable.Deserializer();
+ final RowsAndColumns deserialized = deserializer.deserialize(objectMapper, ByteBufferUtils.wrapLE(bytes));
+
+ Assert.assertTrue(deserialized instanceof RowBasedFrameRowsAndColumns);
+ Assert.assertEquals(frameRac, deserialized);
+ }
+}
diff --git a/processing/src/test/java/org/apache/druid/query/rowsandcols/serde/WireTransferableContextTest.java b/processing/src/test/java/org/apache/druid/query/rowsandcols/serde/WireTransferableContextTest.java
new file mode 100644
index 00000000000..0894c228328
--- /dev/null
+++ b/processing/src/test/java/org/apache/druid/query/rowsandcols/serde/WireTransferableContextTest.java
@@ -0,0 +1,126 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.rowsandcols.serde;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.query.rowsandcols.RowsAndColumns;
+import org.apache.druid.query.rowsandcols.semantic.WireTransferable;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+public class WireTransferableContextTest
+{
+ @Test
+ public void test_defaultLegacyFrameSerialization()
+ {
+ Assert.assertTrue(WireTransferableContext.DEFAULT_LEGACY_FRAME_SERIALIZATION);
+ }
+
+ @Test
+ public void test_constructor_andGetters()
+ {
+ final ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class);
+ final WireTransferable.ConcreteDeserializer mockDeserializer =
+ Mockito.mock(WireTransferable.ConcreteDeserializer.class);
+
+ final WireTransferableContext context = new WireTransferableContext(mockMapper, mockDeserializer, true);
+
+ Assert.assertSame(mockMapper, context.smileMapper());
+ Assert.assertSame(mockDeserializer, context.concreteDeserializer());
+ Assert.assertTrue(context.useLegacyFrameSerialization());
+ }
+
+ @Test
+ public void test_constructor_withNullValues()
+ {
+ final WireTransferableContext context = new WireTransferableContext(null, null, false);
+
+ Assert.assertNull(context.smileMapper());
+ Assert.assertNull(context.concreteDeserializer());
+ Assert.assertFalse(context.useLegacyFrameSerialization());
+ }
+
+ @Test
+ public void test_useLegacyFrameSerialization_true()
+ {
+ final WireTransferableContext context = new WireTransferableContext(null, null, true);
+ Assert.assertTrue(context.useLegacyFrameSerialization());
+ }
+
+ @Test
+ public void test_useLegacyFrameSerialization_false()
+ {
+ final WireTransferableContext context = new WireTransferableContext(null, null, false);
+ Assert.assertFalse(context.useLegacyFrameSerialization());
+ }
+
+ @Test
+ public void test_serializedBytes() throws IOException
+ {
+ final ObjectMapper mockMapper = Mockito.mock(ObjectMapper.class);
+ final WireTransferable mockWireTransferable = Mockito.mock(WireTransferable.class);
+ final WireTransferable.ByteArrayOffsetAndLen expectedResult =
+ new WireTransferable.ByteArrayOffsetAndLen(new byte[]{1, 2, 3}, 0, 3);
+
+ Mockito.when(mockWireTransferable.serializedBytes(mockMapper)).thenReturn(expectedResult);
+
+ final WireTransferableContext context = new WireTransferableContext(mockMapper, null, true);
+ final WireTransferable.ByteArrayOffsetAndLen result = context.serializedBytes(mockWireTransferable);
+
+ Assert.assertSame(expectedResult, result);
+ Mockito.verify(mockWireTransferable).serializedBytes(mockMapper);
+ }
+
+ @Test
+ public void test_deserialize()
+ {
+ final WireTransferable.ConcreteDeserializer mockDeserializer =
+ Mockito.mock(WireTransferable.ConcreteDeserializer.class);
+ final RowsAndColumns mockRac = Mockito.mock(RowsAndColumns.class);
+ final ByteBuffer buffer = ByteBuffer.allocate(10);
+
+ Mockito.when(mockDeserializer.deserialize(buffer)).thenReturn(mockRac);
+
+ final WireTransferableContext context = new WireTransferableContext(null, mockDeserializer, true);
+ final RowsAndColumns result = context.deserialize(buffer);
+
+ Assert.assertSame(mockRac, result);
+ Mockito.verify(mockDeserializer).deserialize(buffer);
+ }
+
+ @Test
+ public void test_deserialize_withNullDeserializer_throwsException()
+ {
+ final WireTransferableContext context = new WireTransferableContext(null, null, true);
+ final ByteBuffer buffer = ByteBuffer.allocate(10);
+
+ final DruidException exception = Assert.assertThrows(
+ DruidException.class,
+ () -> context.deserialize(buffer)
+ );
+ Assert.assertTrue(exception.getMessage().contains("Cannot deserialize"));
+ Assert.assertTrue(exception.getMessage().contains("no concreteDeserializer"));
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org