This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 4824a5fd1f7 Add bounded codec runtime and compression handlers (#19285)
4824a5fd1f7 is described below
commit 4824a5fd1f793b55a02de69f00395497830cf7a7
Author: Xiang Fu <[email protected]>
AuthorDate: Mon Aug 24 19:36:06 2026 -0700
Add bounded codec runtime and compression handlers (#19285)
* Add bounded codec runtime and compression handlers
* Harden Snappy decode bounds, reserve CODEC in the registry, reject ZSTD(0)
- Snappy was the only pipeline codec missing the sanity cap on the
header-declared decompressed size; a corrupt or hostile header must not
drive allocation. Rather than adding a fourth private copy of the cap,
hoist it into CodecBufferUtils.checkDeclaredDecompressedSize and route
all four codecs' guard sites through it, so the bound cannot be omitted
from a future codec or silently drift between implementations. Corrupt-
input tests cover both Snappy decode paths at the cap boundary (cap + 1).
- CodecRegistry.register() rejects the reserved wrapper name via the shared
CodecSpecParser.REMOVED_WRAPPER_NAME constant as defense in depth
alongside the parser and CodecInvocation reservations.
- ZSTD(0) is rejected: zstd treats level 0 as "use the default level", so
it would behave identically to ZSTD(3) under a second canonical spelling,
and the canonical spec is frozen into segment headers. Rejection is
reversible later; acceptance is not.
* Address codec runtime review feedback
* Address codec runtime review nits
---
.../segment/local/io/codec/ChunkCodecHandler.java | 81 ++++
.../segment/local/io/codec/CodecBufferUtils.java | 82 ++++
.../pinot/segment/local/io/codec/CodecContext.java | 39 ++
.../segment/local/io/codec/CodecDefinition.java | 80 ++++
.../pinot/segment/local/io/codec/CodecKind.java | 35 ++
.../pinot/segment/local/io/codec/CodecOptions.java | 35 ++
.../local/io/codec/CodecPipelineExecutor.java | 441 +++++++++++++++++++++
.../local/io/codec/CodecPipelineValidator.java | 90 +++++
.../segment/local/io/codec/CodecRegistry.java | 118 ++++++
.../local/io/codec/GzipCodecDefinition.java | 273 +++++++++++++
.../segment/local/io/codec/Lz4CodecDefinition.java | 174 ++++++++
.../local/io/codec/SnappyCodecDefinition.java | 175 ++++++++
.../local/io/codec/ZstdCodecDefinition.java | 238 +++++++++++
.../local/io/codec/CodecPipelineExecutorTest.java | 254 ++++++++++++
.../local/io/codec/CodecPipelineValidatorTest.java | 188 +++++++++
.../segment/local/io/codec/CodecRegistryTest.java | 141 +++++++
.../io/codec/CompressionCodecCorruptInputTest.java | 236 +++++++++++
.../local/io/codec/ZstdCodecDefinitionTest.java | 49 +++
18 files changed, 2729 insertions(+)
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/ChunkCodecHandler.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/ChunkCodecHandler.java
new file mode 100644
index 00000000000..57d94e3b78a
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/ChunkCodecHandler.java
@@ -0,0 +1,81 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+
+/// Extension of [CodecDefinition] that adds the encode/decode operations
needed to execute a
+/// codec pipeline over forward-index chunks.
+///
+/// All implementations are expected to be stateless and thread-safe.
+///
+/// Buffer contract for all methods:
+///
+/// - [#encode()] and [#decode()]: `src` is ready for read (position=0);
+/// the returned buffer is ready for read and owned by the caller.
+/// - [#decodeInto()]: implementations must treat `dst` as freshly cleared
+/// (calling `dst.clear()` internally is recommended as a defensive
first step);
+/// `dst` is flipped (position=0, limit=decoded bytes) on return.
+///
+/// @param <O> typed [CodecOptions] for this codec
+interface ChunkCodecHandler<O extends CodecOptions> extends CodecDefinition<O>
{
+
+ /// Encodes `src` and returns the encoded bytes ready for read.
+ ///
+ /// **Position contract:** implementations may consume `src` (advance its
position).
+ /// Callers that need to re-read `src` after this call must pass
`src.duplicate()`.
+ ///
+ /// @param options parsed options for this codec invocation
+ /// @param ctx column context (data type, etc.)
+ /// @param src unencoded data, ready for read
+ /// @return encoded buffer ready for read; caller owns this buffer
+ ByteBuffer encode(O options, CodecContext ctx, ByteBuffer src) throws
IOException;
+
+ /// Decodes `src` and returns the decoded bytes ready for read.
+ ///
+ /// **Position contract:** implementations may consume `src` (advance its
position).
+ /// Callers that need to re-read `src` after this call must pass
`src.duplicate()`.
+ ///
+ /// @param options parsed options for this codec invocation
+ /// @param ctx column context
+ /// @param src encoded data, ready for read
+ /// @return decoded buffer ready for read; caller owns this buffer
+ ByteBuffer decode(O options, CodecContext ctx, ByteBuffer src) throws
IOException;
+
+ /// Decodes `src` directly into `dst`, avoiding an extra allocation.
+ /// Implementations must treat `dst` as freshly cleared and flip it before
returning.
+ ///
+ /// Callers must ensure `dst` is a direct [ByteBuffer] when
+ /// [#requiresDirectDstBuffer()] returns `true`.
+ ///
+ /// @param options parsed options for this codec invocation
+ /// @param ctx column context
+ /// @param src encoded data, ready for read
+ /// @param dst output buffer; must be direct when required; must have
sufficient capacity
+ void decodeInto(O options, CodecContext ctx, ByteBuffer src, ByteBuffer dst)
throws IOException;
+
+ /// Returns an upper bound on the encoded byte count for an input of
`inputSize` bytes.
+ int maxEncodedSize(O options, int inputSize);
+
+ /// Returns `true` if [#decodeInto()] requires `dst` to be a direct
+ /// [ByteBuffer] (e.g. codecs that delegate to JNI libraries with
direct-buffer-only APIs).
+ boolean requiresDirectDstBuffer();
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecBufferUtils.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecBufferUtils.java
new file mode 100644
index 00000000000..057c574e98a
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecBufferUtils.java
@@ -0,0 +1,82 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import org.apache.pinot.segment.spi.memory.CleanerUtil;
+
+
+/// Package-private buffer helpers shared across codec handler implementations.
+final class CodecBufferUtils {
+
+ private CodecBufferUtils() {
+ }
+
+ /// Sanity cap (1 GiB) on any decompressed size declared by untrusted
encoded segment data. A
+ /// corrupt or hostile declaration must never drive a giant pre-allocation;
1 GiB is well above
+ /// any realistic chunk size.
+ static final long MAX_DECLARED_DECOMPRESSED_SIZE = 1L << 30;
+
+ /// Validates a decompressed size declared by untrusted segment data and
returns it when in range.
+ ///
+ /// Shared by every codec definition so the bound cannot be omitted from a
new codec or silently
+ /// drift between implementations.
+ ///
+ /// @param declared size read from the encoded data
+ /// @param codec codec display name for the error message
+ /// @param source where the size was read from, e.g. "length prefix" or
"frame header"
+ /// @return `declared`, guaranteed to be in `[0,
MAX_DECLARED_DECOMPRESSED_SIZE]`
+ /// @throws IOException if the declared size is negative or exceeds the cap
+ static int checkDeclaredDecompressedSize(long declared, String codec, String
source)
+ throws IOException {
+ if (declared < 0 || declared > MAX_DECLARED_DECOMPRESSED_SIZE) {
+ throw new IOException(codec + ": declared decompressed size " + declared
+ " in " + source
+ + " is out of range [0, " + MAX_DECLARED_DECOMPRESSED_SIZE + "].
Segment may be corrupt.");
+ }
+ return (int) declared;
+ }
+
+ /// Returns `buf` if already direct; otherwise copies into a new direct
buffer.
+ ///
+ /// **Hot-path note:** pipeline callers should pass direct buffers. The
heap-copy branch is a
+ /// defensive fallback for tests and ad-hoc callers; using it on production
hot paths would add
+ /// a per-call direct-buffer allocation.
+ static ByteBuffer toDirectBuffer(ByteBuffer buf) {
+ if (buf.isDirect()) {
+ return buf;
+ }
+ ByteBuffer direct = ByteBuffer.allocateDirect(buf.remaining());
+ direct.put(buf.duplicate());
+ direct.flip();
+ return direct;
+ }
+
+ /// Releases `converted` when [#toDirectBuffer(ByteBuffer)] had to copy a
heap buffer.
+ static void cleanDirectCopy(ByteBuffer original, ByteBuffer converted) {
+ if (converted != original) {
+ CleanerUtil.cleanQuietly(converted);
+ }
+ }
+
+ /// Releases an owned direct buffer after a failed codec operation.
+ static void cleanQuietly(ByteBuffer buffer) {
+ CleanerUtil.cleanQuietly(buffer);
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecContext.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecContext.java
new file mode 100644
index 00000000000..011f6d958f4
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecContext.java
@@ -0,0 +1,39 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+
+
+/// Context supplied to [CodecDefinition#validateContext()] during pipeline
validation.
+///
+/// Carries the column's stored data type so that codecs can reject
unsupported types at
+/// configuration time rather than at runtime. Instances are immutable and
thread-safe.
+final class CodecContext {
+ private final DataType _dataType;
+
+ CodecContext(DataType dataType) {
+ _dataType = dataType;
+ }
+
+ /// Returns the stored [DataType] of the column being indexed.
+ DataType getDataType() {
+ return _dataType;
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecDefinition.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecDefinition.java
new file mode 100644
index 00000000000..49491d02197
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecDefinition.java
@@ -0,0 +1,80 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.util.List;
+
+
+/// Describes a single codec (transform or compression) that can participate
in a pipeline.
+///
+/// Implementations are registered in a `CodecRegistry` and looked up by
+/// [#name()] during DSL parsing. Each implementation is responsible for:
+///
+/// 1. Parsing its own raw string arguments into typed [CodecOptions].
+/// 1. Validating that the options are consistent with the target column's
context.
+/// 1. Producing a stable canonical string so pipelines round-trip through
metadata.
+///
+/// Implementations must be stateless and thread-safe; a single instance is
reused for
+/// all columns that reference the codec.
+///
+/// @param <O> the concrete [CodecOptions] type for this codec
+interface CodecDefinition<O extends CodecOptions> {
+
+ /// The uppercase name used in the DSL, e.g. `"LZ4"` or `"ZSTD"`.
+ /// Must be unique across all registered codecs.
+ String name();
+
+ /// Whether this codec is a transform or a byte-compression stage.
+ CodecKind kind();
+
+ /// Parses the raw positional arguments from the DSL invocation into typed
options.
+ ///
+ /// @param args positional string arguments; empty list when the codec was
invoked without parens
+ /// or with an empty argument list
+ /// @return parsed options (never `null`)
+ /// @throws IllegalArgumentException if the argument list is invalid for
this codec
+ O parseOptions(List<String> args);
+
+ /// Validates that the codec can be applied to the given column context.
+ ///
+ /// @param options options previously returned by [#parseOptions()]
+ /// @param ctx context describing the target column
+ /// @throws IllegalArgumentException if the codec is incompatible with this
context
+ void validateContext(O options, CodecContext ctx);
+
+ /// Returns the canonical DSL string for the given options, e.g. `"ZSTD(3)"`
or
+ /// `"LZ4"`. The canonical form is stored in segment metadata and used when
+ /// re-opening the segment.
+ String canonicalize(O options);
+
+ /// Whether this codec's `encode` output is a contiguous array of
column-typed values of the same
+ /// width as its input (so that a subsequent [CodecKind#TRANSFORM] stage can
consume it directly).
+ ///
+ /// This governs how stages may be chained in a pipeline:
+ /// - **Typed-layout-preserving transforms** return `true`: they map a typed
+ /// value array to a same-width typed value array, so any number of them
may be chained.
+ /// - **Packing transforms** return `false`: they emit a bit-packed /
+ /// self-framed byte stream that is no longer a typed value array, so they
must be the **last**
+ /// transform in the chain (only [CodecKind#COMPRESSION] stages may
follow).
+ /// - **Compression** codecs are byte→byte and the value returned here is
irrelevant (a compression
+ /// stage always ends the "typed value" domain); the default `false` is
appropriate.
+ default boolean preservesTypedValueLayout() {
+ return false;
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecKind.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecKind.java
new file mode 100644
index 00000000000..0be5ea52685
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecKind.java
@@ -0,0 +1,35 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+/// Classification of a codec within a pipeline.
+///
+/// A pipeline may contain any number of typed-layout-preserving [#TRANSFORM]
stages, followed by at most
+/// one packing transform, followed by any number of [#COMPRESSION] stages. A
packing transform or
+/// compression stage ends the column-typed value domain, so no transform may
follow either one.
+/// Stages run left-to-right on encode and right-to-left on decode.
+enum CodecKind {
+ /// Reversible transformation over column-typed values.
Typed-layout-preserving transforms may
+ /// be chained; a packing transform emits bytes and
+ /// must be the last transform. See
[CodecDefinition#preservesTypedValueLayout()].
+ TRANSFORM,
+ /// Byte-level compression (e.g. ZSTD). Any number of compression stages may
follow the
+ /// transforms; once compression begins, only further compression stages are
allowed.
+ COMPRESSION
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecOptions.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecOptions.java
new file mode 100644
index 00000000000..b565c3c7992
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecOptions.java
@@ -0,0 +1,35 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+/// Marker interface for codec-specific parsed options.
+///
+/// Each [CodecDefinition] defines its own concrete options class that carries
the
+/// validated, typed parameters for that codec (e.g. compression level for
ZSTD).
+///
+/// **Implementation contract:**
+///
+/// - Implementations MUST be immutable and thread-safe — a single instance is
shared across
+/// all encode/decode calls for a given codec invocation.
+/// - `equals`/`hashCode` MUST reflect every configurable parameter so two
+/// [canonical][CodecDefinition#canonicalize()] specs that differ only
in option values
+/// compare unequal. Canonical-spec equality drives
reload-on-config-change detection.
+/// - `toString` should return a human-readable parameter listing for
log/error messages.
+interface CodecOptions {
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecPipelineExecutor.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecPipelineExecutor.java
new file mode 100644
index 00000000000..4e4c8b56877
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecPipelineExecutor.java
@@ -0,0 +1,441 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.pinot.segment.spi.codec.CodecInvocation;
+import org.apache.pinot.segment.spi.codec.CodecPipeline;
+import org.apache.pinot.segment.spi.codec.CodecSpecParser;
+import org.apache.pinot.segment.spi.memory.CleanerUtil;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+
+
+/// Executes a parsed and validated [CodecPipeline] for a single forward-index
chunk.
+///
+/// Write path: values → transforms (in order) → compression stages → bytes
stored on disk.
+/// Read path: bytes from disk → reverse compression/transform stages → values.
+///
+/// The executor is constructed once per column from the canonical `codecSpec`
string
+/// stored in the file header and is thread-safe for concurrent read calls (it
holds no mutable
+/// per-call state).
+///
+/// The executor is codec-agnostic: it holds an ordered list of internal bound
stages, each pairing
+/// a handler with its parsed options. The registry and handlers are a closed,
package-private
+/// runtime; this class is the only public entry point and only drives the
pipeline loop.
+///
+/// ### Buffer contract
+///
+/// - [#encode]: `src` is ready for read (position=0); returns a new
+/// [ByteBuffer] ready for read containing the encoded bytes.
+/// - [#decode(ByteBuffer)]: `src` is ready for read; returns a new
+/// [ByteBuffer] ready for read containing the decoded bytes.
+/// - [#maxEncodedSize]: returns an upper bound on encoded size.
+public final class CodecPipelineExecutor {
+
+ /// A codec handler bound to the options parsed from a specific pipeline
invocation.
+ private static final class BoundStage<O extends CodecOptions> {
+ final ChunkCodecHandler<O> _handler;
+ final O _options;
+ final CodecContext _ctx;
+
+ BoundStage(ChunkCodecHandler<O> handler, O options, CodecContext ctx) {
+ _handler = handler;
+ _options = options;
+ _ctx = ctx;
+ }
+
+ ByteBuffer encode(ByteBuffer src) throws IOException {
+ return _handler.encode(_options, _ctx, src);
+ }
+
+ ByteBuffer decode(ByteBuffer src) throws IOException {
+ return _handler.decode(_options, _ctx, src);
+ }
+
+ void decodeInto(ByteBuffer src, ByteBuffer dst) throws IOException {
+ _handler.decodeInto(_options, _ctx, src, dst);
+ }
+
+ int maxEncodedSize(int inputSize) {
+ return _handler.maxEncodedSize(_options, inputSize);
+ }
+
+ boolean requiresDirectDstBuffer() {
+ return _handler.requiresDirectDstBuffer();
+ }
+
+ boolean isCompression() {
+ return _handler.kind() == CodecKind.COMPRESSION;
+ }
+
+ String canonicalize() {
+ return _handler.canonicalize(_options);
+ }
+ }
+
+ /// Caller-owned, reusable decode workspace.
+ ///
+ /// The workspace retains at most two direct buffers and grows them on
demand. Multi-stage
+ /// decoding alternates between those buffers, so repeated chunk reads avoid
allocating and
+ /// explicitly cleaning a direct buffer for every intermediate stage. This
class is not
+ /// thread-safe; each reader context or calling thread must own a separate
instance.
+ public static final class DecodeScratch implements AutoCloseable {
+ private final ByteBuffer[] _buffers = new ByteBuffer[2];
+ private int[] _stageBounds = new int[0];
+ private int _allocationCount;
+ private boolean _closed;
+
+ private ByteBuffer buffer(int slot, int capacity) {
+ ensureOpen();
+ Preconditions.checkArgument(capacity >= 0, "Scratch capacity must be
non-negative: %s", capacity);
+ ByteBuffer buffer = _buffers[slot];
+ if (buffer == null || buffer.capacity() < capacity) {
+ CleanerUtil.cleanQuietly(buffer);
+ buffer = ByteBuffer.allocateDirect(capacity);
+ _buffers[slot] = buffer;
+ _allocationCount++;
+ }
+ // Limit the returned view's capacity to this stage's validated bound.
Reusing a larger
+ // backing buffer must not let a corrupt inner frame exploit capacity
left from an earlier
+ // larger chunk.
+ ByteBuffer view = buffer.duplicate();
+ view.clear();
+ view.limit(capacity);
+ return view.slice();
+ }
+
+ private int[] stageBounds(int stageCount) {
+ ensureOpen();
+ if (_stageBounds.length < stageCount) {
+ _stageBounds = new int[stageCount];
+ }
+ return _stageBounds;
+ }
+
+ private void ensureOpen() {
+ Preconditions.checkState(!_closed, "DecodeScratch is closed");
+ }
+
+ int allocationCount() {
+ return _allocationCount;
+ }
+
+ @Override
+ public void close() {
+ if (_closed) {
+ return;
+ }
+ _closed = true;
+ for (int i = 0; i < _buffers.length; i++) {
+ CleanerUtil.cleanQuietly(_buffers[i]);
+ _buffers[i] = null;
+ }
+ _stageBounds = new int[0];
+ }
+ }
+
+ private final List<BoundStage<?>> _stages;
+ private final String _canonicalSpec;
+ private final DataType _storedType;
+ private final boolean _hasCompression;
+ private final boolean _requiresDirectDstBuffer;
+
+ /// Creates an executor by parsing and validating the given spec.
+ ///
+ /// @param spec the codec DSL string (e.g. `"LZ4,ZSTD(3)"`)
+ /// @param storedType stored data type used for validation
+ public static CodecPipelineExecutor create(String spec, DataType storedType)
{
+ return create(spec, new CodecContext(storedType), CodecRegistry.DEFAULT);
+ }
+
+ /// Package-scoped construction hook for tests that provide a custom closed
registry.
+ static CodecPipelineExecutor create(String spec, CodecContext ctx,
CodecRegistry registry) {
+ CodecPipeline pipeline = CodecSpecParser.parse(spec);
+ CodecPipelineValidator.validate(pipeline, registry, ctx);
+ return new CodecPipelineExecutor(pipeline, registry, ctx);
+ }
+
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ private CodecPipelineExecutor(CodecPipeline pipeline, CodecRegistry
registry, CodecContext ctx) {
+ List<CodecInvocation> invocations = pipeline.stages();
+ if (invocations.isEmpty()) {
+ throw new IllegalArgumentException("Codec pipeline must contain at least
one stage");
+ }
+ List<BoundStage<?>> stages = new ArrayList<>(invocations.size());
+
+ for (CodecInvocation inv : invocations) {
+ ChunkCodecHandler handler = (ChunkCodecHandler)
registry.getOrThrow(inv.name());
+ CodecOptions opts = handler.parseOptions(inv.args());
+ stages.add(new BoundStage<>(handler, opts, ctx));
+ }
+ _stages = stages;
+ _canonicalSpec = buildCanonical(stages);
+ _storedType = ctx.getDataType();
+ _hasCompression = stages.stream().anyMatch(BoundStage::isCompression);
+ // decodeInto() writes the final output through stage zero. All other
stage outputs use
+ // executor-owned direct scratch buffers, so only stage zero constrains
the caller's dst.
+ _requiresDirectDstBuffer = stages.get(0).requiresDirectDstBuffer();
+ }
+
+ /// Returns the canonical spec string derived from the parsed pipeline.
+ public String getCanonicalSpec() {
+ return _canonicalSpec;
+ }
+
+ /// Returns the stored column type to which this executor was validated and
bound.
+ public DataType getStoredType() {
+ return _storedType;
+ }
+
+ /// Returns an upper bound on the number of bytes that [#encode] may produce
for
+ /// a decoded chunk of the given byte length.
+ int maxEncodedSize(int decodedSize) {
+ return maxEncodedSize(decodedSize, Integer.MAX_VALUE, Long.MAX_VALUE);
+ }
+
+ /// Returns the composed encoded-size bound while requiring every stage's
bound to stay within
+ /// `maxStageSize`. This lets an on-disk format reject a pipeline/chunk-size
combination before
+ /// either the writer produces an unreadable file or the reader allocates
excessive scratch.
+ int maxEncodedSize(int decodedSize, int maxStageSize) {
+ return maxEncodedSize(decodedSize, maxStageSize, Long.MAX_VALUE);
+ }
+
+ /// Returns the composed encoded-size bound while also capping the sum of
all stage-output
+ /// bounds. The cumulative limit bounds CPU and allocation churn for long
pipelines even when
+ /// each individual stage stays below `maxStageSize`.
+ public int maxEncodedSize(int decodedSize, int maxStageSize, long
maxCumulativeSize) {
+ Preconditions.checkArgument(decodedSize >= 0, "decodedSize must be
non-negative: %s", decodedSize);
+ Preconditions.checkArgument(maxStageSize >= decodedSize,
+ "maxStageSize %s must be at least decodedSize %s", maxStageSize,
decodedSize);
+ Preconditions.checkArgument(maxCumulativeSize >= decodedSize,
+ "maxCumulativeSize %s must be at least decodedSize %s",
maxCumulativeSize, decodedSize);
+ int size = decodedSize;
+ long cumulativeSize = 0;
+ for (int i = 0; i < _stages.size(); i++) {
+ size = _stages.get(i).maxEncodedSize(size);
+ if (size < 0 || size > maxStageSize) {
+ throw new IllegalArgumentException(
+ "Codec stage " + i + " maximum encoded size " + size + " is
outside [0, " + maxStageSize
+ + "] for pipeline " + _canonicalSpec);
+ }
+ cumulativeSize += size;
+ if (cumulativeSize > maxCumulativeSize) {
+ throw new IllegalArgumentException(
+ "Codec pipeline cumulative stage-output bound " + cumulativeSize +
" exceeds " + maxCumulativeSize
+ + " at stage " + i + " for pipeline " + _canonicalSpec);
+ }
+ }
+ return size;
+ }
+
+ /// Encodes `src` through the pipeline and returns the encoded bytes ready
for read.
+ ///
+ /// The readable range is interpreted in the persisted big-endian
typed-value order, independent
+ /// of the caller view's byte order. The caller's position, limit, and order
are not modified.
+ ///
+ /// @param src decoded chunk data, ready for read (position=0,
limit=dataSize)
+ /// @return encoded buffer ready for read; the caller owns this buffer
+ public ByteBuffer encode(ByteBuffer src) throws IOException {
+ ByteBuffer input = src.duplicate().order(ByteOrder.BIG_ENDIAN);
+ ByteBuffer current = input;
+ try {
+ for (BoundStage<?> stage : _stages) {
+ ByteBuffer previous = current;
+ current = stage.encode(previous);
+ if (previous != input && previous != current) {
+ CleanerUtil.cleanQuietly(previous);
+ }
+ }
+ return current;
+ } catch (IOException | RuntimeException | Error e) {
+ if (current != input) {
+ CleanerUtil.cleanQuietly(current);
+ }
+ throw e;
+ }
+ }
+
+ /// Decodes `src` through the reversed pipeline and returns the original
bytes.
+ ///
+ /// @param src encoded chunk data, ready for read
+ /// @return decoded buffer ready for read; the caller owns this buffer
+ ByteBuffer decode(ByteBuffer src) throws IOException {
+ ByteBuffer current = src;
+ try {
+ for (int i = _stages.size() - 1; i >= 0; i--) {
+ ByteBuffer previous = current;
+ current = _stages.get(i).decode(previous);
+ if (previous != src && previous != current) {
+ CleanerUtil.cleanQuietly(previous);
+ }
+ }
+ return current;
+ } catch (IOException | RuntimeException | Error e) {
+ if (current != src) {
+ CleanerUtil.cleanQuietly(current);
+ }
+ throw e;
+ }
+ }
+
+ /// Decodes `src` through the reversed pipeline, writing the result directly
into
+ /// `dst`. On return `dst` is flipped and ready for read (position=0,
+ /// limit=decoded size).
+ ///
+ /// For single-stage pipelines (transform-only or compression-only), the
decoded bytes are
+ /// written directly into `dst`, avoiding an intermediate allocation. For
multi-stage
+ /// pipelines an intermediate buffer is allocated for any stage that is not
the last in the
+ /// reversed pipeline; the final stage writes directly into `dst`.
+ ///
+ /// @param src encoded chunk data, ready for read
+ /// @param dst caller-supplied output buffer; must be a *direct*
[ByteBuffer] when
+ /// the pipeline requires it (see
[ChunkCodecHandler#requiresDirectDstBuffer]);
+ /// must have sufficient [ByteBuffer#capacity()] for the decoded
data; its
+ /// position and limit are overwritten before returning
+ /// @throws IOException if decoding fails
+ /// @throws IllegalArgumentException if `dst` is not direct when required,
or does not have
+ /// enough capacity
+ void decode(ByteBuffer src, ByteBuffer dst) throws IOException {
+ decode(src, dst, dst.capacity());
+ }
+
+ /// Decodes into `dst` while bounding every intermediate allocation from the
expected final
+ /// decoded size. This is the segment-reader entry point: codec frames are
untrusted and may
+ /// advertise arbitrary decoded lengths, so an intermediate stage must never
allocate directly
+ /// from its frame header.
+ ///
+ /// @param src encoded chunk data, ready for read
+ /// @param dst caller-owned final output buffer
+ /// @param expectedDecodedSize exact decoded bytes declared and validated by
the outer format
+ void decode(ByteBuffer src, ByteBuffer dst, int expectedDecodedSize) throws
IOException {
+ decode(src, dst, expectedDecodedSize, Integer.MAX_VALUE, Long.MAX_VALUE);
+ }
+
+ /// Variant of [#decode(ByteBuffer, ByteBuffer, int)] that caps every
intermediate scratch
+ /// buffer. Callers reading untrusted persisted data should pass the
format's validated limit.
+ void decode(ByteBuffer src, ByteBuffer dst, int expectedDecodedSize, int
maxIntermediateSize)
+ throws IOException {
+ decode(src, dst, expectedDecodedSize, maxIntermediateSize, Long.MAX_VALUE);
+ }
+
+ /// Variant that also caps the sum of stage-output bounds, limiting total
work and allocation
+ /// churn for an otherwise-valid long pipeline. This convenience overload
owns a temporary
+ /// workspace; production chunk readers should use the caller-owned
[DecodeScratch] overload.
+ void decode(ByteBuffer src, ByteBuffer dst, int expectedDecodedSize, int
maxIntermediateSize,
+ long maxCumulativeSize)
+ throws IOException {
+ try (DecodeScratch scratch = new DecodeScratch()) {
+ decode(src, dst, expectedDecodedSize, maxIntermediateSize,
maxCumulativeSize, scratch);
+ }
+ }
+
+ /// Decodes with caller-owned reusable scratch buffers.
+ ///
+ /// The caller must not share scratch across concurrent decode calls and
must close it when the
+ /// owning reader context is closed.
+ public void decode(ByteBuffer src, ByteBuffer dst, int expectedDecodedSize,
int maxIntermediateSize,
+ long maxCumulativeSize, DecodeScratch scratch)
+ throws IOException {
+ scratch.ensureOpen();
+ Preconditions.checkArgument(!_requiresDirectDstBuffer || dst.isDirect(),
+ "decode(src, dst) requires a direct ByteBuffer for pipeline: %s",
_canonicalSpec);
+ Preconditions.checkArgument(expectedDecodedSize >= 0 &&
expectedDecodedSize <= dst.capacity(),
+ "expectedDecodedSize %s is out of range [0, %s]", expectedDecodedSize,
dst.capacity());
+ Preconditions.checkArgument(maxIntermediateSize >= expectedDecodedSize,
+ "maxIntermediateSize %s must be at least expectedDecodedSize %s",
maxIntermediateSize, expectedDecodedSize);
+ Preconditions.checkArgument(maxCumulativeSize >= expectedDecodedSize,
+ "maxCumulativeSize %s must be at least expectedDecodedSize %s",
maxCumulativeSize, expectedDecodedSize);
+
+ int stageCount = _stages.size();
+ int[] maxOutputAfterStage = scratch.stageBounds(stageCount);
+ int maxSize = expectedDecodedSize;
+ long cumulativeSize = 0;
+ for (int i = 0; i < stageCount; i++) {
+ maxSize = _stages.get(i).maxEncodedSize(maxSize);
+ if (maxSize < 0 || maxSize > maxIntermediateSize) {
+ throw new IllegalArgumentException(
+ "Codec stage " + i + " maximum encoded size " + maxSize + " is
outside [0, "
+ + maxIntermediateSize + "] for pipeline " + _canonicalSpec);
+ }
+ cumulativeSize += maxSize;
+ if (cumulativeSize > maxCumulativeSize) {
+ throw new IllegalArgumentException(
+ "Codec pipeline cumulative stage-output bound " + cumulativeSize +
" exceeds " + maxCumulativeSize
+ + " at stage " + i + " for pipeline " + _canonicalSpec);
+ }
+ maxOutputAfterStage[i] = maxSize;
+ }
+ Preconditions.checkArgument(src.remaining() <=
maxOutputAfterStage[stageCount - 1],
+ "Encoded input size %s exceeds composed pipeline bound %s for expected
decoded size %s",
+ src.remaining(), maxOutputAfterStage[stageCount - 1],
expectedDecodedSize);
+
+ // Decode every stage through decodeInto(). Intermediate stages alternate
between two reusable
+ // direct buffers whose views are capacity-limited to the forward
maxEncodedSize chain rooted
+ // at the validated final decoded size. A corrupt inner frame-controlled
length can therefore
+ // only trigger a bounded "exceeds dst capacity" failure, never a
frame-controlled allocation.
+ ByteBuffer finalOutput = dst.duplicate().order(ByteOrder.BIG_ENDIAN);
+ ByteBuffer current = src;
+ int scratchSlot = 0;
+ for (int i = stageCount - 1; i >= 0; i--) {
+ ByteBuffer output;
+ if (i == 0) {
+ output = finalOutput;
+ } else {
+ output = scratch.buffer(scratchSlot, maxOutputAfterStage[i - 1]);
+ scratchSlot ^= 1;
+ }
+ _stages.get(i).decodeInto(current, output);
+ current = output;
+ }
+ if (finalOutput.remaining() != expectedDecodedSize) {
+ throw new IOException("Codec pipeline decoded " +
finalOutput.remaining() + " bytes but expected "
+ + expectedDecodedSize + " for pipeline " + _canonicalSpec + ".
Segment may be corrupt.");
+ }
+ // decodeInto() flips the final view. Mirror that readable range onto the
caller's buffer while
+ // preserving its independently configured byte order.
+ dst.clear();
+ dst.limit(finalOutput.limit());
+ dst.position(finalOutput.position());
+ }
+
+ /// Returns true if the pipeline has at least one compression stage.
+ boolean isCompressed() {
+ return _hasCompression;
+ }
+
+ // -------------------------------------------------------------------------
+ // Canonical spec builder
+ // -------------------------------------------------------------------------
+
+ private static String buildCanonical(List<BoundStage<?>> stages) {
+ StringBuilder sb = new StringBuilder();
+ for (int i = 0; i < stages.size(); i++) {
+ if (i > 0) {
+ sb.append(',');
+ }
+ sb.append(stages.get(i).canonicalize());
+ }
+ return sb.toString();
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecPipelineValidator.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecPipelineValidator.java
new file mode 100644
index 00000000000..927bfe5a7fa
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecPipelineValidator.java
@@ -0,0 +1,90 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.util.List;
+import org.apache.pinot.segment.spi.codec.CodecInvocation;
+import org.apache.pinot.segment.spi.codec.CodecPipeline;
+
+
+/// Validates a [CodecPipeline] against structural rules and per-codec context.
+///
+/// A pipeline is a chain of the form: **N typed-layout-preserving transforms
→ at most one packing
+/// transform → N compressions**. This is enforced by walking the stages and
tracking whether we
+/// are still operating on column-typed values ("typed domain"):
+///
+/// 1. All codec names must be registered in the supplied [CodecRegistry].
+/// 1. Each codec's [ChunkCodecHandler#validateContext()] must pass for the
column's context.
+/// 1. A [CodecKind#TRANSFORM] stage may only appear while still in the typed
domain — i.e. it must
+/// not follow a packing transform or any compression stage (those no
longer emit a typed value
+/// array, so a transform could not consume their output).
+/// 1. A **typed-layout-preserving** transform (see
+/// [CodecDefinition#preservesTypedValueLayout()]) keeps the pipeline in
the typed domain, so any number
+/// may be chained. A **packing** transform leaves the typed domain, so it
+/// must be the last transform — only compression stages may follow.
+/// 1. Any number of [CodecKind#COMPRESSION] stages may follow, in any order;
each is byte→byte.
+///
+/// Evaluation order: stages run **left-to-right on encode** and
**right-to-left on decode**
+/// — i.e. `A,B` encodes as `B.encode(A.encode(x))` and decodes as
+/// `A.decode(B.decode(y))`. The runtime [CodecPipelineExecutor] enforces
+/// this ordering.
+final class CodecPipelineValidator {
+
+ private CodecPipelineValidator() {
+ }
+
+ /// Validates the pipeline.
+ ///
+ /// @param pipeline pipeline AST to validate
+ /// @param registry registry used to resolve codec names
+ /// @param ctx column context for type validation
+ /// @throws IllegalArgumentException if any validation rule is violated
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ static void validate(CodecPipeline pipeline, CodecRegistry registry,
CodecContext ctx) {
+ List<CodecInvocation> stages = pipeline.stages();
+
+ // `typedDomain` is true while the running value is still a contiguous
array of column-typed
+ // values. A packing transform or any compression stage ends the typed
domain; once ended, only
+ // further compression stages are legal (a transform could not consume
non-typed input).
+ boolean typedDomain = true;
+ for (CodecInvocation invocation : stages) {
+ ChunkCodecHandler codec = registry.getOrThrow(invocation.name());
+ CodecOptions options = codec.parseOptions(invocation.args());
+ codec.validateContext(options, ctx);
+
+ if (codec.kind() == CodecKind.TRANSFORM) {
+ if (!typedDomain) {
+ throw new IllegalArgumentException(
+ "Transform stage '" + invocation.name() + "' must operate on
column values, but it follows a "
+ + "packing transform or a compression stage in: " +
pipeline.toDslString()
+ + ". A packing transform must be the last transform, and all
transforms "
+ + "must precede any compression stage.");
+ }
+ // Packing transforms emit a non-typed byte stream, so nothing typed
may follow them;
+ // typed-layout-preserving transforms keep the typed domain.
+ if (!codec.preservesTypedValueLayout()) {
+ typedDomain = false;
+ }
+ } else {
+ // COMPRESSION: byte→byte, always permitted; ends the typed domain.
+ typedDomain = false;
+ }
+ }
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecRegistry.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecRegistry.java
new file mode 100644
index 00000000000..c24c9471249
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecRegistry.java
@@ -0,0 +1,118 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import com.google.common.annotations.VisibleForTesting;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.pinot.segment.spi.codec.CodecInvocation;
+
+
+/// Registry of known [ChunkCodecHandler] instances, looked up by
+/// [ChunkCodecHandler#name()].
+///
+/// In v1 the built-in set is **closed**: production always uses the immutable
[#DEFAULT]
+/// instance, which holds the hardcoded built-ins and rejects [#register()].
There is
+/// intentionally no plugin registration path yet — a segment whose header
names an unknown codec
+/// cannot be decoded, so opening the set to third parties requires a
controller-side
+/// version/codec compatibility gate first. The mutable constructor +
[#register()] exist only
+/// for tests.
+///
+/// The mutable registry is not thread-safe for concurrent writes; tests must
complete all
+/// registrations before sharing it across threads.
+///
+/// Lookup is case-insensitive.
+final class CodecRegistry {
+
+ /// Default, immutable registry containing all built-in compression codecs.
+ /// Safe for concurrent reads; does not allow [#register()].
+ static final CodecRegistry DEFAULT;
+
+ static {
+ Map<String, ChunkCodecHandler<?>> m = new LinkedHashMap<>();
+ m.put(ZstdCodecDefinition.INSTANCE.name().toUpperCase(Locale.ROOT),
ZstdCodecDefinition.INSTANCE);
+ // Accept the legacy enum spelling ZSTANDARD as an alias for ZSTD so users
familiar with
+ // FieldConfig.CompressionCodec.ZSTANDARD aren't blocked. Canonicalization
still emits "ZSTD",
+ // so the on-disk name is unaffected.
+ m.put("ZSTANDARD", ZstdCodecDefinition.INSTANCE);
+ m.put(Lz4CodecDefinition.INSTANCE.name().toUpperCase(Locale.ROOT),
Lz4CodecDefinition.INSTANCE);
+ m.put(SnappyCodecDefinition.INSTANCE.name().toUpperCase(Locale.ROOT),
SnappyCodecDefinition.INSTANCE);
+ m.put(GzipCodecDefinition.INSTANCE.name().toUpperCase(Locale.ROOT),
GzipCodecDefinition.INSTANCE);
+ DEFAULT = new CodecRegistry(Collections.unmodifiableMap(m));
+ }
+
+ private final Map<String, ChunkCodecHandler<?>> _codecs;
+ private final boolean _immutable;
+
+ /// Creates a new empty, mutable registry. Intended for tests only —
production code should use
+ /// [#DEFAULT]. Custom codecs must complete all [#register()] calls before
sharing the
+ /// registry across threads.
+ @VisibleForTesting
+ CodecRegistry() {
+ _codecs = new LinkedHashMap<>();
+ _immutable = false;
+ }
+
+ private CodecRegistry(Map<String, ChunkCodecHandler<?>> codecs) {
+ _codecs = codecs;
+ _immutable = true;
+ }
+
+ /// Registers a codec. Throws if a codec with the same name
(case-insensitive) is already present,
+ /// or if this registry is immutable.
+ ///
+ /// @return `this` for fluent chaining
+ @VisibleForTesting
+ CodecRegistry register(ChunkCodecHandler<?> codec) {
+ if (_immutable) {
+ throw new UnsupportedOperationException("CodecRegistry.DEFAULT is
immutable");
+ }
+ String key = codec.name().toUpperCase(Locale.ROOT);
+ // Reuse the public AST node's centralized name validation, including the
removed CODEC
+ // wrapper reservation, instead of duplicating grammar constants in the
runtime module.
+ new CodecInvocation(key, List.of());
+ if (_codecs.containsKey(key)) {
+ throw new IllegalArgumentException("A codec named '" + key + "' is
already registered");
+ }
+ _codecs.put(key, codec);
+ return this;
+ }
+
+ /// Looks up a codec by name (case-insensitive).
+ ///
+ /// @return the matching codec handler, or `null` if not found
+ @Nullable
+ ChunkCodecHandler<?> get(String name) {
+ return _codecs.get(name.toUpperCase(Locale.ROOT));
+ }
+
+ /// Looks up a codec by name, throwing [IllegalArgumentException] if not
found.
+ ChunkCodecHandler<?> getOrThrow(String name) {
+ ChunkCodecHandler<?> codec = get(name);
+ if (codec == null) {
+ throw new IllegalArgumentException(
+ "Unknown codec '" + name + "'. Known codecs: " + _codecs.keySet());
+ }
+ return codec;
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/GzipCodecDefinition.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/GzipCodecDefinition.java
new file mode 100644
index 00000000000..1723a7e5ef9
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/GzipCodecDefinition.java
@@ -0,0 +1,273 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.List;
+import java.util.zip.DataFormatException;
+import java.util.zip.Deflater;
+import java.util.zip.Inflater;
+import org.apache.pinot.segment.spi.memory.CleanerUtil;
+
+
+/// Legacy `GZIP`-named compression codec backed by the default
[java.util.zip.Deflater], which
+/// produces a zlib-wrapped DEFLATE stream rather than RFC 1952 `.gz` framing.
+///
+/// DSL form: `GZIP` — no configuration options.
+///
+/// GZIP is a [CodecKind#COMPRESSION] stage. Compression stages may be chained
after all
+/// transforms.
+///
+/// The name {@value #NAME} is a frozen on-disk API contract stored verbatim
in segment file
+/// headers. It must never be changed or reused for a different algorithm.
+///
+/// Wire format: zlib-wrapped DEFLATE payload (including the zlib checksum)
followed by a 4-byte
+/// big-endian footer containing the uncompressed byte count. The footer
allows decompression
+/// without knowing the original size out-of-band.
+///
+/// **Performance note:** The JDK [Deflater]/[Inflater] instances are reused
per thread and
+/// operate directly on [ByteBuffer] inputs and outputs. For write-intensive
workloads prefer
+/// LZ4 or ZSTD.
+final class GzipCodecDefinition implements
ChunkCodecHandler<GzipCodecDefinition.Options> {
+
+ /// On-disk permanent name stored verbatim in segment file headers.
+ /// This string is a frozen on-disk API contract and must never be changed.
+ public static final String NAME = "GZIP";
+
+ /// Thread-local Deflater/Inflater: reset() between uses to amortize JNI
allocation cost.
+ /// Note: these hold native resources that are only released when the thread
dies. For long-lived
+ /// server thread pools this is bounded by the worker count and acceptable.
+ private static final ThreadLocal<Deflater> DEFLATER =
ThreadLocal.withInitial(Deflater::new);
+ private static final ThreadLocal<Inflater> INFLATER =
ThreadLocal.withInitial(Inflater::new);
+ private static final ThreadLocal<byte[]> COMPLETION_PROBE =
ThreadLocal.withInitial(() -> new byte[1]);
+
+ public static final GzipCodecDefinition INSTANCE = new GzipCodecDefinition();
+
+ /// Singleton options — GZIP has no configurable parameters.
+ public static final Options OPTIONS = new Options();
+
+ private GzipCodecDefinition() {
+ }
+
+ /// Typed options for [GzipCodecDefinition]. GZIP has no configurable
parameters.
+ public static final class Options implements CodecOptions {
+ private Options() {
+ }
+ }
+
+ @Override
+ public String name() {
+ return NAME;
+ }
+
+ @Override
+ public CodecKind kind() {
+ return CodecKind.COMPRESSION;
+ }
+
+ @Override
+ public Options parseOptions(List<String> args) {
+ if (!args.isEmpty()) {
+ throw new IllegalArgumentException("GZIP codec does not accept
arguments, got: " + args);
+ }
+ return OPTIONS;
+ }
+
+ @Override
+ public void validateContext(Options options, CodecContext ctx) {
+ // GZIP can compress any data type; no restriction
+ }
+
+ @Override
+ public String canonicalize(Options options) {
+ return NAME;
+ }
+
+ @Override
+ public ByteBuffer encode(Options options, CodecContext ctx, ByteBuffer src)
throws IOException {
+ int uncompressedSize = src.remaining();
+ ByteBuffer out = ByteBuffer.allocateDirect(maxEncodedSize(options,
uncompressedSize));
+ Deflater deflater = DEFLATER.get();
+ boolean succeeded = false;
+ try {
+ out.limit(out.capacity() - Integer.BYTES);
+ deflater.reset();
+ deflater.setInput(src.duplicate());
+ deflater.finish();
+ while (!deflater.finished()) {
+ if (!out.hasRemaining()) {
+ throw new IOException("GZIP encode exceeded maximum encoded size " +
out.capacity()
+ + " before deflater finished. Segment build aborted.");
+ }
+ int encoded = deflater.deflate(out);
+ if (encoded == 0 && !deflater.finished()) {
+ throw new IOException("GZIP deflater made no progress before
finishing");
+ }
+ }
+ out.limit(out.capacity());
+ out.putInt(uncompressedSize);
+ out.flip();
+ succeeded = true;
+ return out;
+ } finally {
+ deflater.reset();
+ if (!succeeded) {
+ CleanerUtil.cleanQuietly(out);
+ }
+ }
+ }
+
+ @Override
+ public ByteBuffer decode(Options options, CodecContext ctx, ByteBuffer src)
throws IOException {
+ int decompressedSize = readDecompressedSize(src);
+ ByteBuffer out = ByteBuffer.allocateDirect(decompressedSize);
+ boolean succeeded = false;
+ try {
+ inflateInto(src, out, decompressedSize);
+ succeeded = true;
+ return out;
+ } finally {
+ if (!succeeded) {
+ CleanerUtil.cleanQuietly(out);
+ }
+ }
+ }
+
+ @Override
+ public void decodeInto(Options options, CodecContext ctx, ByteBuffer src,
ByteBuffer dst) throws IOException {
+ dst.clear();
+ int decompressedSize = readDecompressedSize(src);
+ if (decompressedSize > dst.capacity()) {
+ throw new IllegalArgumentException(
+ "GZIP: decompressed size " + decompressedSize + " exceeds dst
capacity " + dst.capacity());
+ }
+ inflateInto(src, dst, decompressedSize);
+ }
+
+ @Override
+ public int maxEncodedSize(Options options, int inputSize) {
+ // DEFLATE worst-case expansion + 4-byte appended uncompressed-size footer
+ if (inputSize < 0) {
+ throw new IllegalArgumentException("GZIP inputSize must be non-negative:
" + inputSize);
+ }
+ long bound = (long) inputSize + (inputSize >> 12) + (inputSize >> 14) +
(inputSize >> 25)
+ + 13 + Integer.BYTES;
+ if (bound > Integer.MAX_VALUE) {
+ throw new IllegalArgumentException("GZIP maximum encoded size exceeds
Integer.MAX_VALUE: " + bound);
+ }
+ return (int) bound;
+ }
+
+ @Override
+ public boolean requiresDirectDstBuffer() {
+ return false;
+ }
+
+ // -------------------------------------------------------------------------
+ // Private helpers
+ // -------------------------------------------------------------------------
+
+ private static int readDecompressedSize(ByteBuffer src) throws IOException {
+ int payloadLimit = src.limit();
+ int payloadSize = src.remaining();
+ if (payloadSize < Integer.BYTES) {
+ throw new IOException("GZIP payload too short to contain
uncompressed-size footer: " + payloadSize + " bytes");
+ }
+ int decompressedSize =
+ src.duplicate().order(ByteOrder.BIG_ENDIAN).getInt(payloadLimit -
Integer.BYTES);
+ if (decompressedSize < 0) {
+ throw new IOException("GZIP: invalid decompressed size in footer: " +
decompressedSize);
+ }
+ return CodecBufferUtils.checkDeclaredDecompressedSize(decompressedSize,
"GZIP", "footer");
+ }
+
+ private static void inflateInto(ByteBuffer src, ByteBuffer dst, int
decompressedSize) throws IOException {
+ ByteBuffer compressed = src.duplicate();
+ compressed.limit(src.limit() - Integer.BYTES);
+ Inflater inflater = INFLATER.get();
+ inflater.reset();
+ try {
+ inflater.setInput(compressed);
+ dst.limit(decompressedSize);
+ while (dst.position() < decompressedSize) {
+ int n;
+ try {
+ n = inflater.inflate(dst);
+ } catch (DataFormatException e) {
+ throw new IOException("GZIP decompression failed", e);
+ }
+ if (n == 0) {
+ if (inflater.finished()) {
+ break;
+ }
+ if (inflater.needsInput()) {
+ throw new IOException(
+ "GZIP inflater ran out of input before producing " +
decompressedSize + " bytes (produced "
+ + dst.position() + ")");
+ }
+ if (inflater.needsDictionary()) {
+ throw new IOException("GZIP inflater requires a preset dictionary
(not supported)");
+ }
+ throw new IOException(
+ "GZIP inflater returned 0 bytes with no known cause (inflated so
far: " + dst.position() + " / "
+ + decompressedSize + ")");
+ }
+ }
+ if (dst.position() != decompressedSize) {
+ throw new IOException("GZIP: inflated " + dst.position() + " bytes but
expected " + decompressedSize);
+ }
+
+ // Filling the caller-declared output size does not prove the DEFLATE
stream is complete: a
+ // corrupt footer can under-report the real output, including zero, and
the loop above would
+ // otherwise accept a truncated prefix without consuming or validating
the stream checksum.
+ byte[] completionProbe = COMPLETION_PROBE.get();
+ while (!inflater.finished()) {
+ int n;
+ try {
+ n = inflater.inflate(completionProbe);
+ } catch (DataFormatException e) {
+ throw new IOException("GZIP decompression failed while validating
stream completion", e);
+ }
+ if (n > 0) {
+ throw new IOException(
+ "GZIP stream expands beyond footer-declared size " +
decompressedSize + ". Segment may be corrupt.");
+ }
+ if (inflater.finished()) {
+ break;
+ }
+ if (inflater.needsInput()) {
+ throw new IOException("GZIP stream ended before its
checksum/trailer. Segment may be truncated.");
+ }
+ if (inflater.needsDictionary()) {
+ throw new IOException("GZIP inflater requires a preset dictionary
(not supported)");
+ }
+ throw new IOException("GZIP inflater did not finish after producing
the footer-declared output size");
+ }
+ if (inflater.getRemaining() != 0) {
+ throw new IOException(
+ "GZIP stream has " + inflater.getRemaining() + " trailing
compressed bytes. Segment may be corrupt.");
+ }
+ dst.flip();
+ } finally {
+ inflater.reset();
+ }
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/Lz4CodecDefinition.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/Lz4CodecDefinition.java
new file mode 100644
index 00000000000..4ab17a04594
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/Lz4CodecDefinition.java
@@ -0,0 +1,174 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.List;
+import net.jpountz.lz4.LZ4CompressorWithLength;
+import net.jpountz.lz4.LZ4DecompressorWithLength;
+import net.jpountz.lz4.LZ4Factory;
+
+
+/// Compression codec backed by LZ4 (length-prefixed variant).
+///
+/// DSL form: `LZ4` — no configuration options; the length-prefixed format is
always used
+/// so that decompression does not require the original uncompressed size
out-of-band.
+///
+/// LZ4 is a [CodecKind#COMPRESSION] stage. Compression stages may be chained
after all
+/// transforms.
+///
+/// The name {@value #NAME} is a frozen on-disk API contract stored verbatim
in segment file
+/// headers. It must never be changed or reused for a different algorithm.
+final class Lz4CodecDefinition implements
ChunkCodecHandler<Lz4CodecDefinition.Options> {
+
+ /// On-disk permanent name stored verbatim in segment file headers.
+ /// This string is a frozen on-disk API contract and must never be changed.
+ public static final String NAME = "LZ4";
+
+ public static final Lz4CodecDefinition INSTANCE = new Lz4CodecDefinition();
+
+ /// Singleton options — LZ4 has no configurable parameters.
+ public static final Options OPTIONS = new Options();
+
+ /// Lazy holder so a missing/broken LZ4 native library only fails when LZ4
is actually used,
+ /// rather than at [CodecRegistry#DEFAULT] class-init time (which would
break every
+ /// consumer of the registry, including consumers that never use LZ4).
+ private static final class Native {
+ static final LZ4Factory FACTORY = LZ4Factory.fastestInstance();
+ static final LZ4CompressorWithLength COMPRESSOR = new
LZ4CompressorWithLength(FACTORY.fastCompressor());
+ static final LZ4DecompressorWithLength DECOMPRESSOR = new
LZ4DecompressorWithLength(FACTORY.safeDecompressor());
+ }
+
+ private Lz4CodecDefinition() {
+ }
+
+ /// Typed options for [Lz4CodecDefinition]. LZ4 has no configurable
parameters.
+ public static final class Options implements CodecOptions {
+ private Options() {
+ }
+ }
+
+ @Override
+ public String name() {
+ return NAME;
+ }
+
+ @Override
+ public CodecKind kind() {
+ return CodecKind.COMPRESSION;
+ }
+
+ @Override
+ public Options parseOptions(List<String> args) {
+ if (!args.isEmpty()) {
+ throw new IllegalArgumentException("LZ4 codec does not accept arguments,
got: " + args);
+ }
+ return OPTIONS;
+ }
+
+ @Override
+ public void validateContext(Options options, CodecContext ctx) {
+ // LZ4 can compress any data type; no restriction
+ }
+
+ @Override
+ public String canonicalize(Options options) {
+ return NAME;
+ }
+
+ @Override
+ public ByteBuffer encode(Options options, CodecContext ctx, ByteBuffer src)
throws IOException {
+ // LZ4 JNI requires both buffers be direct; mirror the Snappy/Zstd
defensive copy.
+ ByteBuffer directSrc = CodecBufferUtils.toDirectBuffer(src);
+ ByteBuffer out = null;
+ boolean succeeded = false;
+ try {
+ int maxSize =
Native.COMPRESSOR.maxCompressedLength(directSrc.remaining());
+ out = ByteBuffer.allocateDirect(maxSize);
+ Native.COMPRESSOR.compress(directSrc, out);
+ out.flip();
+ succeeded = true;
+ return out;
+ } finally {
+ CodecBufferUtils.cleanDirectCopy(src, directSrc);
+ if (!succeeded) {
+ CodecBufferUtils.cleanQuietly(out);
+ }
+ }
+ }
+
+ @Override
+ public ByteBuffer decode(Options options, CodecContext ctx, ByteBuffer src)
throws IOException {
+ ByteBuffer directSrc = CodecBufferUtils.toDirectBuffer(src);
+ ByteBuffer out = null;
+ boolean succeeded = false;
+ try {
+ int decompressedLength = CodecBufferUtils.checkDeclaredDecompressedSize(
+ LZ4DecompressorWithLength.getDecompressedLength(directSrc), "LZ4",
"length prefix");
+ out = ByteBuffer.allocateDirect(decompressedLength);
+ Native.DECOMPRESSOR.decompress(directSrc, out);
+ if (out.position() != decompressedLength) {
+ throw new IOException("LZ4 decoded " + out.position() + " bytes but
expected " + decompressedLength
+ + ". Segment may be corrupt.");
+ }
+ out.flip();
+ succeeded = true;
+ return out;
+ } finally {
+ CodecBufferUtils.cleanDirectCopy(src, directSrc);
+ if (!succeeded) {
+ CodecBufferUtils.cleanQuietly(out);
+ }
+ }
+ }
+
+ @Override
+ public void decodeInto(Options options, CodecContext ctx, ByteBuffer src,
ByteBuffer dst) throws IOException {
+ dst.clear();
+ ByteBuffer directSrc = CodecBufferUtils.toDirectBuffer(src);
+ try {
+ int decompressedLength = CodecBufferUtils.checkDeclaredDecompressedSize(
+ LZ4DecompressorWithLength.getDecompressedLength(directSrc), "LZ4",
"length prefix");
+ if (decompressedLength > dst.capacity()) {
+ throw new IllegalArgumentException(
+ "LZ4: decompressed size " + decompressedLength + " exceeds dst
capacity " + dst.capacity());
+ }
+ Native.DECOMPRESSOR.decompress(directSrc, dst);
+ if (dst.position() != decompressedLength) {
+ throw new IOException("LZ4 decoded " + dst.position() + " bytes but
expected " + decompressedLength
+ + ". Segment may be corrupt.");
+ }
+ dst.flip();
+ } finally {
+ CodecBufferUtils.cleanDirectCopy(src, directSrc);
+ }
+ }
+
+ @Override
+ public int maxEncodedSize(Options options, int inputSize) {
+ // LZ4CompressorWithLength adds a 4-byte length prefix before the
compressed payload
+ return Native.FACTORY.fastCompressor().maxCompressedLength(inputSize) +
Integer.BYTES;
+ }
+
+ @Override
+ public boolean requiresDirectDstBuffer() {
+ return false;
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/SnappyCodecDefinition.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/SnappyCodecDefinition.java
new file mode 100644
index 00000000000..6a28c480e77
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/SnappyCodecDefinition.java
@@ -0,0 +1,175 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.List;
+import org.xerial.snappy.Snappy;
+
+
+/// Compression codec backed by Snappy (via the `org.xerial.snappy` JNI
library).
+///
+/// DSL form: `SNAPPY` — no configuration options.
+///
+/// Snappy is a [CodecKind#COMPRESSION] stage. Compression stages may be
chained after all
+/// transforms.
+///
+/// The name {@value #NAME} is a frozen on-disk API contract stored verbatim
in segment file
+/// headers. It must never be changed or reused for a different algorithm.
+///
+/// Note: [#decodeInto()] requires a direct [ByteBuffer] for `dst` because the
+/// Snappy JNI `uncompress(ByteBuffer, ByteBuffer)` overload requires both
buffers to be direct.
+/// Also, that JNI overload writes at `dst.position()` but does *not* advance
it; the
+/// returned byte count is used to set `dst.limit()` explicitly.
+final class SnappyCodecDefinition implements
ChunkCodecHandler<SnappyCodecDefinition.Options> {
+
+ /// On-disk permanent name stored verbatim in segment file headers.
+ /// This string is a frozen on-disk API contract and must never be changed.
+ public static final String NAME = "SNAPPY";
+
+ public static final SnappyCodecDefinition INSTANCE = new
SnappyCodecDefinition();
+
+ /// Singleton options — Snappy has no configurable parameters.
+ public static final Options OPTIONS = new Options();
+
+ private SnappyCodecDefinition() {
+ }
+
+ /// Typed options for [SnappyCodecDefinition]. Snappy has no configurable
parameters.
+ public static final class Options implements CodecOptions {
+ private Options() {
+ }
+ }
+
+ @Override
+ public String name() {
+ return NAME;
+ }
+
+ @Override
+ public CodecKind kind() {
+ return CodecKind.COMPRESSION;
+ }
+
+ @Override
+ public Options parseOptions(List<String> args) {
+ if (!args.isEmpty()) {
+ throw new IllegalArgumentException("SNAPPY codec does not accept
arguments, got: " + args);
+ }
+ return OPTIONS;
+ }
+
+ @Override
+ public void validateContext(Options options, CodecContext ctx) {
+ // Snappy can compress any data type; no restriction
+ }
+
+ @Override
+ public String canonicalize(Options options) {
+ return NAME;
+ }
+
+ @Override
+ public ByteBuffer encode(Options options, CodecContext ctx, ByteBuffer src)
throws IOException {
+ ByteBuffer directSrc = CodecBufferUtils.toDirectBuffer(src);
+ ByteBuffer out = null;
+ boolean succeeded = false;
+ try {
+ int maxSize = Snappy.maxCompressedLength(directSrc.remaining());
+ out = ByteBuffer.allocateDirect(maxSize);
+ int compressedSize = Snappy.compress(directSrc, out);
+ out.limit(compressedSize);
+ out.position(0);
+ succeeded = true;
+ return out;
+ } finally {
+ CodecBufferUtils.cleanDirectCopy(src, directSrc);
+ if (!succeeded) {
+ CodecBufferUtils.cleanQuietly(out);
+ }
+ }
+ }
+
+ @Override
+ public ByteBuffer decode(Options options, CodecContext ctx, ByteBuffer src)
throws IOException {
+ // Snappy JNI requires direct buffers — convert before reading the size
header so the heap
+ // fallback path actually works (Snappy.uncompressedLength does not accept
heap input).
+ ByteBuffer directSrc = CodecBufferUtils.toDirectBuffer(src);
+ ByteBuffer out = null;
+ boolean succeeded = false;
+ try {
+ int decompressedSize = CodecBufferUtils.checkDeclaredDecompressedSize(
+ Snappy.uncompressedLength(directSrc.duplicate()), "Snappy", "varint
length header");
+ out = ByteBuffer.allocateDirect(decompressedSize);
+ // Snappy JNI writes at dst.position() but does NOT advance it; use
returned count to set limit
+ int written = Snappy.uncompress(directSrc, out);
+ if (written != decompressedSize) {
+ throw new IOException(
+ "Snappy decode size mismatch: expected " + decompressedSize + ",
got " + written
+ + ". Segment may be corrupt.");
+ }
+ out.limit(written);
+ out.position(0);
+ succeeded = true;
+ return out;
+ } finally {
+ CodecBufferUtils.cleanDirectCopy(src, directSrc);
+ if (!succeeded) {
+ CodecBufferUtils.cleanQuietly(out);
+ }
+ }
+ }
+
+ @Override
+ public void decodeInto(Options options, CodecContext ctx, ByteBuffer src,
ByteBuffer dst) throws IOException {
+ dst.clear();
+ // Snappy JNI requires direct buffers — convert before reading the size
header.
+ ByteBuffer directSrc = CodecBufferUtils.toDirectBuffer(src);
+ try {
+ int decompressedSize = CodecBufferUtils.checkDeclaredDecompressedSize(
+ Snappy.uncompressedLength(directSrc.duplicate()), "Snappy", "varint
length header");
+ if (decompressedSize > dst.capacity()) {
+ throw new IllegalArgumentException(
+ "Snappy: decompressed size " + decompressedSize + " exceeds dst
capacity " + dst.capacity());
+ }
+ // Snappy JNI writes at dst.position() but does NOT advance it; use
returned count to set limit
+ int written = Snappy.uncompress(directSrc, dst);
+ if (written != decompressedSize) {
+ throw new IOException(
+ "Snappy decode size mismatch: expected " + decompressedSize + ",
got " + written
+ + ". Segment may be corrupt.");
+ }
+ dst.position(0);
+ dst.limit(written);
+ } finally {
+ CodecBufferUtils.cleanDirectCopy(src, directSrc);
+ }
+ }
+
+ @Override
+ public int maxEncodedSize(Options options, int inputSize) {
+ return Snappy.maxCompressedLength(inputSize);
+ }
+
+ @Override
+ public boolean requiresDirectDstBuffer() {
+ return true;
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/ZstdCodecDefinition.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/ZstdCodecDefinition.java
new file mode 100644
index 00000000000..559b39475c4
--- /dev/null
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/ZstdCodecDefinition.java
@@ -0,0 +1,238 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import com.github.luben.zstd.Zstd;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Objects;
+
+
+/// Compression codec backed by Zstandard (Zstd) with a configurable
compression level.
+///
+/// DSL forms:
+///
+/// - `ZSTD` — uses default level 3
+/// - `ZSTD(3)` — explicit level in the range `1` through
[Zstd#maxCompressionLevel()]
+///
+/// Zstd also supports negative fast-compression levels, but the initial codec
DSL deliberately
+/// accepts unsigned integer arguments only. Those levels are therefore
outside this version's
+/// public codec contract. Level `0` (zstd's alias for "use the default
level") is rejected so that
+/// each behavior has exactly one canonical spelling; use `ZSTD` or an
explicit level instead.
+///
+/// ZSTD is a [CodecKind#COMPRESSION] stage. Compression stages may be chained
after all
+/// transforms.
+///
+/// Both encode and decode use the Zstd JNI, which requires direct
[ByteBuffer]s.
+final class ZstdCodecDefinition implements
ChunkCodecHandler<ZstdCodecDefinition.Options> {
+
+ /// On-disk permanent name stored verbatim in segment file headers.
+ /// This string is a frozen on-disk API contract and must never be changed.
+ public static final String NAME = "ZSTD";
+
+ public static final ZstdCodecDefinition INSTANCE = new ZstdCodecDefinition();
+
+ /// Default compression level when none is specified.
+ public static final int DEFAULT_LEVEL = 3;
+
+ private ZstdCodecDefinition() {
+ }
+
+ /// Typed options for [ZstdCodecDefinition].
+ public static final class Options implements CodecOptions {
+ private final int _level;
+
+ public Options(int level) {
+ _level = level;
+ }
+
+ /// Returns the Zstd compression level.
+ public int getLevel() {
+ return _level;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof Options)) {
+ return false;
+ }
+ return _level == ((Options) o)._level;
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(_level);
+ }
+ }
+
+ @Override
+ public String name() {
+ return NAME;
+ }
+
+ @Override
+ public CodecKind kind() {
+ return CodecKind.COMPRESSION;
+ }
+
+ @Override
+ public Options parseOptions(List<String> args) {
+ if (args.isEmpty()) {
+ return new Options(DEFAULT_LEVEL);
+ }
+ if (args.size() != 1) {
+ throw new IllegalArgumentException("ZSTD codec accepts at most one
argument (compression level), got: " + args);
+ }
+ int level;
+ try {
+ level = Integer.parseInt(args.get(0));
+ } catch (NumberFormatException e) {
+ throw new IllegalArgumentException("ZSTD codec level must be an integer,
got: " + args.get(0));
+ }
+ // Level 0 is rejected rather than accepted: zstd treats 0 as "use the
default level", which would give
+ // ZSTD(0) and ZSTD(3) identical behavior under two different canonical
spellings. The canonical spec is
+ // frozen into segment headers, so each behavior must have exactly one
spelling; use ZSTD or ZSTD(3).
+ int minLevel = 1;
+ int maxLevel = Zstd.maxCompressionLevel();
+ if (level < minLevel || level > maxLevel) {
+ throw new IllegalArgumentException(
+ "ZSTD level " + level + " is out of range [" + minLevel + ", " +
maxLevel + "]");
+ }
+ return new Options(level);
+ }
+
+ @Override
+ public void validateContext(Options options, CodecContext ctx) {
+ // ZSTD can compress any data type; no restriction
+ }
+
+ @Override
+ public String canonicalize(Options options) {
+ return NAME + "(" + options.getLevel() + ")";
+ }
+
+ @Override
+ public ByteBuffer encode(Options options, CodecContext ctx, ByteBuffer src)
throws IOException {
+ ByteBuffer directSrc = CodecBufferUtils.toDirectBuffer(src);
+ ByteBuffer out = null;
+ boolean succeeded = false;
+ try {
+ long bound = Zstd.compressBound(directSrc.remaining());
+ if (bound > Integer.MAX_VALUE) {
+ throw new IOException("Zstd compressBound " + bound + " exceeds
Integer.MAX_VALUE for input of "
+ + directSrc.remaining() + " bytes");
+ }
+ out = ByteBuffer.allocateDirect((int) bound);
+ long result = Zstd.compress(out, directSrc, options.getLevel());
+ if (Zstd.isError(result)) {
+ throw new IOException("Zstd compression failed: " +
Zstd.getErrorName(result));
+ }
+ out.flip();
+ succeeded = true;
+ return out;
+ } finally {
+ CodecBufferUtils.cleanDirectCopy(src, directSrc);
+ if (!succeeded) {
+ CodecBufferUtils.cleanQuietly(out);
+ }
+ }
+ }
+
+ @Override
+ public ByteBuffer decode(Options options, CodecContext ctx, ByteBuffer src)
throws IOException {
+ ByteBuffer directSrc = CodecBufferUtils.toDirectBuffer(src);
+ ByteBuffer out = null;
+ boolean succeeded = false;
+ try {
+ long decompressedSize = Zstd.getFrameContentSize(directSrc);
+ // Zstd uses negative sentinel values for an unknown or invalid content
size. Zero is a known,
+ // valid content size for an empty frame and must be allowed through to
decompression.
+ if (decompressedSize < 0) {
+ throw new IOException("Zstd: cannot determine decompressed size from
frame header");
+ }
+ out = ByteBuffer.allocateDirect(
+ CodecBufferUtils.checkDeclaredDecompressedSize(decompressedSize,
"Zstd", "frame header"));
+ long result = Zstd.decompress(out, directSrc);
+ if (Zstd.isError(result)) {
+ throw new IOException("Zstd decompression failed: " +
Zstd.getErrorName(result));
+ }
+ if (result != decompressedSize) {
+ throw new IOException("Zstd decoded " + result + " bytes but frame
declared " + decompressedSize
+ + ". Segment may be corrupt.");
+ }
+ out.flip();
+ succeeded = true;
+ return out;
+ } finally {
+ CodecBufferUtils.cleanDirectCopy(src, directSrc);
+ if (!succeeded) {
+ CodecBufferUtils.cleanQuietly(out);
+ }
+ }
+ }
+
+ @Override
+ public void decodeInto(Options options, CodecContext ctx, ByteBuffer src,
ByteBuffer dst) throws IOException {
+ dst.clear();
+ ByteBuffer directSrc = CodecBufferUtils.toDirectBuffer(src);
+ try {
+ long declaredDecompressedSize = Zstd.getFrameContentSize(directSrc);
+ // As in decode(), zero is a valid known size; only Zstd's negative
sentinel values are errors.
+ if (declaredDecompressedSize < 0) {
+ throw new IOException("Zstd: cannot determine decompressed size from
frame header");
+ }
+ int decompressedSize = CodecBufferUtils.checkDeclaredDecompressedSize(
+ declaredDecompressedSize, "Zstd", "frame header");
+ if (decompressedSize > dst.capacity()) {
+ throw new IllegalArgumentException(
+ "Zstd: decompressed size " + decompressedSize + " exceeds dst
capacity " + dst.capacity());
+ }
+ long result = Zstd.decompress(dst, directSrc);
+ if (Zstd.isError(result)) {
+ throw new IOException("Zstd decompression failed: " +
Zstd.getErrorName(result));
+ }
+ if (result != decompressedSize) {
+ throw new IOException("Zstd decoded " + result + " bytes but frame
declared " + decompressedSize
+ + ". Segment may be corrupt.");
+ }
+ dst.flip();
+ } finally {
+ CodecBufferUtils.cleanDirectCopy(src, directSrc);
+ }
+ }
+
+ @Override
+ public int maxEncodedSize(Options options, int inputSize) {
+ long bound = Zstd.compressBound(inputSize);
+ if (bound > Integer.MAX_VALUE) {
+ throw new IllegalArgumentException(
+ "Zstd compressBound " + bound + " exceeds Integer.MAX_VALUE for
inputSize=" + inputSize);
+ }
+ return (int) bound;
+ }
+
+ @Override
+ public boolean requiresDirectDstBuffer() {
+ return true;
+ }
+}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CodecPipelineExecutorTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CodecPipelineExecutorTest.java
new file mode 100644
index 00000000000..ad9bd363c1c
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CodecPipelineExecutorTest.java
@@ -0,0 +1,254 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.io.IOException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Modifier;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
+
+
+/// Format-independent round-trip and resource-bound tests for
[CodecPipelineExecutor].
+public class CodecPipelineExecutorTest {
+ private static final byte[] INPUT = createInput();
+
+ @DataProvider(name = "singleCodecAndInputKind")
+ public Object[][] singleCodecAndInputKind() {
+ return new Object[][]{
+ {"LZ4", false}, {"LZ4", true},
+ {"SNAPPY", false}, {"SNAPPY", true},
+ {"GZIP", false}, {"GZIP", true},
+ {"ZSTD", false}, {"ZSTD", true},
+ {"ZSTD(1)", false}, {"ZSTD(1)", true}
+ };
+ }
+
+ @Test(dataProvider = "singleCodecAndInputKind")
+ public void testSingleCodecRoundTrip(String spec, boolean directInput)
throws Exception {
+ CodecPipelineExecutor executor = CodecPipelineExecutor.create(spec,
DataType.BYTES);
+ ByteBuffer encoded = executor.encode(inputBuffer(directInput));
+
+ assertTrue(encoded.remaining() <= executor.maxEncodedSize(INPUT.length));
+ assertBytesEqual(executor.decode(encoded.duplicate()));
+
+ ByteBuffer destination = ByteBuffer.allocateDirect(INPUT.length);
+ int maxIntermediateSize = executor.maxEncodedSize(INPUT.length);
+ executor.decode(encoded.duplicate(), destination, INPUT.length,
maxIntermediateSize,
+ (long) maxIntermediateSize * 32);
+ assertBytesEqual(destination);
+ }
+
+ @Test
+ public void testMultiCompressionRoundTrips() throws Exception {
+ for (String spec : new String[]{"LZ4,SNAPPY", "GZIP,ZSTD(5)",
"SNAPPY,LZ4,GZIP"}) {
+ CodecPipelineExecutor executor = CodecPipelineExecutor.create(spec,
DataType.BYTES);
+ ByteBuffer encoded = executor.encode(inputBuffer(false));
+ assertBytesEqual(executor.decode(encoded.duplicate()));
+
+ ByteBuffer destination = ByteBuffer.allocateDirect(INPUT.length);
+ int maxIntermediateSize = executor.maxEncodedSize(INPUT.length);
+ executor.decode(encoded.duplicate(), destination, INPUT.length,
maxIntermediateSize,
+ (long) maxIntermediateSize * 32);
+ assertBytesEqual(destination);
+ }
+ }
+
+ @Test
+ public void testCallerOwnedScratchReusesDirectBuffers()
+ throws Exception {
+ CodecPipelineExecutor executor =
CodecPipelineExecutor.create("LZ4,SNAPPY,GZIP", DataType.BYTES);
+ ByteBuffer encoded = executor.encode(inputBuffer(false));
+ int maxIntermediateSize = executor.maxEncodedSize(INPUT.length);
+ try (CodecPipelineExecutor.DecodeScratch scratch = new
CodecPipelineExecutor.DecodeScratch()) {
+ ByteBuffer firstDestination = ByteBuffer.allocateDirect(INPUT.length);
+ executor.decode(encoded.duplicate(), firstDestination, INPUT.length,
maxIntermediateSize,
+ (long) maxIntermediateSize * 32, scratch);
+ assertBytesEqual(firstDestination);
+ assertEquals(scratch.allocationCount(), 2);
+
+ ByteBuffer secondDestination = ByteBuffer.allocateDirect(INPUT.length);
+ executor.decode(encoded.duplicate(), secondDestination, INPUT.length,
maxIntermediateSize,
+ (long) maxIntermediateSize * 32, scratch);
+ assertBytesEqual(secondDestination);
+ assertEquals(scratch.allocationCount(), 2, "Repeated decode must reuse
both ping-pong buffers");
+ }
+ }
+
+ @Test
+ public void testCanonicalNamesAliasesAndOptions() {
+ assertEquals(CodecPipelineExecutor.create("zstd",
DataType.INT).getCanonicalSpec(), "ZSTD(3)");
+ assertEquals(CodecPipelineExecutor.create("zstandard",
DataType.INT).getCanonicalSpec(), "ZSTD(3)");
+ // Zstd treats level 0 as "use the default level", which would give
ZSTD(0) and ZSTD(3) identical
+ // behavior under two canonical spellings; it is rejected so each behavior
has exactly one spelling.
+ assertThrows(IllegalArgumentException.class, () ->
CodecPipelineExecutor.create("zstd(0)", DataType.INT));
+ assertEquals(CodecPipelineExecutor.create("zstd(5),gzip",
DataType.INT).getCanonicalSpec(),
+ "ZSTD(5),GZIP");
+ assertEquals(CodecPipelineExecutor.create("lz4,snappy",
DataType.INT).getCanonicalSpec(),
+ "LZ4,SNAPPY");
+ }
+
+ @Test
+ public void testCompressionClassification() {
+ assertTrue(CodecPipelineExecutor.create("LZ4",
DataType.INT).isCompressed());
+ assertTrue(CodecPipelineExecutor.create("GZIP,ZSTD",
DataType.STRING).isCompressed());
+ }
+
+ @Test
+ public void testDirectDestinationRequirement() throws Exception {
+ for (String spec : new String[]{"SNAPPY", "ZSTD", "SNAPPY,GZIP"}) {
+ CodecPipelineExecutor executor = CodecPipelineExecutor.create(spec,
DataType.BYTES);
+ ByteBuffer encoded = executor.encode(inputBuffer(false));
+ IllegalArgumentException exception =
expectThrows(IllegalArgumentException.class,
+ () -> executor.decode(encoded.duplicate(),
ByteBuffer.allocate(INPUT.length), INPUT.length));
+ assertTrue(exception.getMessage().contains("requires a direct
ByteBuffer"), exception.getMessage());
+ }
+ }
+
+ @Test
+ public void testHeapDestinationWhenNoStageRequiresDirectBuffer() throws
Exception {
+ for (String spec : new String[]{"LZ4", "GZIP", "LZ4,GZIP", "GZIP,SNAPPY"})
{
+ CodecPipelineExecutor executor = CodecPipelineExecutor.create(spec,
DataType.BYTES);
+ ByteBuffer encoded = executor.encode(inputBuffer(true));
+ ByteBuffer destination = ByteBuffer.allocate(INPUT.length);
+ executor.decode(encoded.duplicate(), destination, INPUT.length,
executor.maxEncodedSize(INPUT.length));
+ assertBytesEqual(destination);
+ }
+ }
+
+ @Test
+ public void testComposedMaximumEncodedSizeBoundsActualOutput() throws
Exception {
+ CodecPipelineExecutor executor =
CodecPipelineExecutor.create("LZ4,SNAPPY,GZIP,ZSTD", DataType.BYTES);
+ int bound = executor.maxEncodedSize(INPUT.length);
+ ByteBuffer encoded = executor.encode(inputBuffer(true));
+ assertTrue(encoded.remaining() <= bound, encoded.remaining() + " > " +
bound);
+ }
+
+ @Test
+ public void testPerStageMaximumEncodedSizeCap() {
+ CodecPipelineExecutor executor =
CodecPipelineExecutor.create("LZ4,SNAPPY", DataType.BYTES);
+ IllegalArgumentException exception =
expectThrows(IllegalArgumentException.class,
+ () -> executor.maxEncodedSize(INPUT.length, INPUT.length));
+ assertTrue(exception.getMessage().contains("maximum encoded size"),
exception.getMessage());
+ }
+
+ @Test
+ public void testCumulativeMaximumEncodedSizeCap() {
+ CodecPipelineExecutor executor =
CodecPipelineExecutor.create("LZ4,SNAPPY", DataType.BYTES);
+ IllegalArgumentException exception =
expectThrows(IllegalArgumentException.class,
+ () -> executor.maxEncodedSize(INPUT.length, Integer.MAX_VALUE,
INPUT.length));
+ assertTrue(exception.getMessage().contains("cumulative stage-output
bound"), exception.getMessage());
+ }
+
+ @Test
+ public void testDecodePerStageIntermediateCap() throws Exception {
+ CodecPipelineExecutor executor =
CodecPipelineExecutor.create("LZ4,SNAPPY", DataType.BYTES);
+ ByteBuffer encoded = executor.encode(inputBuffer(false));
+ IllegalArgumentException exception =
expectThrows(IllegalArgumentException.class,
+ () -> executor.decode(encoded.duplicate(),
ByteBuffer.allocateDirect(INPUT.length), INPUT.length,
+ INPUT.length));
+ assertTrue(exception.getMessage().contains("maximum encoded size"),
exception.getMessage());
+ }
+
+ @Test
+ public void testDecodeCumulativeIntermediateCap() throws Exception {
+ CodecPipelineExecutor executor =
CodecPipelineExecutor.create("LZ4,SNAPPY", DataType.BYTES);
+ ByteBuffer encoded = executor.encode(inputBuffer(false));
+ IllegalArgumentException exception =
expectThrows(IllegalArgumentException.class,
+ () -> executor.decode(encoded.duplicate(),
ByteBuffer.allocateDirect(INPUT.length), INPUT.length,
+ Integer.MAX_VALUE, INPUT.length));
+ assertTrue(exception.getMessage().contains("cumulative stage-output
bound"), exception.getMessage());
+ }
+
+ @Test
+ public void testDecodeRejectsUnexpectedFinalSize() throws Exception {
+ CodecPipelineExecutor executor = CodecPipelineExecutor.create("GZIP",
DataType.BYTES);
+ ByteBuffer encoded = executor.encode(inputBuffer(false));
+ int expectedDecodedSize = INPUT.length - 1;
+ int maxIntermediateSize = executor.maxEncodedSize(INPUT.length);
+ IOException exception = expectThrows(IOException.class,
+ () -> executor.decode(encoded.duplicate(),
ByteBuffer.allocate(INPUT.length), expectedDecodedSize,
+ maxIntermediateSize, Long.MAX_VALUE));
+ assertTrue(exception.getMessage().contains("but expected " +
expectedDecodedSize), exception.getMessage());
+ }
+
+ @Test
+ public void testInvalidMaximumSizeArguments() {
+ CodecPipelineExecutor executor = CodecPipelineExecutor.create("LZ4",
DataType.BYTES);
+ assertThrows(IllegalArgumentException.class, () ->
executor.maxEncodedSize(-1));
+ assertThrows(IllegalArgumentException.class,
+ () -> executor.maxEncodedSize(INPUT.length, INPUT.length - 1));
+ assertThrows(IllegalArgumentException.class,
+ () -> executor.maxEncodedSize(INPUT.length, Integer.MAX_VALUE,
INPUT.length - 1L));
+ }
+
+ @Test
+ public void testPublicRuntimeSurfaceIsBounded() throws
ReflectiveOperationException {
+ assertTrue(Modifier.isPublic(CodecPipelineExecutor.class.getModifiers()));
+
assertTrue(Modifier.isPublic(CodecPipelineExecutor.DecodeScratch.class.getModifiers()));
+ for (Class<?> closedType : new Class<?>[]{
+ ChunkCodecHandler.class, CodecContext.class, CodecDefinition.class,
CodecKind.class, CodecOptions.class,
+ CodecRegistry.class, CodecPipelineValidator.class,
Lz4CodecDefinition.class, SnappyCodecDefinition.class,
+ GzipCodecDefinition.class, ZstdCodecDefinition.class}) {
+ assertFalse(Modifier.isPublic(closedType.getModifiers()),
closedType.getName());
+ }
+
+ Method publicFactory =
CodecPipelineExecutor.class.getDeclaredMethod("create", String.class,
DataType.class);
+ assertTrue(Modifier.isPublic(publicFactory.getModifiers()));
+ Method testFactory = CodecPipelineExecutor.class.getDeclaredMethod(
+ "create", String.class, CodecContext.class, CodecRegistry.class);
+ assertFalse(Modifier.isPublic(testFactory.getModifiers()));
+ Method allocationReturningDecode =
CodecPipelineExecutor.class.getDeclaredMethod("decode", ByteBuffer.class);
+ assertFalse(Modifier.isPublic(allocationReturningDecode.getModifiers()));
+ Method boundedDecode =
CodecPipelineExecutor.class.getDeclaredMethod("decode", ByteBuffer.class,
+ ByteBuffer.class, int.class, int.class, long.class);
+ assertFalse(Modifier.isPublic(boundedDecode.getModifiers()));
+ Method reusableBoundedDecode =
CodecPipelineExecutor.class.getDeclaredMethod("decode", ByteBuffer.class,
+ ByteBuffer.class, int.class, int.class, long.class,
CodecPipelineExecutor.DecodeScratch.class);
+ assertTrue(Modifier.isPublic(reusableBoundedDecode.getModifiers()));
+ }
+
+ private static ByteBuffer inputBuffer(boolean direct) {
+ ByteBuffer buffer = direct ? ByteBuffer.allocateDirect(INPUT.length) :
ByteBuffer.allocate(INPUT.length);
+ return buffer.put(INPUT).flip();
+ }
+
+ private static void assertBytesEqual(ByteBuffer actual) {
+ byte[] bytes = new byte[actual.remaining()];
+ actual.duplicate().get(bytes);
+ assertTrue(Arrays.equals(bytes, INPUT));
+ }
+
+ private static byte[] createInput() {
+ byte[] input = new byte[16 * 1024];
+ for (int i = 0; i < input.length; i++) {
+ input[i] = (byte) ((i * 31) ^ (i >>> 3));
+ }
+ return input;
+ }
+}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CodecPipelineValidatorTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CodecPipelineValidatorTest.java
new file mode 100644
index 00000000000..807e3c5e1bb
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CodecPipelineValidatorTest.java
@@ -0,0 +1,188 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+import org.apache.pinot.segment.spi.codec.CodecPipeline;
+import org.apache.pinot.segment.spi.codec.CodecSpecParser;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+
+
+/// Tests for [CodecPipelineValidator]. Synthetic transform handlers keep this
layer independent
+/// from the concrete transform implementations introduced by later stacked
pull requests.
+public class CodecPipelineValidatorTest {
+ private static final TestCodecHandler TYPED_TRANSFORM =
+ new TestCodecHandler("TYPED", CodecKind.TRANSFORM, true);
+ private static final TestCodecHandler PACKING_TRANSFORM =
+ new TestCodecHandler("PACKING", CodecKind.TRANSFORM, false);
+ private static final CodecRegistry REGISTRY = new CodecRegistry()
+ .register(TYPED_TRANSFORM)
+ .register(PACKING_TRANSFORM)
+ .register(Lz4CodecDefinition.INSTANCE)
+ .register(SnappyCodecDefinition.INSTANCE)
+ .register(GzipCodecDefinition.INSTANCE)
+ .register(ZstdCodecDefinition.INSTANCE);
+
+ @Test
+ public void testCompressionPipelinesAreValid() {
+ validate("LZ4", DataType.INT);
+ validate("ZSTD(3)", DataType.LONG);
+ validate("ZSTD", DataType.STRING);
+ validate("LZ4,GZIP,SNAPPY,ZSTD(5)", DataType.BYTES);
+ }
+
+ @Test
+ public void testTypedAndPackingTransformOrdering() {
+ validate("TYPED", DataType.INT);
+ validate("TYPED,TYPED,LZ4", DataType.INT);
+ validate("TYPED,PACKING,ZSTD(3)", DataType.LONG);
+ validate("PACKING,SNAPPY", DataType.STRING);
+ }
+
+ @Test(expectedExceptions = IllegalArgumentException.class,
+ expectedExceptionsMessageRegExp = ".*must operate on column values.*")
+ public void testTransformAfterCompressionRejected() {
+ validate("ZSTD(3),TYPED", DataType.INT);
+ }
+
+ @Test(expectedExceptions = IllegalArgumentException.class,
+ expectedExceptionsMessageRegExp = ".*must operate on column values.*")
+ public void testTransformAfterPackingTransformRejected() {
+ validate("PACKING,TYPED,LZ4", DataType.INT);
+ }
+
+ @Test(expectedExceptions = IllegalArgumentException.class,
+ expectedExceptionsMessageRegExp = ".*must operate on column values.*")
+ public void testTwoPackingTransformsRejected() {
+ validate("PACKING,PACKING", DataType.LONG);
+ }
+
+ @Test(expectedExceptions = IllegalArgumentException.class,
+ expectedExceptionsMessageRegExp = ".*Unknown codec.*")
+ public void testUnknownCodec() {
+ validate("NOSUCHCODEC", DataType.INT);
+ }
+
+ @Test(expectedExceptions = IllegalArgumentException.class,
+ expectedExceptionsMessageRegExp = ".*out of range.*")
+ public void testZstdLevelTooHigh() {
+ validate("ZSTD(99)", DataType.INT);
+ }
+
+ @Test
+ public void testZstdNegativeLevelIsOutsideTheUnsignedDslContract() {
+ assertThrows(IllegalArgumentException.class,
+ () -> ZstdCodecDefinition.INSTANCE.parseOptions(List.of("-1")));
+ }
+
+ @Test(expectedExceptions = IllegalArgumentException.class)
+ public void testZstdBadArgCount() {
+ validate("ZSTD(3,4)", DataType.INT);
+ }
+
+ @Test
+ public void testTypedValueLayoutContract() {
+ assertTrue(TYPED_TRANSFORM.preservesTypedValueLayout());
+ assertFalse(PACKING_TRANSFORM.preservesTypedValueLayout());
+ assertFalse(Lz4CodecDefinition.INSTANCE.preservesTypedValueLayout());
+ }
+
+ private static void validate(String spec, DataType dataType) {
+ CodecPipeline pipeline = CodecSpecParser.parse(spec);
+ CodecPipelineValidator.validate(pipeline, REGISTRY, new
CodecContext(dataType));
+ }
+
+ private static final class TestCodecHandler implements
ChunkCodecHandler<CodecOptions> {
+ private static final CodecOptions OPTIONS = new CodecOptions() {
+ };
+
+ private final String _name;
+ private final CodecKind _kind;
+ private final boolean _preservesTypedValueLayout;
+
+ private TestCodecHandler(String name, CodecKind kind, boolean
preservesTypedValueLayout) {
+ _name = name;
+ _kind = kind;
+ _preservesTypedValueLayout = preservesTypedValueLayout;
+ }
+
+ @Override
+ public String name() {
+ return _name;
+ }
+
+ @Override
+ public CodecKind kind() {
+ return _kind;
+ }
+
+ @Override
+ public CodecOptions parseOptions(List<String> args) {
+ if (!args.isEmpty()) {
+ throw new IllegalArgumentException(_name + " does not accept
arguments");
+ }
+ return OPTIONS;
+ }
+
+ @Override
+ public void validateContext(CodecOptions options, CodecContext ctx) {
+ }
+
+ @Override
+ public String canonicalize(CodecOptions options) {
+ return _name;
+ }
+
+ @Override
+ public boolean preservesTypedValueLayout() {
+ return _preservesTypedValueLayout;
+ }
+
+ @Override
+ public ByteBuffer encode(CodecOptions options, CodecContext ctx,
ByteBuffer src) {
+ throw new UnsupportedOperationException("Validation-only test codec");
+ }
+
+ @Override
+ public ByteBuffer decode(CodecOptions options, CodecContext ctx,
ByteBuffer src) {
+ throw new UnsupportedOperationException("Validation-only test codec");
+ }
+
+ @Override
+ public void decodeInto(CodecOptions options, CodecContext ctx, ByteBuffer
src, ByteBuffer dst) {
+ throw new UnsupportedOperationException("Validation-only test codec");
+ }
+
+ @Override
+ public int maxEncodedSize(CodecOptions options, int inputSize) {
+ return inputSize;
+ }
+
+ @Override
+ public boolean requiresDirectDstBuffer() {
+ return false;
+ }
+ }
+}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CodecRegistryTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CodecRegistryTest.java
new file mode 100644
index 00000000000..cb0564d1dba
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CodecRegistryTest.java
@@ -0,0 +1,141 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.io.IOException;
+import java.lang.reflect.Modifier;
+import java.nio.ByteBuffer;
+import java.util.List;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
+
+
+/// Tests for the immutable production registry and its package-scoped mutable
test fixture.
+public class CodecRegistryTest {
+
+ @Test
+ public void testDefaultRegistryContainsOnlyCompressionHandlersAndAlias() {
+ assertSame(CodecRegistry.DEFAULT.get("lz4"), Lz4CodecDefinition.INSTANCE);
+ assertSame(CodecRegistry.DEFAULT.get("SNAPPY"),
SnappyCodecDefinition.INSTANCE);
+ assertSame(CodecRegistry.DEFAULT.get("gzip"),
GzipCodecDefinition.INSTANCE);
+ assertSame(CodecRegistry.DEFAULT.get("zstd"),
ZstdCodecDefinition.INSTANCE);
+ assertSame(CodecRegistry.DEFAULT.get("zstandard"),
ZstdCodecDefinition.INSTANCE);
+ assertNull(CodecRegistry.DEFAULT.get("DELTA"));
+ assertNull(CodecRegistry.DEFAULT.get("T64"));
+ }
+
+ @Test
+ public void
testMutableRegistryEntryPointsRemainAvailableToSamePackageTests() {
+ CodecRegistry registry = new CodecRegistry();
+
+ assertSame(registry.register(Lz4CodecDefinition.INSTANCE), registry);
+ assertSame(registry.get("lz4"), Lz4CodecDefinition.INSTANCE);
+ }
+
+ @Test
+ public void testMutableRegistryEntryPointsAreNotPublic() throws
ReflectiveOperationException {
+ assertFalse(Modifier.isPublic(CodecRegistry.class.getModifiers()));
+
assertFalse(Modifier.isPublic(CodecRegistry.class.getDeclaredConstructor().getModifiers()));
+ assertFalse(Modifier.isPublic(
+ CodecRegistry.class.getDeclaredMethod("register",
ChunkCodecHandler.class).getModifiers()));
+ assertFalse(Modifier.isPublic(
+ CodecRegistry.class.getDeclaredMethod("get",
String.class).getModifiers()));
+ }
+
+ @Test
+ public void testDefaultRegistryRejectsMutation() {
+ assertThrows(UnsupportedOperationException.class,
+ () -> CodecRegistry.DEFAULT.register(Lz4CodecDefinition.INSTANCE));
+ }
+
+ @Test
+ public void testMutableRegistryRejectsDuplicateNameIgnoringCase() {
+ CodecRegistry registry = new
CodecRegistry().register(Lz4CodecDefinition.INSTANCE);
+ assertThrows(IllegalArgumentException.class, () ->
registry.register(Lz4CodecDefinition.INSTANCE));
+ }
+
+ @Test
+ public void testMutableRegistryRejectsReservedWrapperName() {
+ CodecRegistry registry = new CodecRegistry();
+ IllegalArgumentException exception =
+ expectThrows(IllegalArgumentException.class, () ->
registry.register(new ReservedNameStub()));
+ assertTrue(exception.getMessage().contains("reserved"),
exception.getMessage());
+ }
+
+ /// Stub whose name collides with the reserved wrapper keyword; every
behavior method is unreachable.
+ private static final class ReservedNameStub implements
ChunkCodecHandler<CodecOptions> {
+ @Override
+ public String name() {
+ return "codec";
+ }
+
+ @Override
+ public CodecKind kind() {
+ return CodecKind.COMPRESSION;
+ }
+
+ @Override
+ public CodecOptions parseOptions(List<String> args) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void validateContext(CodecOptions options, CodecContext ctx) {
+ }
+
+ @Override
+ public String canonicalize(CodecOptions options) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public ByteBuffer encode(CodecOptions options, CodecContext ctx,
ByteBuffer src)
+ throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public ByteBuffer decode(CodecOptions options, CodecContext ctx,
ByteBuffer src)
+ throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void decodeInto(CodecOptions options, CodecContext ctx, ByteBuffer
src, ByteBuffer dst)
+ throws IOException {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int maxEncodedSize(CodecOptions options, int inputSize) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public boolean requiresDirectDstBuffer() {
+ return false;
+ }
+ }
+}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CompressionCodecCorruptInputTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CompressionCodecCorruptInputTest.java
new file mode 100644
index 00000000000..99a0e2a0914
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/CompressionCodecCorruptInputTest.java
@@ -0,0 +1,236 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.Arrays;
+import java.util.List;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.testng.annotations.Test;
+import org.xerial.snappy.Snappy;
+
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
+
+
+/// Negative-path coverage for the compression codecs' corrupt-frame defenses
(length-prefix /
+/// footer / sanity-cap branches). Happy-path round-trips are covered by
[CodecPipelineExecutorTest].
+public class CompressionCodecCorruptInputTest {
+
+ private static final CodecContext CTX = new CodecContext(DataType.INT);
+
+ private static ByteBuffer garbage(int n) {
+ ByteBuffer b = ByteBuffer.allocateDirect(n);
+ for (int i = 0; i < n; i++) {
+ b.put((byte) 0xFF);
+ }
+ b.flip();
+ return b;
+ }
+
+ @Test
+ public void testLz4DecodeRejectsGarbage() {
+ // A 0xFF length-prefix decodes to a negative/huge length → out-of-range
guard must fire.
+ assertThrows(Exception.class,
+ () -> Lz4CodecDefinition.INSTANCE.decode(Lz4CodecDefinition.OPTIONS,
CTX, garbage(64)));
+ }
+
+ @Test
+ public void testZstdDecodeRejectsGarbage() {
+ assertThrows(Exception.class,
+ () ->
ZstdCodecDefinition.INSTANCE.decode(ZstdCodecDefinition.INSTANCE.parseOptions(List.of()),
CTX,
+ garbage(64)));
+ }
+
+ @Test
+ public void testZstdDecodeIntoRejectsOversizedDeclaredSize() {
+ // decodeInto must apply the shared sanity cap before its later
destination-capacity check.
+ ByteBuffer dst = ByteBuffer.allocateDirect(16);
+ IOException exception = expectThrows(IOException.class,
+ () -> ZstdCodecDefinition.INSTANCE.decodeInto(
+ ZstdCodecDefinition.INSTANCE.parseOptions(List.of()), CTX,
oversizedZstdFrame(), dst));
+ assertTrue(exception.getMessage().contains("out of range"),
exception.getMessage());
+ }
+
+ /// Frame header whose content-size field declares (1 << 30) + 1
decompressed bytes.
+ private static ByteBuffer oversizedZstdFrame() {
+ ByteBuffer frame = ByteBuffer.allocateDirect(9);
+ frame.put(new byte[]{
+ 0x28, (byte) 0xB5, 0x2F, (byte) 0xFD, (byte) 0xA0, 0x01, 0x00, 0x00,
0x40
+ });
+ frame.flip();
+ return frame;
+ }
+
+ @Test
+ public void testSnappyDecodeRejectsGarbage() {
+ assertThrows(Exception.class,
+ () ->
SnappyCodecDefinition.INSTANCE.decode(SnappyCodecDefinition.OPTIONS, CTX,
garbage(64)));
+ }
+
+ @Test
+ public void testSnappyDecodeRejectsOversizedDeclaredSize() {
+ // The sanity cap must fire before any allocation driven by the untrusted
header. The declared size is
+ // the smallest rejected value (cap + 1), pinning the inclusive boundary:
exactly 1 GiB is accepted.
+ IOException exception = expectThrows(IOException.class,
+ () ->
SnappyCodecDefinition.INSTANCE.decode(SnappyCodecDefinition.OPTIONS, CTX,
oversizedSnappyFrame()));
+ assertTrue(exception.getMessage().contains("out of range"),
exception.getMessage());
+ }
+
+ @Test
+ public void testSnappyDecodeIntoRejectsOversizedDeclaredSize() {
+ // decodeInto is the path the bounded executor uses in production; it must
reject with the size-cap
+ // IOException, not the later dst-capacity IllegalArgumentException.
+ ByteBuffer dst = ByteBuffer.allocateDirect(16);
+ IOException exception = expectThrows(IOException.class,
+ () ->
SnappyCodecDefinition.INSTANCE.decodeInto(SnappyCodecDefinition.OPTIONS, CTX,
+ oversizedSnappyFrame(), dst));
+ assertTrue(exception.getMessage().contains("out of range"),
exception.getMessage());
+ }
+
+ /// Frame whose varint length header declares (1 << 30) + 1 decompressed
bytes — one past the cap.
+ private static ByteBuffer oversizedSnappyFrame() {
+ ByteBuffer frame = ByteBuffer.allocateDirect(8);
+ frame.put(new byte[]{(byte) 0x81, (byte) 0x80, (byte) 0x80, (byte) 0x80,
0x04, 0x00, 0x00, 0x00});
+ frame.flip();
+ return frame;
+ }
+
+ @Test
+ public void testGzipDecodeRejectsGarbage() {
+ IOException exception = expectThrows(IOException.class,
+ () -> GzipCodecDefinition.INSTANCE.decode(GzipCodecDefinition.OPTIONS,
CTX, garbage(64)));
+ assertTrue(exception.getMessage().contains("invalid decompressed size"),
exception.getMessage());
+ }
+
+ @Test
+ public void testGzipDecodeRejectsTruncatedFooter() {
+ // Fewer than 4 bytes cannot carry the uncompressed-size footer GZIP
appends.
+ IOException exception = expectThrows(IOException.class,
+ () -> GzipCodecDefinition.INSTANCE.decode(GzipCodecDefinition.OPTIONS,
CTX, garbage(2)));
+ assertTrue(exception.getMessage().contains("too short"),
exception.getMessage());
+ }
+
+ @Test
+ public void testGzipDecodeRejectsUnderReportedSize() throws Exception {
+ ByteBuffer encoded = encodeGzip(new byte[256]);
+ encoded.putInt(encoded.limit() - Integer.BYTES, 128);
+ IOException exception = expectThrows(IOException.class,
+ () -> GzipCodecDefinition.INSTANCE.decode(GzipCodecDefinition.OPTIONS,
CTX, encoded));
+ assertTrue(exception.getMessage().contains("expands beyond footer-declared
size"), exception.getMessage());
+ }
+
+ @Test
+ public void testGzipDecodeRejectsZeroSizeForNonEmptyStream() throws
Exception {
+ ByteBuffer encoded = encodeGzip(new byte[256]);
+ encoded.putInt(encoded.limit() - Integer.BYTES, 0);
+ IOException exception = expectThrows(IOException.class,
+ () -> GzipCodecDefinition.INSTANCE.decode(GzipCodecDefinition.OPTIONS,
CTX, encoded));
+ assertTrue(exception.getMessage().contains("expands beyond footer-declared
size"), exception.getMessage());
+ }
+
+ @Test
+ public void testGzipDecodeRejectsCorruptChecksum() throws Exception {
+ ByteBuffer encoded = encodeGzip(new byte[256]);
+ int checksumByte = encoded.limit() - Integer.BYTES - 1;
+ encoded.put(checksumByte, (byte) (encoded.get(checksumByte) ^ 0x01));
+ IOException exception = expectThrows(IOException.class,
+ () -> GzipCodecDefinition.INSTANCE.decode(GzipCodecDefinition.OPTIONS,
CTX, encoded));
+ assertTrue(exception.getMessage().contains("GZIP decompression failed"),
exception.getMessage());
+ }
+
+ @Test
+ public void testGzipDecodeReadsBigEndianFooterFromLittleEndianView()
+ throws Exception {
+ byte[] values = new byte[256];
+ Arrays.fill(values, (byte) 0x5A);
+ ByteBuffer encoded = encodeGzip(values).order(ByteOrder.LITTLE_ENDIAN);
+
+ assertDecodedEquals(values,
+ GzipCodecDefinition.INSTANCE.decode(GzipCodecDefinition.OPTIONS, CTX,
encoded));
+ }
+
+ @Test
+ public void testGzipDecodeHonorsNonZeroSourcePosition()
+ throws Exception {
+ byte[] values = new byte[256];
+ Arrays.fill(values, (byte) 0x6B);
+ ByteBuffer encoded = encodeGzip(values);
+ int prefixLength = 7;
+ ByteBuffer framed = ByteBuffer.allocateDirect(prefixLength +
encoded.remaining());
+ framed.position(prefixLength);
+ framed.put(encoded.duplicate()).flip();
+ framed.position(prefixLength);
+
+ assertDecodedEquals(values,
+ GzipCodecDefinition.INSTANCE.decode(GzipCodecDefinition.OPTIONS, CTX,
framed));
+ }
+
+ @Test
+ public void testGzipMaxEncodedSizeRejectsOverflow() {
+ assertThrows(IllegalArgumentException.class,
+ () ->
GzipCodecDefinition.INSTANCE.maxEncodedSize(GzipCodecDefinition.OPTIONS,
Integer.MAX_VALUE));
+ }
+
+ @Test
+ public void testMultiStageDecodeBoundsInnerSnappyOutput() throws Exception {
+ // A valid Snappy frame that expands to 1 MiB is far larger than the
LZ4-encoded intermediate
+ // for the four-byte final chunk declared by the outer format. The
executor must use
+ // Snappy.decodeInto() with scratch derived from LZ4's bound instead of
allocating from the
+ // inner Snappy header.
+ byte[] compressed = Snappy.compress(new byte[1024 * 1024]);
+ // Keep only a small prefix: it contains Snappy's claimed decoded length
and stays below the
+ // executor's outer encoded-size bound, so this specifically exercises the
inner-frame guard.
+ ByteBuffer encoded = ByteBuffer.allocateDirect(16);
+ encoded.put(compressed, 0, encoded.capacity()).flip();
+ ByteBuffer decoded = ByteBuffer.allocateDirect(Integer.BYTES);
+ CodecPipelineExecutor executor = CodecPipelineExecutor.create(
+ "LZ4,SNAPPY", CTX, CodecRegistry.DEFAULT);
+
+ IllegalArgumentException exception =
expectThrows(IllegalArgumentException.class,
+ () -> executor.decode(encoded, decoded, Integer.BYTES));
+ assertTrue(exception.getMessage().contains("Snappy: decompressed size"),
exception.getMessage());
+ }
+
+ @Test
+ public void testBoundedExecutorRejectsOversizedOuterGzipFrameBeforeCopy() {
+ ByteBuffer encoded = garbage(64);
+ ByteBuffer decoded = ByteBuffer.allocateDirect(Integer.BYTES);
+ CodecPipelineExecutor executor = CodecPipelineExecutor.create("GZIP", CTX,
CodecRegistry.DEFAULT);
+
+ IllegalArgumentException exception =
expectThrows(IllegalArgumentException.class,
+ () -> executor.decode(encoded, decoded, Integer.BYTES, 128));
+ assertTrue(exception.getMessage().contains("Encoded input size"),
exception.getMessage());
+ }
+
+ private static ByteBuffer encodeGzip(byte[] values) throws Exception {
+ ByteBuffer source = ByteBuffer.allocateDirect(values.length);
+ source.put(values).flip();
+ return GzipCodecDefinition.INSTANCE.encode(GzipCodecDefinition.OPTIONS,
CTX, source);
+ }
+
+ private static void assertDecodedEquals(byte[] expected, ByteBuffer decoded)
{
+ byte[] actual = new byte[decoded.remaining()];
+ decoded.get(actual);
+ assertTrue(Arrays.equals(actual, expected));
+ }
+}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/ZstdCodecDefinitionTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/ZstdCodecDefinitionTest.java
new file mode 100644
index 00000000000..d478ce70417
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/codec/ZstdCodecDefinitionTest.java
@@ -0,0 +1,49 @@
+/**
+ * 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.pinot.segment.local.io.codec;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+public class ZstdCodecDefinitionTest {
+ private static final ZstdCodecDefinition CODEC =
ZstdCodecDefinition.INSTANCE;
+ private static final ZstdCodecDefinition.Options OPTIONS =
CODEC.parseOptions(List.of());
+ private static final CodecContext CONTEXT = new CodecContext(DataType.INT);
+
+ @Test
+ public void testEmptyInputRoundTrip() throws Exception {
+ ByteBuffer encoded = CODEC.encode(OPTIONS, CONTEXT,
ByteBuffer.allocateDirect(0));
+ assertTrue(encoded.hasRemaining(), "An empty input should still produce a
Zstd frame");
+
+ ByteBuffer decoded = CODEC.decode(OPTIONS, CONTEXT, encoded.duplicate());
+ assertEquals(decoded.remaining(), 0);
+
+ ByteBuffer destination = ByteBuffer.allocateDirect(0);
+ CODEC.decodeInto(OPTIONS, CONTEXT, encoded.duplicate(), destination);
+ assertEquals(destination.position(), 0);
+ assertEquals(destination.limit(), 0);
+ assertEquals(destination.remaining(), 0);
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
