This is an automated email from the ASF dual-hosted git repository.
wgtmac pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/parquet-java.git
The following commit(s) were added to refs/heads/master by this push:
new 8e4f57163 PARQUET-3459: Per column compression (#3526)
8e4f57163 is described below
commit 8e4f571630276201d6db97e28c635fcce7f44270
Author: mengna-lin <[email protected]>
AuthorDate: Mon Jul 13 19:34:50 2026 -0700
PARQUET-3459: Per column compression (#3526)
---
.../apache/parquet/column/ParquetProperties.java | 79 ++++-
.../parquet/column/TestParquetProperties.java | 141 ++++++++
.../compression/CompressionCodecFactory.java | 24 ++
.../org/apache/parquet/hadoop/CodecFactory.java | 150 +++++++-
.../parquet/hadoop/ColumnChunkPageWriteStore.java | 149 +++++++-
.../apache/parquet/hadoop/DirectCodecFactory.java | 30 +-
.../hadoop/InternalParquetRecordWriter.java | 26 +-
.../apache/parquet/hadoop/ParquetOutputFormat.java | 17 +
.../apache/parquet/hadoop/ParquetRecordWriter.java | 6 +-
.../org/apache/parquet/hadoop/ParquetWriter.java | 36 +-
.../hadoop/TestColumnChunkPageWriteStore.java | 133 +++++++
.../parquet/hadoop/TestDirectCodecFactory.java | 159 +++++++++
.../apache/parquet/hadoop/TestParquetWriter.java | 388 +++++++++++++++++++++
13 files changed, 1293 insertions(+), 45 deletions(-)
diff --git
a/parquet-column/src/main/java/org/apache/parquet/column/ParquetProperties.java
b/parquet-column/src/main/java/org/apache/parquet/column/ParquetProperties.java
index f29214b45..8fe45e01e 100644
---
a/parquet-column/src/main/java/org/apache/parquet/column/ParquetProperties.java
+++
b/parquet-column/src/main/java/org/apache/parquet/column/ParquetProperties.java
@@ -39,6 +39,7 @@ import
org.apache.parquet.column.values.factory.DefaultValuesWriterFactory;
import org.apache.parquet.column.values.factory.ValuesWriterFactory;
import org.apache.parquet.column.values.rle.RunLengthBitPackingHybridEncoder;
import
org.apache.parquet.column.values.rle.RunLengthBitPackingHybridValuesWriter;
+import org.apache.parquet.hadoop.metadata.CompressionCodecName;
import org.apache.parquet.schema.MessageType;
/**
@@ -135,6 +136,8 @@ public class ParquetProperties {
private final Map<String, String> extraMetaData;
private final ColumnProperty<Boolean> statistics;
private final ColumnProperty<Boolean> sizeStatistics;
+ private final ColumnProperty<CompressionCodecName> columnCodecs;
+ private final ColumnProperty<Integer> columnCompressionLevels;
private ParquetProperties(Builder builder) {
this.pageSizeThreshold = builder.pageSize;
@@ -167,6 +170,8 @@ public class ParquetProperties {
this.extraMetaData = builder.extraMetaData;
this.statistics = builder.statistics.build();
this.sizeStatistics = builder.sizeStatistics.build();
+ this.columnCodecs = builder.columnCodecs.build();
+ this.columnCompressionLevels = builder.columnCompressionLevels.build();
}
public static Builder builder() {
@@ -370,9 +375,31 @@ public class ParquetProperties {
return sizeStatisticsEnabled;
}
+ /**
+ * Returns the compression codec configured for the given column, or {@code
null} if no
+ * column-specific codec has been set (the caller should fall back to the
job-level codec).
+ *
+ * @param column the column descriptor
+ * @return the per-column codec, or {@code null} if not set
+ */
+ public CompressionCodecName getColumnCodec(ColumnDescriptor column) {
+ return columnCodecs.getValue(column);
+ }
+
+ /**
+ * Returns the compression level configured for the given column, or {@code
null} if no
+ * column-specific level has been set.
+ *
+ * @param column the column descriptor
+ * @return the per-column compression level, or {@code null} if not set
+ */
+ public Integer getColumnCompressionLevel(ColumnDescriptor column) {
+ return columnCompressionLevels.getValue(column);
+ }
+
@Override
public String toString() {
- return "Parquet page size to " + getPageSizeThreshold() + '\n'
+ String result = "Parquet page size to " + getPageSizeThreshold() + '\n'
+ "Parquet dictionary page size to " +
getDictionaryPageSizeThreshold() + '\n'
+ "Dictionary is " + dictionaryEnabled + '\n'
+ "Writer version is: " + getWriterVersion() + '\n'
@@ -389,6 +416,18 @@ public class ParquetProperties {
+ "Writing page checksums is: " + (getPageWriteChecksumEnabled() ?
"on" : "off") + '\n'
+ "Statistics enabled: " + statisticsEnabled + '\n'
+ "Size statistics enabled: " + sizeStatisticsEnabled;
+ String perColumn = "";
+ if
(!columnCodecs.toString().equals(Objects.toString(columnCodecs.getDefaultValue())))
{
+ perColumn = "Per-column codecs: " + columnCodecs;
+ }
+ if
(!columnCompressionLevels.toString().equals(Objects.toString(columnCompressionLevels.getDefaultValue())))
{
+ perColumn += perColumn.isEmpty() ? "Per-column compression levels: " :
", compression levels: ";
+ perColumn += columnCompressionLevels;
+ }
+ if (!perColumn.isEmpty()) {
+ result += '\n' + perColumn;
+ }
+ return result;
}
public static class Builder {
@@ -419,6 +458,8 @@ public class ParquetProperties {
private Map<String, String> extraMetaData = new HashMap<>();
private final ColumnProperty.Builder<Boolean> statistics;
private final ColumnProperty.Builder<Boolean> sizeStatistics;
+ private final ColumnProperty.Builder<CompressionCodecName> columnCodecs;
+ private final ColumnProperty.Builder<Integer> columnCompressionLevels;
private Builder() {
enableDict =
ColumnProperty.<Boolean>builder().withDefaultValue(DEFAULT_IS_DICTIONARY_ENABLED);
@@ -436,6 +477,8 @@ public class ParquetProperties {
ColumnProperty.<Integer>builder().withDefaultValue(DEFAULT_BLOOM_FILTER_CANDIDATES_NUMBER);
statistics =
ColumnProperty.<Boolean>builder().withDefaultValue(DEFAULT_STATISTICS_ENABLED);
sizeStatistics =
ColumnProperty.<Boolean>builder().withDefaultValue(DEFAULT_SIZE_STATISTICS_ENABLED);
+ columnCodecs =
ColumnProperty.<CompressionCodecName>builder().withDefaultValue(null);
+ columnCompressionLevels =
ColumnProperty.<Integer>builder().withDefaultValue(null);
}
private Builder(ParquetProperties toCopy) {
@@ -448,7 +491,11 @@ public class ParquetProperties {
this.estimateNextSizeCheck = toCopy.estimateNextSizeCheck;
this.valuesWriterFactory = toCopy.valuesWriterFactory;
this.allocator = toCopy.allocator;
+ this.pageValueCountThreshold = toCopy.pageValueCountThreshold;
this.pageRowCountLimit = toCopy.pageRowCountLimit;
+ this.rowGroupRowCountLimit = toCopy.rowGroupRowCountLimit;
+ this.columnIndexTruncateLength = toCopy.columnIndexTruncateLength;
+ this.statisticsTruncateLength = toCopy.statisticsTruncateLength;
this.pageWriteChecksumEnabled = toCopy.pageWriteChecksumEnabled;
this.bloomFilterNDVs = ColumnProperty.builder(toCopy.bloomFilterNDVs);
this.bloomFilterFPPs = ColumnProperty.builder(toCopy.bloomFilterFPPs);
@@ -459,7 +506,11 @@ public class ParquetProperties {
this.byteStreamSplitEnabled =
ColumnProperty.builder(toCopy.byteStreamSplitEnabled);
this.extraMetaData = toCopy.extraMetaData;
this.statistics = ColumnProperty.builder(toCopy.statistics);
+ this.statisticsEnabled = toCopy.statisticsEnabled;
this.sizeStatistics = ColumnProperty.builder(toCopy.sizeStatistics);
+ this.sizeStatisticsEnabled = toCopy.sizeStatisticsEnabled;
+ this.columnCodecs = ColumnProperty.builder(toCopy.columnCodecs);
+ this.columnCompressionLevels =
ColumnProperty.builder(toCopy.columnCompressionLevels);
}
/**
@@ -756,6 +807,32 @@ public class ParquetProperties {
return this;
}
+ /**
+ * Set the compression codec for the specified column.
+ *
+ * @param columnPath the path of the column (dot-string)
+ * @param codec the compression codec to use for this column
+ * @return this builder for method chaining
+ */
+ public Builder withCompressionCodec(String columnPath,
CompressionCodecName codec) {
+ this.columnCodecs.withValue(columnPath, Objects.requireNonNull(codec,
"codec cannot be null"));
+ return this;
+ }
+
+ /**
+ * Set the compression level for the specified column.
+ * The valid range is codec-specific (e.g. ZSTD: negative levels trade
ratio for speed, up to 22 for best compression (default 3), GZIP: 0–9 or -1
(default 6), BROTLI: 0–11 (default 1)).
+ * Pass {@code null} to unset a previously configured level for that
column.
+ *
+ * @param columnPath the path of the column (dot-string)
+ * @param level the compression level, or {@code null} to unset
+ * @return this builder for method chaining
+ */
+ public Builder withCompressionLevel(String columnPath, Integer level) {
+ this.columnCompressionLevels.withValue(columnPath, level);
+ return this;
+ }
+
public ParquetProperties build() {
ParquetProperties properties = new ParquetProperties(this);
// we pass a constructed but uninitialized factory to ParquetProperties
above as currently
diff --git
a/parquet-column/src/test/java/org/apache/parquet/column/TestParquetProperties.java
b/parquet-column/src/test/java/org/apache/parquet/column/TestParquetProperties.java
new file mode 100644
index 000000000..8456e42f4
--- /dev/null
+++
b/parquet-column/src/test/java/org/apache/parquet/column/TestParquetProperties.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.parquet.column;
+
+import static org.apache.parquet.hadoop.metadata.CompressionCodecName.GZIP;
+import static org.apache.parquet.hadoop.metadata.CompressionCodecName.SNAPPY;
+import static org.apache.parquet.hadoop.metadata.CompressionCodecName.ZSTD;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+
+import org.apache.parquet.schema.MessageType;
+import org.apache.parquet.schema.MessageTypeParser;
+import org.junit.Before;
+import org.junit.Test;
+
+public class TestParquetProperties {
+
+ private static final MessageType SCHEMA = MessageTypeParser.parseMessageType(
+ "message test { required binary col_a; required int32 col_b; required
double col_c; }");
+
+ private ColumnDescriptor colA;
+ private ColumnDescriptor colB;
+ private ColumnDescriptor colC;
+
+ @Before
+ public void setUp() {
+ colA = SCHEMA.getColumns().get(0);
+ colB = SCHEMA.getColumns().get(1);
+ colC = SCHEMA.getColumns().get(2);
+ }
+
+ @Test
+ public void columnCodec_notSet_returnsNull() {
+ ParquetProperties props = ParquetProperties.builder().build();
+ assertNull(props.getColumnCodec(colA));
+ }
+
+ @Test
+ public void columnLevel_notSet_returnsNull() {
+ ParquetProperties props = ParquetProperties.builder().build();
+ assertNull(props.getColumnCompressionLevel(colA));
+ }
+
+ @Test
+ public void columnCodec_setForColumn_returnsConfiguredCodec() {
+ ParquetProperties props =
+ ParquetProperties.builder().withCompressionCodec("col_a",
ZSTD).build();
+ assertEquals(ZSTD, props.getColumnCodec(colA));
+ }
+
+ @Test
+ public void columnCodec_setMultipleTimes_lastValueWins() {
+ ParquetProperties props = ParquetProperties.builder()
+ .withCompressionCodec("col_a", ZSTD)
+ .withCompressionCodec("col_a", SNAPPY)
+ .build();
+ assertEquals(SNAPPY, props.getColumnCodec(colA));
+ }
+
+ @Test
+ public void columnCodec_otherColumnsUnaffected() {
+ ParquetProperties props =
+ ParquetProperties.builder().withCompressionCodec("col_a",
ZSTD).build();
+ assertNull(props.getColumnCodec(colB));
+ assertNull(props.getColumnCodec(colC));
+ }
+
+ @Test
+ public void columnLevel_setForColumn_returnsConfiguredLevel() {
+ ParquetProperties props =
+ ParquetProperties.builder().withCompressionLevel("col_a", 10).build();
+ assertEquals(Integer.valueOf(10), props.getColumnCompressionLevel(colA));
+ }
+
+ @Test
+ public void columnLevel_otherColumnsUnaffected() {
+ ParquetProperties props =
+ ParquetProperties.builder().withCompressionLevel("col_a", 10).build();
+ assertNull(props.getColumnCompressionLevel(colB));
+ }
+
+ @Test
+ public void columnCodecAndLevel_multipleColumns_eachGetsOwn() {
+ ParquetProperties props = ParquetProperties.builder()
+ .withCompressionCodec("col_a", ZSTD)
+ .withCompressionLevel("col_a", 10)
+ .withCompressionCodec("col_b", SNAPPY)
+ .withCompressionCodec("col_c", GZIP)
+ .withCompressionLevel("col_c", 5)
+ .build();
+
+ assertEquals(ZSTD, props.getColumnCodec(colA));
+ assertEquals(Integer.valueOf(10), props.getColumnCompressionLevel(colA));
+
+ assertEquals(SNAPPY, props.getColumnCodec(colB));
+ assertNull(props.getColumnCompressionLevel(colB));
+
+ assertEquals(GZIP, props.getColumnCodec(colC));
+ assertEquals(Integer.valueOf(5), props.getColumnCompressionLevel(colC));
+ }
+
+ @Test
+ public void withCompressionCodec_nullCodec_throwsNullPointerException() {
+ assertThrows(
+ NullPointerException.class, () ->
ParquetProperties.builder().withCompressionCodec("col_a", null));
+ }
+
+ @Test
+ public void copyBuilder_preservesColumnCodecAndLevel() {
+ ParquetProperties original = ParquetProperties.builder()
+ .withCompressionCodec("col_a", ZSTD)
+ .withCompressionLevel("col_a", 7)
+ .withCompressionCodec("col_b", SNAPPY)
+ .build();
+
+ ParquetProperties copy = ParquetProperties.copy(original).build();
+
+ assertEquals(ZSTD, copy.getColumnCodec(colA));
+ assertEquals(Integer.valueOf(7), copy.getColumnCompressionLevel(colA));
+ assertEquals(SNAPPY, copy.getColumnCodec(colB));
+ assertNull(copy.getColumnCompressionLevel(colB));
+ assertNull(copy.getColumnCodec(colC));
+ }
+}
diff --git
a/parquet-common/src/main/java/org/apache/parquet/compression/CompressionCodecFactory.java
b/parquet-common/src/main/java/org/apache/parquet/compression/CompressionCodecFactory.java
index 561dcb899..bc5a26ab8 100644
---
a/parquet-common/src/main/java/org/apache/parquet/compression/CompressionCodecFactory.java
+++
b/parquet-common/src/main/java/org/apache/parquet/compression/CompressionCodecFactory.java
@@ -46,6 +46,30 @@ public interface CompressionCodecFactory {
*/
BytesInputCompressor getCompressor(CompressionCodecName codecName);
+ /**
+ * Returns a {@link BytesInputCompressor} instance for the specified codec
name and compression level.
+ * <p>
+ * The compression level controls the trade-off between compression speed
and ratio. The valid range
+ * and meaning of the level is codec-specific:
+ * <ul>
+ * <li>ZSTD: negative levels trade compression ratio for higher speed, up
to 22 for best compression, default 3; the exact bounds are defined by the zstd
library</li>
+ * <li>GZIP: 1 (fastest) to 9 (best compression), default 6</li>
+ * <li>BROTLI: 0 (fastest) to 11 (best compression), default 1</li>
+ * </ul>
+ * Implementations that do not support compression levels should ignore the
{@code level} parameter
+ * and delegate to {@link #getCompressor(CompressionCodecName)}.
+ * <p>
+ * The compressor is not thread-safe, so one instance for each working
thread is required.
+ *
+ * @param codecName the codec name which the compressor instance is to be
returned
+ * @param level the compression level; codec-specific, ignored if the
codec does not support levels
+ * @return the compressor instance for the specified codec name and level
+ * @see BytesInputCompressor#release()
+ */
+ default BytesInputCompressor getCompressor(CompressionCodecName codecName,
int level) {
+ return getCompressor(codecName);
+ }
+
/**
* Returns a {@link BytesInputDecompressor} instance for the specified codec
name to be used for decompressing page
* data.
diff --git
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/CodecFactory.java
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/CodecFactory.java
index 98b49835a..c9391201f 100644
--- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/CodecFactory.java
+++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/CodecFactory.java
@@ -18,6 +18,7 @@
*/
package org.apache.parquet.hadoop;
+import com.github.luben.zstd.Zstd;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
@@ -32,6 +33,7 @@ import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.compress.CompressionOutputStream;
import org.apache.hadoop.io.compress.Compressor;
import org.apache.hadoop.io.compress.Decompressor;
+import org.apache.hadoop.io.compress.zlib.ZlibCompressor;
import org.apache.hadoop.util.ReflectionUtils;
import org.apache.parquet.Preconditions;
import org.apache.parquet.bytes.ByteBufferAllocator;
@@ -43,13 +45,20 @@ import org.apache.parquet.hadoop.codec.Lz4RawCodec;
import org.apache.parquet.hadoop.codec.ZstandardCodec;
import org.apache.parquet.hadoop.metadata.CompressionCodecName;
import org.apache.parquet.hadoop.util.ConfigurationUtil;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
public class CodecFactory implements CompressionCodecFactory {
+ private static final Logger LOG =
LoggerFactory.getLogger(CodecFactory.class);
+
protected static final Map<String, CompressionCodec> CODEC_BY_NAME =
Collections.synchronizedMap(new HashMap<String, CompressionCodec>());
- private final Map<CompressionCodecName, BytesCompressor> compressors = new
HashMap<>();
+ static final String GZIP_COMPRESS_LEVEL = "zlib.compress.level";
+ static final String BROTLI_COMPRESS_QUALITY = "compression.brotli.quality";
+
+ private final Map<String, BytesCompressor> compressors = new HashMap<>();
private final Map<CompressionCodecName, BytesDecompressor> decompressors =
new HashMap<>();
protected final ParquetConfiguration conf;
@@ -252,10 +261,22 @@ public class CodecFactory implements
CompressionCodecFactory {
@Override
public BytesCompressor getCompressor(CompressionCodecName codecName) {
- BytesCompressor comp = compressors.get(codecName);
+ String key = cacheKey(codecName);
+ BytesCompressor comp = compressors.get(key);
if (comp == null) {
comp = createCompressor(codecName);
- compressors.put(codecName, comp);
+ compressors.put(key, comp);
+ }
+ return comp;
+ }
+
+ @Override
+ public BytesCompressor getCompressor(CompressionCodecName codecName, int
level) {
+ String key = cacheKey(codecName, level);
+ BytesCompressor comp = compressors.get(key);
+ if (comp == null) {
+ comp = createCompressorAtLevel(codecName, level);
+ compressors.put(key, comp);
}
return comp;
}
@@ -271,8 +292,7 @@ public class CodecFactory implements
CompressionCodecFactory {
}
protected BytesCompressor createCompressor(CompressionCodecName codecName) {
- CompressionCodec codec = getCodec(codecName);
- return codec == null ? NO_OP_COMPRESSOR : new
HeapBytesCompressor(codecName, codec);
+ return compressorForCodec(codecName, getCodec(codecName));
}
protected BytesDecompressor createDecompressor(CompressionCodecName
codecName) {
@@ -280,6 +300,117 @@ public class CodecFactory implements
CompressionCodecFactory {
return codec == null ? NO_OP_DECOMPRESSOR : new
HeapBytesDecompressor(codec);
}
+ protected BytesCompressor createCompressorAtLevel(CompressionCodecName
codecName, int level) {
+ return compressorForCodec(codecName, getCodecAtLevel(codecName, level));
+ }
+
+ private BytesCompressor compressorForCodec(CompressionCodecName codecName,
CompressionCodec codec) {
+ return codec == null ? NO_OP_COMPRESSOR : new
HeapBytesCompressor(codecName, codec);
+ }
+
+ static void validateZstdLevel(int level) {
+ if (level < Zstd.minCompressionLevel() || level >
Zstd.maxCompressionLevel()) {
+ throw new BadConfigurationException(
+ "Unsupported ZSTD compression level: " + level + ". Valid range is "
+ Zstd.minCompressionLevel()
+ + " (fastest) to " + Zstd.maxCompressionLevel() + " (best
compression).");
+ }
+ }
+
+ private static void validateBrotliLevel(int level) {
+ if (level < 0 || level > 11) {
+ throw new BadConfigurationException("Unsupported Brotli compression
level: " + level
+ + ". Valid range is 0 (fastest) to 11 (best compression).");
+ }
+ }
+
+ private static void validateGzipLevel(int level) {
+ if (level != -1 && (level < 0 || level > 9)) {
+ throw new BadConfigurationException("Unsupported GZIP compression level:
" + level
+ + ". Valid range is 0 (no compression) to 9 (best compression), or
-1 for default.");
+ }
+ }
+
+ private static ZlibCompressor.CompressionLevel zlibCompressionLevel(int
level) {
+ switch (level) {
+ case -1:
+ return ZlibCompressor.CompressionLevel.DEFAULT_COMPRESSION;
+ case 0:
+ return ZlibCompressor.CompressionLevel.NO_COMPRESSION;
+ case 1:
+ return ZlibCompressor.CompressionLevel.BEST_SPEED;
+ case 2:
+ return ZlibCompressor.CompressionLevel.TWO;
+ case 3:
+ return ZlibCompressor.CompressionLevel.THREE;
+ case 4:
+ return ZlibCompressor.CompressionLevel.FOUR;
+ case 5:
+ return ZlibCompressor.CompressionLevel.FIVE;
+ case 6:
+ return ZlibCompressor.CompressionLevel.SIX;
+ case 7:
+ return ZlibCompressor.CompressionLevel.SEVEN;
+ case 8:
+ return ZlibCompressor.CompressionLevel.EIGHT;
+ case 9:
+ return ZlibCompressor.CompressionLevel.BEST_COMPRESSION;
+ default:
+ throw new BadConfigurationException("Unsupported GZIP compression
level: " + level
+ + ". Valid range is 0 (no compression) to 9 (best compression), or
-1 for default.");
+ }
+ }
+
+ /**
+ * Returns a {@link CompressionCodec} instance configured at the specified
compression level.
+ * A level-specific {@link Configuration} snapshot is built so that the
codec is initialized
+ * with the requested level rather than the global configuration value.
+ * For codecs that do not support levels the method falls back to {@link
#getCodec(CompressionCodecName)}.
+ */
+ private CompressionCodec getCodecAtLevel(CompressionCodecName codecName, int
level) {
+ String codecClassName = codecName.getHadoopCompressionCodecClassName();
+ if (codecClassName == null) {
+ return null;
+ }
+ String key = cacheKey(codecName, level);
+ CompressionCodec codec = CODEC_BY_NAME.get(key);
+ if (codec != null) {
+ return codec;
+ }
+ // Build a Configuration snapshot with the level explicitly set for this
codec.
+ Configuration levelConf = new
Configuration(ConfigurationUtil.createHadoopConfiguration(conf));
+ switch (codecName) {
+ case ZSTD:
+ validateZstdLevel(level);
+ levelConf.setInt(ZstandardCodec.PARQUET_COMPRESS_ZSTD_LEVEL, level);
+ break;
+ case GZIP:
+ validateGzipLevel(level);
+ levelConf.setEnum(GZIP_COMPRESS_LEVEL, zlibCompressionLevel(level));
+ break;
+ case BROTLI:
+ validateBrotliLevel(level);
+ levelConf.setInt(BROTLI_COMPRESS_QUALITY, level);
+ break;
+ default:
+ // Codec does not support levels; fall back to the default codec
instance.
+ LOG.warn("Compression level {} is not supported for codec {} and will
be ignored.", level, codecName);
+ return getCodec(codecName);
+ }
+ try {
+ Class<?> codecClass;
+ try {
+ codecClass = Class.forName(codecClassName);
+ } catch (ClassNotFoundException e) {
+ codecClass = new
Configuration(false).getClassLoader().loadClass(codecClassName);
+ }
+ codec = (CompressionCodec) ReflectionUtils.newInstance(codecClass,
levelConf);
+ CODEC_BY_NAME.put(key, codec);
+ return codec;
+ } catch (ClassNotFoundException e) {
+ throw new BadConfigurationException("Class " + codecClassName + " was
not found", e);
+ }
+ }
+
/**
* @param codecName the requested codec
* @return the corresponding hadoop codec. null if UNCOMPRESSED
@@ -316,10 +447,10 @@ public class CodecFactory implements
CompressionCodecFactory {
String level = null;
switch (codecName) {
case GZIP:
- level = conf.get("zlib.compress.level");
+ level = conf.get(GZIP_COMPRESS_LEVEL);
break;
case BROTLI:
- level = conf.get("compression.brotli.quality");
+ level = conf.get(BROTLI_COMPRESS_QUALITY);
break;
case ZSTD:
level = conf.get("parquet.compression.codec.zstd.level");
@@ -331,6 +462,11 @@ public class CodecFactory implements
CompressionCodecFactory {
return level == null ? codecClass : codecClass + ":" + level;
}
+ private String cacheKey(CompressionCodecName codecName, int level) {
+ String codecClass = codecName.getHadoopCompressionCodecClassName();
+ return (codecClass == null ? codecName.name() : codecClass) + ":" + level;
+ }
+
@Override
public void release() {
for (BytesCompressor compressor : compressors.values()) {
diff --git
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ColumnChunkPageWriteStore.java
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ColumnChunkPageWriteStore.java
index c0cf216cc..69c28b2cd 100644
---
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ColumnChunkPageWriteStore.java
+++
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ColumnChunkPageWriteStore.java
@@ -25,8 +25,10 @@ import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Objects;
import java.util.Optional;
import java.util.Set;
+import java.util.function.Function;
import java.util.zip.CRC32;
import org.apache.parquet.bytes.ByteBufferAllocator;
import org.apache.parquet.bytes.ByteBufferReleaser;
@@ -44,6 +46,7 @@ import
org.apache.parquet.column.statistics.geospatial.GeospatialStatistics;
import org.apache.parquet.column.values.bloomfilter.BloomFilter;
import org.apache.parquet.column.values.bloomfilter.BloomFilterWriteStore;
import org.apache.parquet.column.values.bloomfilter.BloomFilterWriter;
+import org.apache.parquet.compression.CompressionCodecFactory;
import
org.apache.parquet.compression.CompressionCodecFactory.BytesInputCompressor;
import org.apache.parquet.crypto.AesCipher;
import org.apache.parquet.crypto.InternalColumnEncryptionSetup;
@@ -53,6 +56,7 @@ import org.apache.parquet.format.BlockCipher;
import org.apache.parquet.format.converter.ParquetMetadataConverter;
import org.apache.parquet.hadoop.CodecFactory.BytesCompressor;
import org.apache.parquet.hadoop.metadata.ColumnPath;
+import org.apache.parquet.hadoop.metadata.CompressionCodecName;
import org.apache.parquet.internal.column.columnindex.ColumnIndexBuilder;
import org.apache.parquet.internal.column.columnindex.OffsetIndexBuilder;
import org.apache.parquet.io.ParquetEncodingException;
@@ -583,21 +587,7 @@ public class ColumnChunkPageWriteStore implements
PageWriteStore, BloomFilterWri
int columnIndexTruncateLength,
boolean pageWriteChecksumEnabled) {
this.schema = schema;
- for (ColumnDescriptor path : schema.getColumns()) {
- writers.put(
- path,
- new ColumnChunkPageWriter(
- path,
- compressor,
- allocator,
- columnIndexTruncateLength,
- pageWriteChecksumEnabled,
- null,
- null,
- null,
- -1,
- -1));
- }
+ initWriters(col -> compressor, allocator, columnIndexTruncateLength,
pageWriteChecksumEnabled, null, -1);
}
@Deprecated
@@ -628,13 +618,138 @@ public class ColumnChunkPageWriteStore implements
PageWriteStore, BloomFilterWri
InternalFileEncryptor fileEncryptor,
int rowGroupOrdinal) {
this.schema = schema;
+ initWriters(
+ col -> compressor,
+ allocator,
+ columnIndexTruncateLength,
+ pageWriteChecksumEnabled,
+ fileEncryptor,
+ rowGroupOrdinal);
+ }
+
+ private ColumnChunkPageWriteStore(
+ Function<ColumnDescriptor, BytesInputCompressor> compressorProvider,
+ MessageType schema,
+ ByteBufferAllocator allocator,
+ int columnIndexTruncateLength,
+ boolean pageWriteChecksumEnabled,
+ InternalFileEncryptor fileEncryptor,
+ int rowGroupOrdinal) {
+ this.schema = schema;
+ initWriters(
+ compressorProvider,
+ allocator,
+ columnIndexTruncateLength,
+ pageWriteChecksumEnabled,
+ fileEncryptor,
+ rowGroupOrdinal);
+ }
+
+ static Function<ColumnDescriptor, BytesInputCompressor> compressorProvider(
+ CompressionCodecFactory codecFactory,
+ CompressionCodecName defaultCodec,
+ ParquetProperties parquetProperties) {
+ return column -> resolveCompressor(column, codecFactory, defaultCodec,
parquetProperties);
+ }
+
+ private static BytesInputCompressor resolveCompressor(
+ ColumnDescriptor column,
+ CompressionCodecFactory codecFactory,
+ CompressionCodecName defaultCodec,
+ ParquetProperties parquetProperties) {
+ CompressionCodecName columnCodec =
parquetProperties.getColumnCodec(column);
+ CompressionCodecName codec = columnCodec != null ? columnCodec :
defaultCodec;
+ Integer level = parquetProperties.getColumnCompressionLevel(column);
+ if (level != null && columnCodec == null) {
+ LOG.warn(
+ "Column '{}': compression level {} set without a per-column codec; "
+ + "applying level to the default codec ({}).",
+ ColumnPath.get(column.getPath()),
+ level,
+ defaultCodec);
+ }
+ return level != null ? codecFactory.getCompressor(codec, level) :
codecFactory.getCompressor(codec);
+ }
+
+ public static Builder builder() {
+ return new Builder();
+ }
+
+ public static class Builder {
+ private Function<ColumnDescriptor, BytesInputCompressor>
compressorProvider;
+ private MessageType schema;
+ private ByteBufferAllocator allocator;
+ private int columnIndexTruncateLength =
ParquetProperties.DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH;
+ private boolean pageWriteChecksumEnabled =
ParquetProperties.DEFAULT_PAGE_WRITE_CHECKSUM_ENABLED;
+ private InternalFileEncryptor fileEncryptor = null;
+ private int rowGroupOrdinal = -1;
+
+ private Builder() {}
+
+ public Builder withCompressorProvider(Function<ColumnDescriptor,
BytesInputCompressor> compressorProvider) {
+ this.compressorProvider = compressorProvider;
+ return this;
+ }
+
+ public Builder withSchema(MessageType schema) {
+ this.schema = schema;
+ return this;
+ }
+
+ public Builder withAllocator(ByteBufferAllocator allocator) {
+ this.allocator = allocator;
+ return this;
+ }
+
+ public Builder withColumnIndexTruncateLength(int
columnIndexTruncateLength) {
+ this.columnIndexTruncateLength = columnIndexTruncateLength;
+ return this;
+ }
+
+ public Builder withPageWriteChecksumEnabled(boolean
pageWriteChecksumEnabled) {
+ this.pageWriteChecksumEnabled = pageWriteChecksumEnabled;
+ return this;
+ }
+
+ public Builder withFileEncryptor(InternalFileEncryptor fileEncryptor) {
+ this.fileEncryptor = fileEncryptor;
+ return this;
+ }
+
+ public Builder withRowGroupOrdinal(int rowGroupOrdinal) {
+ this.rowGroupOrdinal = rowGroupOrdinal;
+ return this;
+ }
+
+ public ColumnChunkPageWriteStore build() {
+ Objects.requireNonNull(compressorProvider, "compressorProvider cannot be
null");
+ Objects.requireNonNull(schema, "schema cannot be null");
+ Objects.requireNonNull(allocator, "allocator cannot be null");
+ return new ColumnChunkPageWriteStore(
+ compressorProvider,
+ schema,
+ allocator,
+ columnIndexTruncateLength,
+ pageWriteChecksumEnabled,
+ fileEncryptor,
+ rowGroupOrdinal);
+ }
+ }
+
+ private void initWriters(
+ Function<ColumnDescriptor, BytesInputCompressor> compressorFn,
+ ByteBufferAllocator allocator,
+ int columnIndexTruncateLength,
+ boolean pageWriteChecksumEnabled,
+ InternalFileEncryptor fileEncryptor,
+ int rowGroupOrdinal) {
if (null == fileEncryptor) {
for (ColumnDescriptor path : schema.getColumns()) {
writers.put(
path,
new ColumnChunkPageWriter(
path,
- compressor,
+ compressorFn.apply(path),
allocator,
columnIndexTruncateLength,
pageWriteChecksumEnabled,
@@ -666,7 +781,7 @@ public class ColumnChunkPageWriteStore implements
PageWriteStore, BloomFilterWri
path,
new ColumnChunkPageWriter(
path,
- compressor,
+ compressorFn.apply(path),
allocator,
columnIndexTruncateLength,
pageWriteChecksumEnabled,
diff --git
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/DirectCodecFactory.java
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/DirectCodecFactory.java
index b2b5233ee..312ccedf8 100644
---
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/DirectCodecFactory.java
+++
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/DirectCodecFactory.java
@@ -110,6 +110,20 @@ class DirectCodecFactory extends CodecFactory implements
AutoCloseable {
}
}
+ @Override
+ protected BytesCompressor createCompressorAtLevel(final CompressionCodecName
codecName, final int level) {
+ switch (codecName) {
+ case SNAPPY:
+ // Snappy has no compression level; keep using the direct compressor.
+ LOG.warn("Compression level {} is not supported for codec {} and will
be ignored.", level, codecName);
+ return new SnappyCompressor();
+ case ZSTD:
+ return new ZstdCompressor(level);
+ default:
+ return super.createCompressorAtLevel(codecName, level);
+ }
+ }
+
@Override
protected BytesDecompressor createDecompressor(final CompressionCodecName
codecName) {
switch (codecName) {
@@ -409,11 +423,21 @@ class DirectCodecFactory extends CodecFactory implements
AutoCloseable {
private final ZstdCompressCtx context;
ZstdCompressor() {
- context = new ZstdCompressCtx();
- context.setLevel(conf.getInt(
+ context = newContext(conf.getInt(
ZstandardCodec.PARQUET_COMPRESS_ZSTD_LEVEL,
ZstandardCodec.DEFAULT_PARQUET_COMPRESS_ZSTD_LEVEL));
- context.setWorkers(conf.getInt(
+ }
+
+ ZstdCompressor(int level) {
+ validateZstdLevel(level);
+ context = newContext(level);
+ }
+
+ private ZstdCompressCtx newContext(int level) {
+ ZstdCompressCtx ctx = new ZstdCompressCtx();
+ ctx.setLevel(level);
+ ctx.setWorkers(conf.getInt(
ZstandardCodec.PARQUET_COMPRESS_ZSTD_WORKERS,
ZstandardCodec.DEFAULTPARQUET_COMPRESS_ZSTD_WORKERS));
+ return ctx;
}
@Override
diff --git
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/InternalParquetRecordWriter.java
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/InternalParquetRecordWriter.java
index dd51d1ef0..d503cbb4e 100644
---
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/InternalParquetRecordWriter.java
+++
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/InternalParquetRecordWriter.java
@@ -25,6 +25,8 @@ import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
+import java.util.function.Function;
+import org.apache.parquet.column.ColumnDescriptor;
import org.apache.parquet.column.ColumnWriteStore;
import org.apache.parquet.column.ParquetProperties;
import org.apache.parquet.column.values.bloomfilter.BloomFilterWriteStore;
@@ -51,7 +53,7 @@ class InternalParquetRecordWriter<T> {
private long rowGroupSizeThreshold;
private final int rowGroupRecordCountThreshold;
private long nextRowGroupSize;
- private final BytesInputCompressor compressor;
+ private final Function<ColumnDescriptor, BytesInputCompressor>
compressorProvider;
private final boolean validating;
private final ParquetProperties props;
@@ -76,7 +78,6 @@ class InternalParquetRecordWriter<T> {
* @param schema the schema of the records
* @param extraMetaData extra meta data to write in the footer of the
file
* @param rowGroupSize the size of a block in the file (this will be
approximate)
- * @param compressor the codec used to compress
*/
public InternalParquetRecordWriter(
ParquetFileWriter parquetFileWriter,
@@ -84,7 +85,7 @@ class InternalParquetRecordWriter<T> {
MessageType schema,
Map<String, String> extraMetaData,
long rowGroupSize,
- BytesInputCompressor compressor,
+ Function<ColumnDescriptor, BytesInputCompressor> compressorProvider,
boolean validating,
ParquetProperties props) {
this.parquetFileWriter = parquetFileWriter;
@@ -94,7 +95,7 @@ class InternalParquetRecordWriter<T> {
this.rowGroupSizeThreshold = rowGroupSize;
this.rowGroupRecordCountThreshold = props.getRowGroupRowCountLimit();
this.nextRowGroupSize = rowGroupSizeThreshold;
- this.compressor = compressor;
+ this.compressorProvider = compressorProvider;
this.validating = validating;
this.props = props;
this.fileEncryptor = parquetFileWriter.getEncryptor();
@@ -108,14 +109,15 @@ class InternalParquetRecordWriter<T> {
}
private void initStore() {
- ColumnChunkPageWriteStore columnChunkPageWriteStore = new
ColumnChunkPageWriteStore(
- compressor,
- schema,
- props.getAllocator(),
- props.getColumnIndexTruncateLength(),
- props.getPageWriteChecksumEnabled(),
- fileEncryptor,
- rowGroupOrdinal);
+ ColumnChunkPageWriteStore columnChunkPageWriteStore =
ColumnChunkPageWriteStore.builder()
+ .withCompressorProvider(compressorProvider)
+ .withSchema(schema)
+ .withAllocator(props.getAllocator())
+ .withColumnIndexTruncateLength(props.getColumnIndexTruncateLength())
+ .withPageWriteChecksumEnabled(props.getPageWriteChecksumEnabled())
+ .withFileEncryptor(fileEncryptor)
+ .withRowGroupOrdinal(rowGroupOrdinal)
+ .build();
pageStore = columnChunkPageWriteStore;
bloomFilterWriteStore = columnChunkPageWriteStore;
diff --git
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetOutputFormat.java
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetOutputFormat.java
index 868ae634c..4db288f45 100644
---
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetOutputFormat.java
+++
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetOutputFormat.java
@@ -163,6 +163,7 @@ public class ParquetOutputFormat<T> extends
FileOutputFormat<Void, T> {
public static final String PAGE_WRITE_CHECKSUM_ENABLED =
"parquet.page.write-checksum.enabled";
public static final String STATISTICS_ENABLED =
"parquet.column.statistics.enabled";
public static final String SIZE_STATISTICS_ENABLED =
"parquet.size.statistics.enabled";
+ public static final String COLUMN_COMPRESSION_LEVEL_PREFIX =
"parquet.compression.level";
public static JobSummaryLevel getJobSummaryLevel(Configuration conf) {
String level = conf.get(JOB_SUMMARY_LEVEL);
@@ -222,6 +223,14 @@ public class ParquetOutputFormat<T> extends
FileOutputFormat<Void, T> {
getConfiguration(job).set(COMPRESSION, compression.name());
}
+ public static void setColumnCompression(Job job, String columnPath,
CompressionCodecName codec) {
+ getConfiguration(job).set(COMPRESSION + '#' + columnPath, codec.name());
+ }
+
+ public static void setColumnCompressionLevel(Job job, String columnPath, int
level) {
+ getConfiguration(job).setInt(COLUMN_COMPRESSION_LEVEL_PREFIX + '#' +
columnPath, level);
+ }
+
public static void setEnableDictionary(Job job, boolean enableDictionary) {
getConfiguration(job).setBoolean(ENABLE_DICTIONARY, enableDictionary);
}
@@ -546,6 +555,14 @@ public class ParquetOutputFormat<T> extends
FileOutputFormat<Void, T> {
STATISTICS_ENABLED,
key -> conf.getBoolean(key,
ParquetProperties.DEFAULT_STATISTICS_ENABLED),
propsBuilder::withStatisticsEnabled)
+ .withColumnConfig(
+ COMPRESSION,
+ key -> CompressionCodecName.fromConf(conf.get(key)),
+ propsBuilder::withCompressionCodec)
+ .withColumnConfig(
+ COLUMN_COMPRESSION_LEVEL_PREFIX,
+ key -> conf.getInt(key, -1),
+ propsBuilder::withCompressionLevel)
.parseConfig(conf);
ParquetProperties props = propsBuilder.build();
diff --git
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetRecordWriter.java
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetRecordWriter.java
index 51528b10b..80a6c69d1 100644
---
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetRecordWriter.java
+++
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetRecordWriter.java
@@ -104,7 +104,7 @@ public class ParquetRecordWriter<T> extends
RecordWriter<Void, T> {
.withWriterVersion(writerVersion)
.build();
internalWriter = new InternalParquetRecordWriter<T>(
- w, writeSupport, schema, extraMetaData, blockSize, compressor,
validating, props);
+ w, writeSupport, schema, extraMetaData, blockSize, col -> compressor,
validating, props);
this.memoryManager = null;
this.codecFactory = null;
}
@@ -173,7 +173,7 @@ public class ParquetRecordWriter<T> extends
RecordWriter<Void, T> {
.withWriterVersion(writerVersion)
.build();
internalWriter = new InternalParquetRecordWriter<T>(
- w, writeSupport, schema, extraMetaData, blockSize, compressor,
validating, props);
+ w, writeSupport, schema, extraMetaData, blockSize, col -> compressor,
validating, props);
this.memoryManager = Objects.requireNonNull(memoryManager, "memoryManager
cannot be null");
memoryManager.addWriter(internalWriter, blockSize);
this.codecFactory = null;
@@ -207,7 +207,7 @@ public class ParquetRecordWriter<T> extends
RecordWriter<Void, T> {
schema,
extraMetaData,
blockSize,
- codecFactory.getCompressor(codec),
+ ColumnChunkPageWriteStore.compressorProvider(codecFactory, codec,
props),
validating,
props);
this.memoryManager = Objects.requireNonNull(memoryManager, "memoryManager
cannot be null");
diff --git
a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetWriter.java
b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetWriter.java
index 8eb5f7f17..7dcbb3188 100644
--- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetWriter.java
+++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetWriter.java
@@ -395,7 +395,6 @@ public class ParquetWriter<T> implements Closeable {
fileWriter.start();
this.codecFactory = codecFactory;
- CompressionCodecFactory.BytesInputCompressor compressor =
codecFactory.getCompressor(compressionCodecName);
final Map<String, String> extraMetadata;
if (encodingProps.getExtraMetaData() == null
@@ -418,7 +417,14 @@ public class ParquetWriter<T> implements Closeable {
}
this.writer = new InternalParquetRecordWriter<T>(
- fileWriter, writeSupport, schema, extraMetadata, rowGroupSize,
compressor, validating, encodingProps);
+ fileWriter,
+ writeSupport,
+ schema,
+ extraMetadata,
+ rowGroupSize,
+ ColumnChunkPageWriteStore.compressorProvider(codecFactory,
compressionCodecName, encodingProps),
+ validating,
+ encodingProps);
}
public void write(T object) throws IOException {
@@ -569,6 +575,32 @@ public class ParquetWriter<T> implements Closeable {
return self();
}
+ /**
+ * Override the compression codec for a specific column.
+ *
+ * @param columnPath dot-string path of the column (e.g. {@code "my_col"})
+ * @param codecName the codec to use for that column
+ * @return this builder for method chaining.
+ */
+ public SELF withCompressionCodec(String columnPath, CompressionCodecName
codecName) {
+ encodingPropsBuilder.withCompressionCodec(columnPath, codecName);
+ return self();
+ }
+
+ /**
+ * Set the compression level for a specific column.
+ * The valid range is codec-specific (e.g. ZSTD: negative levels trade
ratio for speed, up to 22 for best compression (default 3), GZIP: 0–9 or -1
(default 6), BROTLI: 0–11 (default 1)).
+ * Pass {@code null} to unset a previously configured level for that
column.
+ *
+ * @param columnPath dot-string path of the column (e.g. {@code "my_col"})
+ * @param level codec-specific compression level, or {@code null} to
unset
+ * @return this builder for method chaining.
+ */
+ public SELF withCompressionLevel(String columnPath, Integer level) {
+ encodingPropsBuilder.withCompressionLevel(columnPath, level);
+ return self();
+ }
+
/**
* Set the {@link CompressionCodecFactory codec factory} used by the
* constructed writer.
diff --git
a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestColumnChunkPageWriteStore.java
b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestColumnChunkPageWriteStore.java
index 8d735b55a..bd7bbf484 100644
---
a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestColumnChunkPageWriteStore.java
+++
b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestColumnChunkPageWriteStore.java
@@ -22,7 +22,9 @@ import static org.apache.parquet.column.Encoding.PLAIN;
import static org.apache.parquet.column.Encoding.RLE;
import static
org.apache.parquet.format.converter.ParquetMetadataConverter.NO_FILTER;
import static org.apache.parquet.hadoop.metadata.CompressionCodecName.GZIP;
+import static org.apache.parquet.hadoop.metadata.CompressionCodecName.SNAPPY;
import static
org.apache.parquet.hadoop.metadata.CompressionCodecName.UNCOMPRESSED;
+import static org.apache.parquet.hadoop.metadata.CompressionCodecName.ZSTD;
import static org.apache.parquet.schema.OriginalType.UTF8;
import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.BINARY;
import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.DOUBLE;
@@ -32,6 +34,8 @@ import static
org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT64;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.ArgumentMatchers.isNull;
@@ -41,6 +45,7 @@ import static org.mockito.Mockito.inOrder;
import java.io.IOException;
import java.nio.ByteOrder;
import java.util.HashMap;
+import java.util.Map;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
@@ -312,4 +317,132 @@ public class TestColumnChunkPageWriteStore {
private BytesInputCompressor compressor(CompressionCodecName codec) {
return new CodecFactory(conf, pageSize).getCompressor(codec);
}
+
+ @Test
+ public void perColumnCodecDefaultUsedWhenNotSet() throws Exception {
+ MessageType schema =
+ MessageTypeParser.parseMessageType("message test { required binary
col_a; required int32 col_b; }");
+ ParquetProperties props = ParquetProperties.builder().build();
+
+ Map<String, CompressionCodecName> codecs = writeAndReadCodecs(schema,
SNAPPY, props);
+
+ assertEquals(SNAPPY, codecs.get("col_a"));
+ assertEquals(SNAPPY, codecs.get("col_b"));
+ }
+
+ @Test
+ public void perColumnCodecOverridesDefaultForOneColumn() throws Exception {
+ MessageType schema =
+ MessageTypeParser.parseMessageType("message test { required binary
col_a; required int32 col_b; }");
+ ParquetProperties props =
+ ParquetProperties.builder().withCompressionCodec("col_a",
ZSTD).build();
+
+ Map<String, CompressionCodecName> codecs = writeAndReadCodecs(schema,
SNAPPY, props);
+
+ assertEquals(ZSTD, codecs.get("col_a"));
+ assertEquals(SNAPPY, codecs.get("col_b"));
+ }
+
+ @Test
+ public void perColumnCodecAllColumnsOverridden() throws Exception {
+ MessageType schema =
+ MessageTypeParser.parseMessageType("message test { required binary
col_a; required int32 col_b; }");
+ ParquetProperties props = ParquetProperties.builder()
+ .withCompressionCodec("col_a", ZSTD)
+ .withCompressionCodec("col_b", GZIP)
+ .build();
+
+ Map<String, CompressionCodecName> codecs = writeAndReadCodecs(schema,
SNAPPY, props);
+
+ assertEquals(ZSTD, codecs.get("col_a"));
+ assertEquals(GZIP, codecs.get("col_b"));
+ }
+
+ @Test
+ public void perColumnLevelWithCodecRoundTrip() throws Exception {
+ MessageType schema =
+ MessageTypeParser.parseMessageType("message test { required binary
col_a; required int32 col_b; }");
+ ParquetProperties props = ParquetProperties.builder()
+ .withCompressionCodec("col_a", ZSTD)
+ .withCompressionLevel("col_a", 10)
+ .build();
+
+ Map<String, CompressionCodecName> codecs = writeAndReadCodecs(schema,
SNAPPY, props);
+
+ assertEquals(ZSTD, codecs.get("col_a"));
+ assertEquals(SNAPPY, codecs.get("col_b"));
+ }
+
+ @Test
+ public void perColumnLevelInvalidZstdLevelThrows() throws Exception {
+ MessageType schema = MessageTypeParser.parseMessageType("message test {
required binary col_a; }");
+ ParquetProperties props = ParquetProperties.builder()
+ .withCompressionCodec("col_a", ZSTD)
+ .withCompressionLevel("col_a", 23)
+ .build();
+
+ BadConfigurationException ex =
+ assertThrows(BadConfigurationException.class, () ->
writeAndReadCodecs(schema, SNAPPY, props));
+ assertTrue(ex.getMessage().contains("23"));
+ }
+
+ @Test
+ public void perColumnLevelInvalidGzipLevelThrows() throws Exception {
+ MessageType schema = MessageTypeParser.parseMessageType("message test {
required binary col_a; }");
+ ParquetProperties props = ParquetProperties.builder()
+ .withCompressionCodec("col_a", GZIP)
+ .withCompressionLevel("col_a", 10)
+ .build();
+
+ BadConfigurationException ex =
+ assertThrows(BadConfigurationException.class, () ->
writeAndReadCodecs(schema, SNAPPY, props));
+ assertTrue(ex.getMessage().contains("10"));
+ }
+
+ private Map<String, CompressionCodecName> writeAndReadCodecs(
+ MessageType schema, CompressionCodecName defaultCodec, ParquetProperties
props) throws Exception {
+ Path file = new
Path("target/test/TestColumnChunkPageWriteStore/perColumnCodec.parquet");
+ FileSystem fs = file.getFileSystem(conf);
+ fs.delete(file, false);
+ fs.mkdirs(file.getParent());
+
+ allocator = TrackingByteBufferAllocator.wrap(new
HeapByteBufferAllocator());
+ CodecFactory codecFactory = new CodecFactory(conf, pageSize);
+
+ OutputFileForTesting outputFile = new OutputFileForTesting(file, conf);
+ ParquetFileWriter writer = new ParquetFileWriter(
+ outputFile,
+ schema,
+ Mode.CREATE,
+ ParquetWriter.DEFAULT_BLOCK_SIZE,
+ ParquetWriter.MAX_PADDING_SIZE_DEFAULT,
+ null,
+ ParquetProperties.builder().withAllocator(allocator).build());
+ writer.start();
+ writer.startBlock(1);
+
+ try (ColumnChunkPageWriteStore store = ColumnChunkPageWriteStore.builder()
+
.withCompressorProvider(ColumnChunkPageWriteStore.compressorProvider(codecFactory,
defaultCodec, props))
+ .withSchema(schema)
+ .withAllocator(allocator)
+ .withColumnIndexTruncateLength(Integer.MAX_VALUE)
+ .build()) {
+ for (ColumnDescriptor col : schema.getColumns()) {
+ Statistics<?> stats =
+ Statistics.getBuilderForReading(col.getPrimitiveType()).build();
+ store.getPageWriter(col).writePage(BytesInput.fromInt(42), 1, 1,
stats, RLE, RLE, PLAIN);
+ }
+ store.flushToFileWriter(writer);
+ }
+
+ writer.endBlock();
+ writer.end(new HashMap<>());
+
+ ParquetMetadata footer = ParquetFileReader.readFooter(conf, file,
NO_FILTER);
+ Map<String, CompressionCodecName> result = new HashMap<>();
+ for (ColumnChunkMetaData col : footer.getBlocks().get(0).getColumns()) {
+ result.put(col.getPath().toDotString(), col.getCodec());
+ }
+ return result;
+ }
}
diff --git
a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestDirectCodecFactory.java
b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestDirectCodecFactory.java
index c78ee09ec..e4be00947 100644
---
a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestDirectCodecFactory.java
+++
b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestDirectCodecFactory.java
@@ -18,12 +18,17 @@
package org.apache.parquet.hadoop;
import static org.apache.parquet.hadoop.metadata.CompressionCodecName.BROTLI;
+import static org.apache.parquet.hadoop.metadata.CompressionCodecName.GZIP;
import static org.apache.parquet.hadoop.metadata.CompressionCodecName.LZ4;
import static org.apache.parquet.hadoop.metadata.CompressionCodecName.LZ4_RAW;
import static org.apache.parquet.hadoop.metadata.CompressionCodecName.LZO;
+import static org.apache.parquet.hadoop.metadata.CompressionCodecName.SNAPPY;
+import static
org.apache.parquet.hadoop.metadata.CompressionCodecName.UNCOMPRESSED;
+import static org.apache.parquet.hadoop.metadata.CompressionCodecName.ZSTD;
import java.io.IOException;
import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
import java.util.HashSet;
import java.util.Random;
import java.util.Set;
@@ -285,4 +290,158 @@ public class TestDirectCodecFactory {
Assert.assertEquals(codec_2_1, codec_2_2);
Assert.assertNotEquals(codec_2_1, codec_5_1);
}
+
+ @Test
+ public void leveledCompressorCachedForSameLevel() {
+ CodecFactory factory = new CodecFactory(new Configuration(), pageSize);
+ BytesInputCompressor c1 = factory.getCompressor(ZSTD, 3);
+ BytesInputCompressor c2 = factory.getCompressor(ZSTD, 3);
+ Assert.assertSame("Same codec+level should return the cached instance",
c1, c2);
+ factory.release();
+ }
+
+ @Test
+ public void leveledCompressorDiffersByLevel() {
+ CodecFactory factory = new CodecFactory(new Configuration(), pageSize);
+ BytesInputCompressor c1 = factory.getCompressor(ZSTD, 1);
+ BytesInputCompressor c3 = factory.getCompressor(ZSTD, 3);
+ Assert.assertNotSame("Different levels should return different compressor
instances", c1, c3);
+ factory.release();
+ }
+
+ @Test
+ public void leveledCacheIsolatedFromNoLevelCache() {
+ CodecFactory factory = new CodecFactory(new Configuration(), pageSize);
+ BytesInputCompressor noLevel = factory.getCompressor(ZSTD);
+ BytesInputCompressor withLevel = factory.getCompressor(ZSTD, 3);
+ Assert.assertNotSame(
+ "Level-aware and no-level compressors should use separate cache
entries", noLevel, withLevel);
+ factory.release();
+ }
+
+ @Test
+ public void leveledUncompressedReturnsNoOp() {
+ CodecFactory factory = new CodecFactory(new Configuration(), pageSize);
+ BytesInputCompressor comp = factory.getCompressor(UNCOMPRESSED, 5);
+ Assert.assertSame(CodecFactory.NO_OP_COMPRESSOR, comp);
+ factory.release();
+ }
+
+ @Test
+ public void leveledSnappyIgnoresLevel() {
+ CodecFactory factory = new CodecFactory(new Configuration(), pageSize);
+ BytesInputCompressor comp = factory.getCompressor(SNAPPY, 99);
+ Assert.assertNotNull(comp);
+ Assert.assertEquals(SNAPPY, comp.getCodecName());
+ factory.release();
+ }
+
+ @Test
+ public void leveledGzipInvalidLevelThrows() {
+ CodecFactory factory = new CodecFactory(new Configuration(), pageSize);
+ try {
+ BadConfigurationException ex =
+ Assert.assertThrows(BadConfigurationException.class, () ->
factory.getCompressor(GZIP, 99));
+ Assert.assertTrue(ex.getMessage().contains("99"));
+ } finally {
+ factory.release();
+ }
+ }
+
+ @Test
+ public void leveledGzipBoundaryLevelsValid() {
+ CodecFactory factory = new CodecFactory(new Configuration(), pageSize);
+ for (int level : new int[] {-1, 0, 1, 9}) {
+ BytesInputCompressor comp = factory.getCompressor(GZIP, level);
+ Assert.assertNotNull("Compressor should not be null for GZIP level " +
level, comp);
+ Assert.assertEquals("Codec name should be GZIP for level " + level,
GZIP, comp.getCodecName());
+ }
+ factory.release();
+ }
+
+ @Test
+ public void leveledZstdRoundTrip() throws IOException {
+ CodecFactory factory = new CodecFactory(new Configuration(), pageSize);
+ byte[] original = "hello parquet per-column
compression".getBytes(StandardCharsets.UTF_8);
+ BytesInputDecompressor decompressor = factory.getDecompressor(ZSTD);
+ for (int level : new int[] {-5, 0, 1, 3, 10, 22}) {
+ BytesInput compressed = factory.getCompressor(ZSTD,
level).compress(BytesInput.from(original));
+ byte[] result = decompressor.decompress(compressed,
original.length).toByteArray();
+ Assert.assertArrayEquals("Round-trip failed at ZSTD level " + level,
original, result);
+ }
+ factory.release();
+ }
+
+ @Test
+ public void leveledGzipRoundTrip() throws IOException {
+ CodecFactory factory = new CodecFactory(new Configuration(), pageSize);
+ byte[] original = "hello parquet per-column
compression".getBytes(StandardCharsets.UTF_8);
+ BytesInputDecompressor decompressor = factory.getDecompressor(GZIP);
+ for (int level : new int[] {1, 5, 9}) {
+ BytesInput compressed = factory.getCompressor(GZIP,
level).compress(BytesInput.from(original));
+ byte[] result = decompressor.decompress(compressed,
original.length).toByteArray();
+ Assert.assertArrayEquals("Round-trip failed at GZIP level " + level,
original, result);
+ }
+ factory.release();
+ }
+
+ @Test
+ public void directFactoryLeveledZstdRoundTrip() throws IOException {
+ CodecFactory heap = new CodecFactory(new Configuration(), pageSize);
+ CodecFactory direct =
+ CodecFactory.createDirectCodecFactory(new Configuration(), new
DirectByteBufferAllocator(), pageSize);
+ try {
+ // Sanity: the heap factory's leveled path yields a HeapBytesCompressor.
+ Assert.assertTrue(
+ "heap factory should produce a HeapBytesCompressor",
+ heap.getCompressor(ZSTD, 3) instanceof
CodecFactory.HeapBytesCompressor);
+
+ // The direct factory must not fall back to the heap/Hadoop path for
leveled ZSTD/SNAPPY.
+ BytesInputCompressor directZstd = direct.getCompressor(ZSTD, 3);
+ Assert.assertFalse(
+ "direct factory ZSTD(level) should not fall back to
HeapBytesCompressor",
+ directZstd instanceof CodecFactory.HeapBytesCompressor);
+ Assert.assertEquals(ZSTD, directZstd.getCodecName());
+ Assert.assertEquals(
+ "leveled ZSTD should use the same direct compressor type as the
no-level path",
+ direct.getCompressor(ZSTD).getClass(),
+ directZstd.getClass());
+
+ BytesInputCompressor directSnappy = direct.getCompressor(SNAPPY, 5);
+ Assert.assertFalse(
+ "direct factory SNAPPY(level) should not fall back to
HeapBytesCompressor",
+ directSnappy instanceof CodecFactory.HeapBytesCompressor);
+ Assert.assertEquals(
+ "leveled SNAPPY should use the same direct compressor type as the
no-level path",
+ direct.getCompressor(SNAPPY).getClass(),
+ directSnappy.getClass());
+
+ // The direct ZSTD level path must compress/decompress correctly at each
level.
+ byte[] original = "hello parquet per-column zstd direct
compression".getBytes(StandardCharsets.UTF_8);
+ BytesInputDecompressor decompressor = heap.getDecompressor(ZSTD);
+ for (int level : new int[] {1, 3, 22}) {
+ byte[] compressed = direct.getCompressor(ZSTD, level)
+ .compress(BytesInput.from(original))
+ .toByteArray();
+ byte[] result = decompressor
+ .decompress(BytesInput.from(compressed), original.length)
+ .toByteArray();
+ Assert.assertArrayEquals("Direct ZSTD round-trip failed at level " +
level, original, result);
+ }
+ } finally {
+ heap.release();
+ direct.release();
+ }
+ }
+
+ @Test
+ public void directFactoryInvalidZstdLevelThrows() {
+ CodecFactory direct =
+ CodecFactory.createDirectCodecFactory(new Configuration(), new
DirectByteBufferAllocator(), pageSize);
+ try {
+ Assert.assertThrows(BadConfigurationException.class, () ->
direct.getCompressor(ZSTD, 23));
+ } finally {
+ direct.release();
+ }
+ }
}
diff --git
a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetWriter.java
b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetWriter.java
index 6b3b0aae2..02e114a0f 100644
---
a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetWriter.java
+++
b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetWriter.java
@@ -27,7 +27,10 @@ import static
org.apache.parquet.column.ParquetProperties.WriterVersion.PARQUET_
import static
org.apache.parquet.format.converter.ParquetMetadataConverter.NO_FILTER;
import static org.apache.parquet.hadoop.ParquetFileReader.readFooter;
import static org.apache.parquet.hadoop.TestUtils.enforceEmptyDir;
+import static org.apache.parquet.hadoop.metadata.CompressionCodecName.GZIP;
+import static org.apache.parquet.hadoop.metadata.CompressionCodecName.SNAPPY;
import static
org.apache.parquet.hadoop.metadata.CompressionCodecName.UNCOMPRESSED;
+import static org.apache.parquet.hadoop.metadata.CompressionCodecName.ZSTD;
import static org.apache.parquet.schema.LogicalTypeAnnotation.stringType;
import static org.apache.parquet.schema.MessageTypeParser.parseMessageType;
import static org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.BINARY;
@@ -55,6 +58,8 @@ import org.apache.commons.lang3.RandomStringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.mapreduce.Job;
+import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.parquet.ParquetReadOptions;
import org.apache.parquet.bytes.HeapByteBufferAllocator;
import org.apache.parquet.bytes.TrackingByteBufferAllocator;
@@ -86,6 +91,7 @@ import org.apache.parquet.hadoop.example.GroupWriteSupport;
import org.apache.parquet.hadoop.metadata.BlockMetaData;
import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
import org.apache.parquet.hadoop.metadata.ColumnPath;
+import org.apache.parquet.hadoop.metadata.CompressionCodecName;
import org.apache.parquet.hadoop.metadata.ParquetMetadata;
import org.apache.parquet.hadoop.util.HadoopInputFile;
import org.apache.parquet.hadoop.util.HadoopOutputFile;
@@ -826,6 +832,231 @@ public class TestParquetWriter {
ExampleParquetWriter.builder(path).withFile(outputFile).build());
}
+ @Test
+ public void perColumnCodecOverridesDefaultForOneColumn() throws Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .required(INT32)
+ .named("col_b")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+
+ try (ParquetWriter<Group> writer = ExampleParquetWriter.builder(path)
+ .withAllocator(allocator)
+ .withType(schema)
+ .withCompressionCodec(SNAPPY)
+ .withCompressionCodec("col_a", ZSTD)
+ .build()) {
+ SimpleGroupFactory f = new SimpleGroupFactory(schema);
+ writer.write(f.newGroup().append("col_a", "hello").append("col_b", 1));
+ }
+
+ ParquetMetadata footer = readFooter(new Configuration(), path, NO_FILTER);
+ Map<String, CompressionCodecName> codecs = new HashMap<>();
+ for (ColumnChunkMetaData col : footer.getBlocks().get(0).getColumns()) {
+ codecs.put(col.getPath().toDotString(), col.getCodec());
+ }
+ assertEquals(ZSTD, codecs.get("col_a"));
+ assertEquals(SNAPPY, codecs.get("col_b"));
+ }
+
+ @Test
+ public void perColumnCodecDefaultUsedWhenNoOverride() throws Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .required(INT32)
+ .named("col_b")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+
+ try (ParquetWriter<Group> writer = ExampleParquetWriter.builder(path)
+ .withAllocator(allocator)
+ .withType(schema)
+ .withCompressionCodec(GZIP)
+ .build()) {
+ SimpleGroupFactory f = new SimpleGroupFactory(schema);
+ writer.write(f.newGroup().append("col_a", "hello").append("col_b", 1));
+ }
+
+ ParquetMetadata footer = readFooter(new Configuration(), path, NO_FILTER);
+ for (ColumnChunkMetaData col : footer.getBlocks().get(0).getColumns()) {
+ assertEquals("Column " + col.getPath().toDotString() + " should use
default codec", GZIP, col.getCodec());
+ }
+ }
+
+ @Test
+ public void perColumnLevelDataRoundTrips() throws Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .required(INT32)
+ .named("col_b")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+
+ try (ParquetWriter<Group> writer = ExampleParquetWriter.builder(path)
+ .withAllocator(allocator)
+ .withType(schema)
+ .withCompressionCodec(SNAPPY)
+ .withCompressionCodec("col_a", ZSTD)
+ .withCompressionLevel("col_a", 1)
+ .build()) {
+ SimpleGroupFactory f = new SimpleGroupFactory(schema);
+ writer.write(f.newGroup().append("col_a", "hello").append("col_b", 42));
+ }
+
+ try (ParquetReader<Group> reader =
+ ParquetReader.builder(new GroupReadSupport(), path).build()) {
+ Group group = reader.read();
+ assertEquals("hello", group.getBinary("col_a", 0).toStringUsingUTF8());
+ assertEquals(42, group.getInteger("col_b", 0));
+ assertNull(reader.read());
+ }
+
+ ParquetMetadata footer = readFooter(new Configuration(), path, NO_FILTER);
+ Map<String, CompressionCodecName> codecs = new HashMap<>();
+ for (ColumnChunkMetaData col : footer.getBlocks().get(0).getColumns()) {
+ codecs.put(col.getPath().toDotString(), col.getCodec());
+ }
+ assertEquals(ZSTD, codecs.get("col_a"));
+ assertEquals(SNAPPY, codecs.get("col_b"));
+ }
+
+ @Test
+ public void perColumnLevelInvalidZstdLevelThrows() throws Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+
+ // ZSTD only supports levels 1-22; level 23 is invalid
+ Assert.assertThrows(BadConfigurationException.class, () -> {
+ try (ParquetWriter<Group> writer = ExampleParquetWriter.builder(path)
+ .withAllocator(allocator)
+ .withType(schema)
+ .withCompressionCodec("col_a", ZSTD)
+ .withCompressionLevel("col_a", 23)
+ .build()) {
+ // exception expected before first write
+ }
+ });
+ }
+
+ @Test
+ public void perColumnLevelInvalidLevelThrows() throws Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+
+ // GZIP only supports levels -1 (default) through 9; level 10 is invalid
+ Assert.assertThrows(BadConfigurationException.class, () -> {
+ try (ParquetWriter<Group> writer = ExampleParquetWriter.builder(path)
+ .withAllocator(allocator)
+ .withType(schema)
+ .withCompressionCodec("col_a", GZIP)
+ .withCompressionLevel("col_a", 10)
+ .build()) {
+ // exception expected before first write
+ }
+ });
+ }
+
+ @Test
+ public void perColumnLevelDifferentLevelsPerColumnDataRoundTrips() throws
Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .required(BINARY)
+ .as(stringType())
+ .named("col_b")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+
+ // Both columns use ZSTD (from default), but at different levels.
+ // This exercises the level-only override path in resolveCompressor().
+ try (ParquetWriter<Group> writer = ExampleParquetWriter.builder(path)
+ .withAllocator(allocator)
+ .withType(schema)
+ .withCompressionCodec(ZSTD)
+ .withCompressionLevel("col_a", 1)
+ .withCompressionLevel("col_b", 10)
+ .build()) {
+ SimpleGroupFactory f = new SimpleGroupFactory(schema);
+ writer.write(f.newGroup().append("col_a", "fast").append("col_b",
"best"));
+ }
+
+ // Both columns must report ZSTD in the footer
+ ParquetMetadata footer = readFooter(new Configuration(), path, NO_FILTER);
+ for (ColumnChunkMetaData col : footer.getBlocks().get(0).getColumns()) {
+ assertEquals("Column " + col.getPath().toDotString() + " should use
ZSTD", ZSTD, col.getCodec());
+ }
+
+ // Data must survive the round-trip at both levels
+ try (ParquetReader<Group> reader =
+ ParquetReader.builder(new GroupReadSupport(), path).build()) {
+ Group group = reader.read();
+ assertEquals("fast", group.getBinary("col_a", 0).toStringUsingUTF8());
+ assertEquals("best", group.getBinary("col_b", 0).toStringUsingUTF8());
+ assertNull(reader.read());
+ }
+ }
+
+ @Test
+ public void perColumnCodecAllColumnsOverridden() throws Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .required(INT32)
+ .named("col_b")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+
+ try (ParquetWriter<Group> writer = ExampleParquetWriter.builder(path)
+ .withAllocator(allocator)
+ .withType(schema)
+ .withCompressionCodec(SNAPPY)
+ .withCompressionCodec("col_a", ZSTD)
+ .withCompressionCodec("col_b", GZIP)
+ .build()) {
+ SimpleGroupFactory f = new SimpleGroupFactory(schema);
+ writer.write(f.newGroup().append("col_a", "hello").append("col_b", 1));
+ }
+
+ ParquetMetadata footer = readFooter(new Configuration(), path, NO_FILTER);
+ Map<String, CompressionCodecName> codecs = new HashMap<>();
+ for (ColumnChunkMetaData col : footer.getBlocks().get(0).getColumns()) {
+ codecs.put(col.getPath().toDotString(), col.getCodec());
+ }
+ assertEquals(ZSTD, codecs.get("col_a"));
+ assertEquals(GZIP, codecs.get("col_b"));
+ }
+
@Test
public void testNoFlushAfterException() throws Exception {
final File testDir = temp.newFile();
@@ -864,6 +1095,163 @@ public class TestParquetWriter {
assertTrue(!fs.exists(file) || fs.getFileStatus(file).getLen() == 0);
}
+ @Test
+ public void outputFormatSetColumnCompressionOverridesDefaultForOneColumn()
throws Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .required(INT32)
+ .named("col_b")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+ Job job = Job.getInstance();
+ ParquetOutputFormat.setColumnCompression(job, "col_a", ZSTD);
+
+ Map<String, CompressionCodecName> codecs =
writeAndReadCodecsViaOutputFormat(schema, SNAPPY, job, path);
+
+ assertEquals(ZSTD, codecs.get("col_a"));
+ assertEquals(SNAPPY, codecs.get("col_b"));
+ }
+
+ @Test
+ public void outputFormatSetColumnCompressionAllColumnsOverridden() throws
Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .required(INT32)
+ .named("col_b")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+ Job job = Job.getInstance();
+ ParquetOutputFormat.setColumnCompression(job, "col_a", ZSTD);
+ ParquetOutputFormat.setColumnCompression(job, "col_b", GZIP);
+
+ Map<String, CompressionCodecName> codecs =
writeAndReadCodecsViaOutputFormat(schema, SNAPPY, job, path);
+
+ assertEquals(ZSTD, codecs.get("col_a"));
+ assertEquals(GZIP, codecs.get("col_b"));
+ }
+
+ @Test
+ public void outputFormatSetColumnCompressionDefaultUsedWhenNotSet() throws
Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .required(INT32)
+ .named("col_b")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+ Job job = Job.getInstance();
+
+ Map<String, CompressionCodecName> codecs =
writeAndReadCodecsViaOutputFormat(schema, GZIP, job, path);
+
+ assertEquals(GZIP, codecs.get("col_a"));
+ assertEquals(GZIP, codecs.get("col_b"));
+ }
+
+ @Test
+ public void outputFormatSetColumnCompressionLevelWithCodecDataRoundTrips()
throws Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .required(INT32)
+ .named("col_b")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+ Job job = Job.getInstance();
+ ParquetOutputFormat.setColumnCompression(job, "col_a", ZSTD);
+ ParquetOutputFormat.setColumnCompressionLevel(job, "col_a", 1);
+
+ Map<String, CompressionCodecName> codecs =
writeAndReadCodecsViaOutputFormat(schema, SNAPPY, job, path);
+ assertEquals(ZSTD, codecs.get("col_a"));
+ assertEquals(SNAPPY, codecs.get("col_b"));
+
+ try (ParquetReader<Group> reader =
+ ParquetReader.builder(new GroupReadSupport(), path).build()) {
+ Group group = reader.read();
+ assertEquals("hello", group.getBinary("col_a", 0).toStringUsingUTF8());
+ assertEquals(42, group.getInteger("col_b", 0));
+ assertNull(reader.read());
+ }
+ }
+
+ @Test
+ public void outputFormatSetColumnCompressionLevelDifferentLevelsPerColumn()
throws Exception {
+ MessageType schema = Types.buildMessage()
+ .required(BINARY)
+ .as(stringType())
+ .named("col_a")
+ .required(BINARY)
+ .as(stringType())
+ .named("col_b")
+ .named("test");
+ File file = temp.newFile();
+ file.delete();
+ Path path = new Path(file.getAbsolutePath());
+ Job job = Job.getInstance();
+ ParquetOutputFormat.setColumnCompressionLevel(job, "col_a", 1);
+ ParquetOutputFormat.setColumnCompressionLevel(job, "col_b", 10);
+
+ Map<String, CompressionCodecName> codecs =
writeAndReadCodecsViaOutputFormat(schema, ZSTD, job, path);
+ assertEquals(ZSTD, codecs.get("col_a"));
+ assertEquals(ZSTD, codecs.get("col_b"));
+
+ try (ParquetReader<Group> reader =
+ ParquetReader.builder(new GroupReadSupport(), path).build()) {
+ Group group = reader.read();
+ assertEquals("hello", group.getBinary("col_a", 0).toStringUsingUTF8());
+ assertEquals("fast", group.getBinary("col_b", 0).toStringUsingUTF8());
+ assertNull(reader.read());
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private Map<String, CompressionCodecName> writeAndReadCodecsViaOutputFormat(
+ MessageType schema, CompressionCodecName defaultCodec, Job job, Path
file) throws Exception {
+ Configuration conf = job.getConfiguration();
+ GroupWriteSupport.setSchema(schema, conf);
+ SimpleGroupFactory f = new SimpleGroupFactory(schema);
+
+ Group group = f.newGroup();
+ for (int i = 0; i < schema.getFieldCount(); i++) {
+ String name = schema.getFieldName(i);
+ switch (schema.getType(i).asPrimitiveType().getPrimitiveTypeName()) {
+ case BINARY:
+ group.append(name, name.equals("col_a") ? "hello" : "fast");
+ break;
+ case INT32:
+ group.append(name, 42);
+ break;
+ default:
+ break;
+ }
+ }
+
+ ParquetOutputFormat<Group> outputFormat = new ParquetOutputFormat<>(new
GroupWriteSupport());
+ RecordWriter<Void, Group> writer = outputFormat.getRecordWriter(conf,
file, defaultCodec);
+ writer.write(null, group);
+ writer.close(null);
+
+ ParquetMetadata footer = readFooter(conf, file, NO_FILTER);
+ Map<String, CompressionCodecName> result = new HashMap<>();
+ for (ColumnChunkMetaData col : footer.getBlocks().get(0).getColumns()) {
+ result.put(col.getPath().toDotString(), col.getCodec());
+ }
+ return result;
+ }
+
@Test
public void testV2PageNullCountWithStatisticsDisabled() throws Exception {
// Regression test: when using PARQUET_2_0 with statistics disabled on a
nullable column,
