This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new a8074ed60c [parquet] Support column compression configuration (#8279)
a8074ed60c is described below
commit a8074ed60c760972fb2a6ac9fb8218f6cd7d96a3
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Jun 18 17:09:04 2026 +0800
[parquet] Support column compression configuration (#8279)
Support configuring Parquet compression codecs per column. This lets
users keep a table-level codec such as ZSTD while overriding selected
columns to another codec, including `none` / `UNCOMPRESSED`.
---
.../parquet/writer/RowDataParquetBuilder.java | 15 +-
.../parquet/hadoop/ColumnCompressionCodecs.java | 52 ++
.../hadoop/ColumnCompressionPageWriteStore.java | 660 +++++++++++++++++++++
.../hadoop/ColumnCompressionRecordWriter.java | 242 ++++++++
.../org/apache/parquet/hadoop/ParquetWriter.java | 26 +-
.../format/parquet/ParquetFormatReadWriteTest.java | 43 ++
6 files changed, 1031 insertions(+), 7 deletions(-)
diff --git
a/paimon-format/src/main/java/org/apache/paimon/format/parquet/writer/RowDataParquetBuilder.java
b/paimon-format/src/main/java/org/apache/paimon/format/parquet/writer/RowDataParquetBuilder.java
index 7d72c8a1a9..2e84df0932 100644
---
a/paimon-format/src/main/java/org/apache/paimon/format/parquet/writer/RowDataParquetBuilder.java
+++
b/paimon-format/src/main/java/org/apache/paimon/format/parquet/writer/RowDataParquetBuilder.java
@@ -19,6 +19,7 @@
package org.apache.paimon.format.parquet.writer;
import org.apache.paimon.data.InternalRow;
+import org.apache.paimon.format.HadoopCompressionType;
import org.apache.paimon.format.parquet.ColumnConfigParser;
import org.apache.paimon.format.parquet.VariantUtils;
import org.apache.paimon.options.Options;
@@ -60,8 +61,7 @@ public class RowDataParquetBuilder implements
ParquetBuilder<InternalRow> {
ParquetRowDataBuilder builder =
new ParquetRowDataBuilder(out, rowType, shreddingSchemas)
.withConf(conf)
- .withCompressionCodec(
-
CompressionCodecName.fromConf(getCompression(compression)))
+
.withCompressionCodec(getCompressionCodec(getCompression(compression)))
.withRowGroupSize(
conf.getLong(
ParquetOutputFormat.BLOCK_SIZE,
@@ -114,6 +114,10 @@ public class RowDataParquetBuilder implements
ParquetBuilder<InternalRow> {
ParquetOutputFormat.COLUMN_INDEX_TRUNCATE_LENGTH,
ParquetProperties.DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH));
new ColumnConfigParser()
+ .withColumnConfig(
+ ParquetOutputFormat.COMPRESSION,
+ key -> getCompressionCodec(conf.get(key)),
+ builder::withCompressionCodec)
.withColumnConfig(
ParquetOutputFormat.ENABLE_DICTIONARY,
key -> conf.getBoolean(key, false),
@@ -137,4 +141,11 @@ public class RowDataParquetBuilder implements
ParquetBuilder<InternalRow> {
public String getCompression(String compression) {
return conf.get("parquet.compression", compression);
}
+
+ private CompressionCodecName getCompressionCodec(String compression) {
+ if (HadoopCompressionType.NONE.value().equalsIgnoreCase(compression)) {
+ return CompressionCodecName.UNCOMPRESSED;
+ }
+ return CompressionCodecName.fromConf(compression);
+ }
}
diff --git
a/paimon-format/src/main/java/org/apache/parquet/hadoop/ColumnCompressionCodecs.java
b/paimon-format/src/main/java/org/apache/parquet/hadoop/ColumnCompressionCodecs.java
new file mode 100644
index 0000000000..4f9c6842d8
--- /dev/null
+++
b/paimon-format/src/main/java/org/apache/parquet/hadoop/ColumnCompressionCodecs.java
@@ -0,0 +1,52 @@
+/*
+ * 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.hadoop;
+
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.hadoop.metadata.ColumnPath;
+import org.apache.parquet.hadoop.metadata.CompressionCodecName;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/** Compression codecs configured per parquet column path. */
+class ColumnCompressionCodecs {
+
+ private final CompressionCodecName defaultCodec;
+ private final Map<ColumnPath, CompressionCodecName> columnCodecs;
+
+ ColumnCompressionCodecs(CompressionCodecName defaultCodec) {
+ this(defaultCodec, Collections.<ColumnPath,
CompressionCodecName>emptyMap());
+ }
+
+ ColumnCompressionCodecs(
+ CompressionCodecName defaultCodec, Map<ColumnPath,
CompressionCodecName> columnCodecs) {
+ this.defaultCodec = defaultCodec;
+ this.columnCodecs =
+ columnCodecs.isEmpty()
+ ? Collections.<ColumnPath,
CompressionCodecName>emptyMap()
+ : new HashMap<>(columnCodecs);
+ }
+
+ CompressionCodecName getCodec(ColumnDescriptor descriptor) {
+ CompressionCodecName codec =
columnCodecs.get(ColumnPath.get(descriptor.getPath()));
+ return codec == null ? defaultCodec : codec;
+ }
+}
diff --git
a/paimon-format/src/main/java/org/apache/parquet/hadoop/ColumnCompressionPageWriteStore.java
b/paimon-format/src/main/java/org/apache/parquet/hadoop/ColumnCompressionPageWriteStore.java
new file mode 100644
index 0000000000..12250aa3f8
--- /dev/null
+++
b/paimon-format/src/main/java/org/apache/parquet/hadoop/ColumnCompressionPageWriteStore.java
@@ -0,0 +1,660 @@
+/*
+ * 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.hadoop;
+
+import org.apache.parquet.bytes.ByteBufferAllocator;
+import org.apache.parquet.bytes.ByteBufferReleaser;
+import org.apache.parquet.bytes.BytesInput;
+import org.apache.parquet.bytes.ConcatenatingByteBufferCollector;
+import org.apache.parquet.column.ColumnDescriptor;
+import org.apache.parquet.column.Encoding;
+import org.apache.parquet.column.page.DictionaryPage;
+import org.apache.parquet.column.page.PageWriteStore;
+import org.apache.parquet.column.page.PageWriter;
+import org.apache.parquet.column.statistics.SizeStatistics;
+import org.apache.parquet.column.statistics.Statistics;
+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.crypto.AesCipher;
+import org.apache.parquet.crypto.InternalColumnEncryptionSetup;
+import org.apache.parquet.crypto.InternalFileEncryptor;
+import org.apache.parquet.crypto.ModuleCipherFactory.ModuleType;
+import org.apache.parquet.format.BlockCipher;
+import org.apache.parquet.format.converter.ParquetMetadataConverter;
+import org.apache.parquet.hadoop.metadata.ColumnPath;
+import org.apache.parquet.internal.column.columnindex.ColumnIndexBuilder;
+import org.apache.parquet.internal.column.columnindex.OffsetIndexBuilder;
+import org.apache.parquet.io.ParquetEncodingException;
+import org.apache.parquet.schema.MessageType;
+import org.apache.parquet.util.AutoCloseables;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.zip.CRC32;
+
+/**
+ * Page write store with per-column compression.
+ *
+ * <p>NOTE: The file is adapted from Apache parquet's {@link
ColumnChunkPageWriteStore}.
+ */
+public class ColumnCompressionPageWriteStore implements PageWriteStore,
BloomFilterWriteStore {
+
+ private static final Logger LOG =
+ LoggerFactory.getLogger(ColumnCompressionPageWriteStore.class);
+
+ private static final ParquetMetadataConverter PARQUET_METADATA_CONVERTER =
+ new ParquetMetadataConverter();
+
+ private static final class ColumnChunkPageWriter implements PageWriter,
BloomFilterWriter {
+
+ private final ColumnDescriptor path;
+ private final CompressionCodecFactory.BytesInputCompressor compressor;
+
+ private final ByteArrayOutputStream tempOutputStream = new
ByteArrayOutputStream();
+ private final ConcatenatingByteBufferCollector buf;
+ private DictionaryPage dictionaryPage;
+
+ private long uncompressedLength;
+ private long compressedLength;
+ private long totalValueCount;
+ private int pageCount;
+
+ private Set<Encoding> rlEncodings = new HashSet<>();
+ private Set<Encoding> dlEncodings = new HashSet<>();
+ private List<Encoding> dataEncodings = new ArrayList<>();
+
+ private BloomFilter bloomFilter;
+ private ColumnIndexBuilder columnIndexBuilder;
+ private OffsetIndexBuilder offsetIndexBuilder;
+ private Statistics totalStatistics;
+ private final SizeStatistics totalSizeStatistics;
+ private final GeospatialStatistics totalGeospatialStatistics;
+ private final ByteBufferReleaser releaser;
+
+ private final CRC32 crc;
+ private final boolean pageWriteChecksumEnabled;
+
+ private final BlockCipher.Encryptor headerBlockEncryptor;
+ private final BlockCipher.Encryptor pageBlockEncryptor;
+ private final int rowGroupOrdinal;
+ private final int columnOrdinal;
+ private int pageOrdinal;
+ private final byte[] dataPageAAD;
+ private final byte[] dataPageHeaderAAD;
+ private final byte[] fileAAD;
+
+ private ColumnChunkPageWriter(
+ ColumnDescriptor path,
+ CompressionCodecFactory.BytesInputCompressor compressor,
+ ByteBufferAllocator allocator,
+ int columnIndexTruncateLength,
+ boolean pageWriteChecksumEnabled,
+ BlockCipher.Encryptor headerBlockEncryptor,
+ BlockCipher.Encryptor pageBlockEncryptor,
+ byte[] fileAAD,
+ int rowGroupOrdinal,
+ int columnOrdinal) {
+ this.path = path;
+ this.compressor = compressor;
+ this.releaser = new ByteBufferReleaser(allocator);
+ this.buf = new ConcatenatingByteBufferCollector(allocator);
+ this.columnIndexBuilder =
+ ColumnIndexBuilder.getBuilder(
+ path.getPrimitiveType(),
columnIndexTruncateLength);
+ this.offsetIndexBuilder = OffsetIndexBuilder.getBuilder();
+ this.totalSizeStatistics =
+ SizeStatistics.newBuilder(
+ path.getPrimitiveType(),
+ path.getMaxRepetitionLevel(),
+ path.getMaxDefinitionLevel())
+ .build();
+ this.totalGeospatialStatistics =
+
GeospatialStatistics.newBuilder(path.getPrimitiveType()).build();
+ this.pageWriteChecksumEnabled = pageWriteChecksumEnabled;
+ this.crc = pageWriteChecksumEnabled ? new CRC32() : null;
+ this.headerBlockEncryptor = headerBlockEncryptor;
+ this.pageBlockEncryptor = pageBlockEncryptor;
+ this.fileAAD = fileAAD;
+ this.rowGroupOrdinal = rowGroupOrdinal;
+ this.columnOrdinal = columnOrdinal;
+ this.pageOrdinal = -1;
+ if (headerBlockEncryptor != null) {
+ dataPageHeaderAAD =
+ AesCipher.createModuleAAD(
+ fileAAD,
+ ModuleType.DataPageHeader,
+ rowGroupOrdinal,
+ columnOrdinal,
+ 0);
+ } else {
+ dataPageHeaderAAD = null;
+ }
+ if (pageBlockEncryptor != null) {
+ dataPageAAD =
+ AesCipher.createModuleAAD(
+ fileAAD, ModuleType.DataPage, rowGroupOrdinal,
columnOrdinal, 0);
+ } else {
+ dataPageAAD = null;
+ }
+ }
+
+ @Override
+ @Deprecated
+ public void writePage(
+ BytesInput bytesInput,
+ int valueCount,
+ Statistics<?> statistics,
+ Encoding rlEncoding,
+ Encoding dlEncoding,
+ Encoding valuesEncoding)
+ throws IOException {
+ columnIndexBuilder = ColumnIndexBuilder.getNoOpBuilder();
+ offsetIndexBuilder = OffsetIndexBuilder.getNoOpBuilder();
+
+ writePage(
+ bytesInput, valueCount, -1, statistics, rlEncoding,
dlEncoding, valuesEncoding);
+ }
+
+ @Override
+ public void writePage(
+ BytesInput bytes,
+ int valueCount,
+ int rowCount,
+ Statistics<?> statistics,
+ Encoding rlEncoding,
+ Encoding dlEncoding,
+ Encoding valuesEncoding)
+ throws IOException {
+ writePage(
+ bytes,
+ valueCount,
+ rowCount,
+ statistics,
+ null,
+ null,
+ rlEncoding,
+ dlEncoding,
+ valuesEncoding);
+ }
+
+ @Override
+ public void writePage(
+ BytesInput bytes,
+ int valueCount,
+ int rowCount,
+ Statistics<?> statistics,
+ SizeStatistics sizeStatistics,
+ GeospatialStatistics geospatialStatistics,
+ Encoding rlEncoding,
+ Encoding dlEncoding,
+ Encoding valuesEncoding)
+ throws IOException {
+ pageOrdinal++;
+ long uncompressedSize = bytes.size();
+ if (uncompressedSize > Integer.MAX_VALUE || uncompressedSize < 0) {
+ throw new ParquetEncodingException(
+ "Cannot write page larger than Integer.MAX_VALUE or
negative bytes: "
+ + uncompressedSize);
+ }
+ BytesInput compressedBytes = compressor.compress(bytes);
+ if (pageBlockEncryptor != null) {
+ AesCipher.quickUpdatePageAAD(dataPageAAD, pageOrdinal);
+ compressedBytes =
+ BytesInput.from(
+ pageBlockEncryptor.encrypt(
+ compressedBytes.toByteArray(),
dataPageAAD));
+ }
+ long compressedSize = compressedBytes.size();
+ if (compressedSize > Integer.MAX_VALUE) {
+ throw new ParquetEncodingException(
+ "Cannot write compressed page larger than
Integer.MAX_VALUE bytes: "
+ + compressedSize);
+ }
+ tempOutputStream.reset();
+ if (headerBlockEncryptor != null) {
+ AesCipher.quickUpdatePageAAD(dataPageHeaderAAD, pageOrdinal);
+ }
+ if (pageWriteChecksumEnabled) {
+ crc.reset();
+ crc.update(compressedBytes.toByteArray());
+ PARQUET_METADATA_CONVERTER.writeDataPageV1Header(
+ (int) uncompressedSize,
+ (int) compressedSize,
+ valueCount,
+ rlEncoding,
+ dlEncoding,
+ valuesEncoding,
+ (int) crc.getValue(),
+ tempOutputStream,
+ headerBlockEncryptor,
+ dataPageHeaderAAD);
+ } else {
+ PARQUET_METADATA_CONVERTER.writeDataPageV1Header(
+ (int) uncompressedSize,
+ (int) compressedSize,
+ valueCount,
+ rlEncoding,
+ dlEncoding,
+ valuesEncoding,
+ tempOutputStream,
+ headerBlockEncryptor,
+ dataPageHeaderAAD);
+ }
+ this.uncompressedLength += uncompressedSize;
+ this.compressedLength += compressedSize;
+ this.totalValueCount += valueCount;
+ this.pageCount += 1;
+
+ mergeColumnStatistics(statistics, sizeStatistics,
geospatialStatistics);
+ offsetIndexBuilder.add(
+ toIntWithCheck(tempOutputStream.size() + compressedSize),
+ rowCount,
+ sizeStatistics != null
+ ? sizeStatistics.getUnencodedByteArrayDataBytes()
+ : Optional.empty());
+
+ buf.collect(BytesInput.concat(BytesInput.from(tempOutputStream),
compressedBytes));
+ rlEncodings.add(rlEncoding);
+ dlEncodings.add(dlEncoding);
+ dataEncodings.add(valuesEncoding);
+ }
+
+ @Override
+ public void writePageV2(
+ int rowCount,
+ int nullCount,
+ int valueCount,
+ BytesInput repetitionLevels,
+ BytesInput definitionLevels,
+ Encoding dataEncoding,
+ BytesInput data,
+ Statistics<?> statistics)
+ throws IOException {
+ writePageV2(
+ rowCount,
+ nullCount,
+ valueCount,
+ repetitionLevels,
+ definitionLevels,
+ dataEncoding,
+ data,
+ statistics,
+ null,
+ null);
+ }
+
+ @Override
+ public void writePageV2(
+ int rowCount,
+ int nullCount,
+ int valueCount,
+ BytesInput repetitionLevels,
+ BytesInput definitionLevels,
+ Encoding dataEncoding,
+ BytesInput data,
+ Statistics<?> statistics,
+ SizeStatistics sizeStatistics,
+ GeospatialStatistics geospatialStatistics)
+ throws IOException {
+ pageOrdinal++;
+
+ int rlByteLength = toIntWithCheck(repetitionLevels.size());
+ int dlByteLength = toIntWithCheck(definitionLevels.size());
+ int uncompressedSize =
+ toIntWithCheck(data.size() + repetitionLevels.size() +
definitionLevels.size());
+ boolean compressed = false;
+ BytesInput compressedData = BytesInput.empty();
+ if (data.size() > 0) {
+ compressedData = compressor.compress(data);
+ compressed = true;
+ }
+ if (pageBlockEncryptor != null) {
+ AesCipher.quickUpdatePageAAD(dataPageAAD, pageOrdinal);
+ compressedData =
+ BytesInput.from(
+ pageBlockEncryptor.encrypt(
+ compressedData.toByteArray(),
dataPageAAD));
+ }
+ int compressedSize =
+ toIntWithCheck(
+ compressedData.size()
+ + repetitionLevels.size()
+ + definitionLevels.size());
+ tempOutputStream.reset();
+ if (headerBlockEncryptor != null) {
+ AesCipher.quickUpdatePageAAD(dataPageHeaderAAD, pageOrdinal);
+ }
+ if (pageWriteChecksumEnabled) {
+ crc.reset();
+ if (repetitionLevels.size() > 0) {
+ crc.update(repetitionLevels.toByteArray());
+ }
+ if (definitionLevels.size() > 0) {
+ crc.update(definitionLevels.toByteArray());
+ }
+ if (compressedData.size() > 0) {
+ crc.update(compressedData.toByteArray());
+ }
+ PARQUET_METADATA_CONVERTER.writeDataPageV2Header(
+ uncompressedSize,
+ compressedSize,
+ valueCount,
+ nullCount,
+ rowCount,
+ dataEncoding,
+ rlByteLength,
+ dlByteLength,
+ compressed,
+ (int) crc.getValue(),
+ tempOutputStream,
+ headerBlockEncryptor,
+ dataPageHeaderAAD);
+ } else {
+ PARQUET_METADATA_CONVERTER.writeDataPageV2Header(
+ uncompressedSize,
+ compressedSize,
+ valueCount,
+ nullCount,
+ rowCount,
+ dataEncoding,
+ rlByteLength,
+ dlByteLength,
+ compressed,
+ tempOutputStream,
+ headerBlockEncryptor,
+ dataPageHeaderAAD);
+ }
+ this.uncompressedLength += uncompressedSize;
+ this.compressedLength += compressedSize;
+ this.totalValueCount += valueCount;
+ this.pageCount += 1;
+
+ mergeColumnStatistics(statistics, sizeStatistics,
geospatialStatistics);
+ offsetIndexBuilder.add(
+ toIntWithCheck((long) tempOutputStream.size() +
compressedSize),
+ rowCount,
+ sizeStatistics != null
+ ? sizeStatistics.getUnencodedByteArrayDataBytes()
+ : Optional.empty());
+
+ buf.collect(
+ BytesInput.concat(
+ BytesInput.from(tempOutputStream),
+ repetitionLevels,
+ definitionLevels,
+ compressedData));
+ dataEncodings.add(dataEncoding);
+ }
+
+ private int toIntWithCheck(long size) {
+ if (size > Integer.MAX_VALUE) {
+ throw new ParquetEncodingException(
+ "Cannot write page larger than " + Integer.MAX_VALUE +
" bytes: " + size);
+ }
+ return (int) size;
+ }
+
+ private void mergeColumnStatistics(
+ Statistics<?> statistics,
+ SizeStatistics sizeStatistics,
+ GeospatialStatistics geospatialStatistics) {
+ totalSizeStatistics.mergeStatistics(sizeStatistics);
+ if (!totalSizeStatistics.isValid()) {
+ sizeStatistics = null;
+ }
+
+ totalGeospatialStatistics.merge(geospatialStatistics);
+
+ if (totalStatistics != null && totalStatistics.isEmpty()) {
+ return;
+ }
+
+ if (statistics == null || statistics.isEmpty()) {
+ totalStatistics =
Statistics.getBuilderForReading(path.getPrimitiveType()).build();
+ columnIndexBuilder = ColumnIndexBuilder.getNoOpBuilder();
+ } else if (totalStatistics == null) {
+ totalStatistics = statistics.copy();
+ columnIndexBuilder.add(statistics, sizeStatistics);
+ } else {
+ totalStatistics.mergeStatistics(statistics);
+ columnIndexBuilder.add(statistics, sizeStatistics);
+ }
+ }
+
+ @Override
+ public long getMemSize() {
+ return buf.size();
+ }
+
+ public void writeToFileWriter(ParquetFileWriter writer) throws
IOException {
+ if (headerBlockEncryptor == null) {
+ writer.writeColumnChunk(
+ path,
+ totalValueCount,
+ compressor.getCodecName(),
+ dictionaryPage,
+ buf,
+ uncompressedLength,
+ compressedLength,
+ totalStatistics,
+ totalSizeStatistics,
+ totalGeospatialStatistics,
+ columnIndexBuilder,
+ offsetIndexBuilder,
+ bloomFilter,
+ rlEncodings,
+ dlEncodings,
+ dataEncodings);
+ } else {
+ writer.writeColumnChunk(
+ path,
+ totalValueCount,
+ compressor.getCodecName(),
+ dictionaryPage,
+ buf,
+ uncompressedLength,
+ compressedLength,
+ totalStatistics,
+ totalSizeStatistics,
+ totalGeospatialStatistics,
+ columnIndexBuilder,
+ offsetIndexBuilder,
+ bloomFilter,
+ rlEncodings,
+ dlEncodings,
+ dataEncodings,
+ headerBlockEncryptor,
+ rowGroupOrdinal,
+ columnOrdinal,
+ fileAAD);
+ }
+ if (LOG.isDebugEnabled()) {
+ LOG.debug(
+ String.format(
+ "written %,dB for %s: %,d values, %,dB
raw, %,dB comp, %d pages, encodings: %s",
+ buf.size(),
+ path,
+ totalValueCount,
+ uncompressedLength,
+ compressedLength,
+ pageCount,
+ new HashSet<>(dataEncodings))
+ + (dictionaryPage != null
+ ? String.format(
+ ", dic { %,d entries, %,dB
raw, %,dB comp}",
+
dictionaryPage.getDictionarySize(),
+
dictionaryPage.getUncompressedSize(),
+
dictionaryPage.getDictionarySize())
+ : ""));
+ }
+ rlEncodings.clear();
+ dlEncodings.clear();
+ dataEncodings.clear();
+ pageCount = 0;
+ pageOrdinal = -1;
+ }
+
+ @Override
+ public long allocatedSize() {
+ return buf.size();
+ }
+
+ @Override
+ public void writeDictionaryPage(DictionaryPage dictionaryPage) throws
IOException {
+ if (this.dictionaryPage != null) {
+ throw new ParquetEncodingException("Only one dictionary page
is allowed");
+ }
+ BytesInput dictionaryBytes = dictionaryPage.getBytes();
+ int uncompressedSize = (int) dictionaryBytes.size();
+ BytesInput compressedBytes = compressor.compress(dictionaryBytes);
+ if (pageBlockEncryptor != null) {
+ byte[] dictionaryPageAAD =
+ AesCipher.createModuleAAD(
+ fileAAD,
+ ModuleType.DictionaryPage,
+ rowGroupOrdinal,
+ columnOrdinal,
+ -1);
+ compressedBytes =
+ BytesInput.from(
+ pageBlockEncryptor.encrypt(
+ compressedBytes.toByteArray(),
dictionaryPageAAD));
+ }
+ this.dictionaryPage =
+ new DictionaryPage(
+ compressedBytes.copy(releaser),
+ uncompressedSize,
+ dictionaryPage.getDictionarySize(),
+ dictionaryPage.getEncoding());
+ }
+
+ @Override
+ public String memUsageString(String prefix) {
+ return buf.memUsageString(prefix + " ColumnChunkPageWriter");
+ }
+
+ @Override
+ public void close() {
+ AutoCloseables.uncheckedClose(buf, releaser);
+ }
+
+ @Override
+ public void writeBloomFilter(BloomFilter bloomFilter) {
+ this.bloomFilter = bloomFilter;
+ }
+ }
+
+ private final Map<ColumnDescriptor, ColumnChunkPageWriter> writers = new
HashMap<>();
+ private final MessageType schema;
+
+ public ColumnCompressionPageWriteStore(
+ CompressionCodecFactory codecFactory,
+ ColumnCompressionCodecs compressionCodecs,
+ MessageType schema,
+ ByteBufferAllocator allocator,
+ int columnIndexTruncateLength,
+ boolean pageWriteChecksumEnabled,
+ InternalFileEncryptor fileEncryptor,
+ int rowGroupOrdinal) {
+ this.schema = schema;
+ if (fileEncryptor == null) {
+ for (ColumnDescriptor path : schema.getColumns()) {
+ writers.put(
+ path,
+ new ColumnChunkPageWriter(
+ path,
+
codecFactory.getCompressor(compressionCodecs.getCodec(path)),
+ allocator,
+ columnIndexTruncateLength,
+ pageWriteChecksumEnabled,
+ null,
+ null,
+ null,
+ -1,
+ -1));
+ }
+ return;
+ }
+
+ int columnOrdinal = -1;
+ byte[] fileAAD = fileEncryptor.getFileAAD();
+ for (ColumnDescriptor path : schema.getColumns()) {
+ columnOrdinal++;
+ BlockCipher.Encryptor headerBlockEncryptor = null;
+ BlockCipher.Encryptor pageBlockEncryptor = null;
+ ColumnPath columnPath = ColumnPath.get(path.getPath());
+
+ InternalColumnEncryptionSetup columnSetup =
+ fileEncryptor.getColumnSetup(columnPath, true,
columnOrdinal);
+ if (columnSetup.isEncrypted()) {
+ headerBlockEncryptor = columnSetup.getMetaDataEncryptor();
+ pageBlockEncryptor = columnSetup.getDataEncryptor();
+ }
+
+ writers.put(
+ path,
+ new ColumnChunkPageWriter(
+ path,
+
codecFactory.getCompressor(compressionCodecs.getCodec(path)),
+ allocator,
+ columnIndexTruncateLength,
+ pageWriteChecksumEnabled,
+ headerBlockEncryptor,
+ pageBlockEncryptor,
+ fileAAD,
+ rowGroupOrdinal,
+ columnOrdinal));
+ }
+ }
+
+ @Override
+ public PageWriter getPageWriter(ColumnDescriptor path) {
+ return writers.get(path);
+ }
+
+ @Override
+ public void close() {
+ AutoCloseables.uncheckedClose(writers.values());
+ writers.clear();
+ }
+
+ @Override
+ public BloomFilterWriter getBloomFilterWriter(ColumnDescriptor path) {
+ return writers.get(path);
+ }
+
+ public void flushToFileWriter(ParquetFileWriter writer) throws IOException
{
+ for (ColumnDescriptor path : schema.getColumns()) {
+ ColumnChunkPageWriter pageWriter = writers.get(path);
+ pageWriter.writeToFileWriter(writer);
+ }
+ }
+}
diff --git
a/paimon-format/src/main/java/org/apache/parquet/hadoop/ColumnCompressionRecordWriter.java
b/paimon-format/src/main/java/org/apache/parquet/hadoop/ColumnCompressionRecordWriter.java
new file mode 100644
index 0000000000..d8afcf4a43
--- /dev/null
+++
b/paimon-format/src/main/java/org/apache/parquet/hadoop/ColumnCompressionRecordWriter.java
@@ -0,0 +1,242 @@
+/*
+ * 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.hadoop;
+
+import org.apache.parquet.column.ColumnWriteStore;
+import org.apache.parquet.column.ParquetProperties;
+import org.apache.parquet.column.values.bloomfilter.BloomFilterWriteStore;
+import org.apache.parquet.compression.CompressionCodecFactory;
+import org.apache.parquet.crypto.InternalFileEncryptor;
+import org.apache.parquet.hadoop.api.WriteSupport;
+import org.apache.parquet.hadoop.api.WriteSupport.FinalizedWriteContext;
+import org.apache.parquet.hadoop.metadata.ParquetMetadata;
+import org.apache.parquet.io.ColumnIOFactory;
+import org.apache.parquet.io.MessageColumnIO;
+import org.apache.parquet.io.api.RecordConsumer;
+import org.apache.parquet.schema.MessageType;
+import org.apache.parquet.util.AutoCloseables;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+
+import static java.lang.Math.max;
+import static java.lang.Math.min;
+
+/**
+ * Record writer with per-column compression.
+ *
+ * <p>NOTE: The file is adapted from Apache parquet's
InternalParquetRecordWriter.
+ */
+class ColumnCompressionRecordWriter<T> {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(ColumnCompressionRecordWriter.class);
+
+ private final ParquetFileWriter parquetFileWriter;
+ private final WriteSupport<T> writeSupport;
+ private final MessageType schema;
+ private final Map<String, String> extraMetaData;
+ private long rowGroupSizeThreshold;
+ private final int rowGroupRecordCountThreshold;
+ private long nextRowGroupSize;
+ private final CompressionCodecFactory codecFactory;
+ private final ColumnCompressionCodecs compressionCodecs;
+ private final boolean validating;
+ private final ParquetProperties props;
+
+ private boolean closed;
+
+ private long recordCount = 0;
+ private long recordCountForNextMemCheck;
+ private long lastRowGroupEndPos = 0;
+
+ private ColumnWriteStore columnStore;
+ private ColumnCompressionPageWriteStore pageStore;
+ private BloomFilterWriteStore bloomFilterWriteStore;
+ private RecordConsumer recordConsumer;
+
+ private InternalFileEncryptor fileEncryptor;
+ private int rowGroupOrdinal;
+ private boolean aborted;
+
+ ColumnCompressionRecordWriter(
+ ParquetFileWriter parquetFileWriter,
+ WriteSupport<T> writeSupport,
+ MessageType schema,
+ Map<String, String> extraMetaData,
+ long rowGroupSize,
+ CompressionCodecFactory codecFactory,
+ ColumnCompressionCodecs compressionCodecs,
+ boolean validating,
+ ParquetProperties props) {
+ this.parquetFileWriter = parquetFileWriter;
+ this.writeSupport = Objects.requireNonNull(writeSupport, "writeSupport
cannot be null");
+ this.schema = schema;
+ this.extraMetaData = extraMetaData;
+ this.rowGroupSizeThreshold = rowGroupSize;
+ this.rowGroupRecordCountThreshold = props.getRowGroupRowCountLimit();
+ this.nextRowGroupSize = rowGroupSizeThreshold;
+ this.codecFactory = codecFactory;
+ this.compressionCodecs = compressionCodecs;
+ this.validating = validating;
+ this.props = props;
+ this.fileEncryptor = parquetFileWriter.getEncryptor();
+ this.rowGroupOrdinal = 0;
+ initStore();
+ recordCountForNextMemCheck = props.getMinRowCountForPageSizeCheck();
+ }
+
+ public ParquetMetadata getFooter() {
+ return parquetFileWriter.getFooter();
+ }
+
+ private void initStore() {
+ ColumnCompressionPageWriteStore columnCompressionPageWriteStore =
+ new ColumnCompressionPageWriteStore(
+ codecFactory,
+ compressionCodecs,
+ schema,
+ props.getAllocator(),
+ props.getColumnIndexTruncateLength(),
+ props.getPageWriteChecksumEnabled(),
+ fileEncryptor,
+ rowGroupOrdinal);
+ pageStore = columnCompressionPageWriteStore;
+ bloomFilterWriteStore = columnCompressionPageWriteStore;
+
+ columnStore = props.newColumnWriteStore(schema, pageStore,
bloomFilterWriteStore);
+ MessageColumnIO columnIO = new
ColumnIOFactory(validating).getColumnIO(schema);
+ this.recordConsumer = columnIO.getRecordWriter(columnStore);
+ writeSupport.prepareForWrite(recordConsumer);
+ }
+
+ public void close() throws IOException, InterruptedException {
+ if (!closed) {
+ try {
+ if (aborted) {
+ return;
+ }
+ flushRowGroupToStore();
+ FinalizedWriteContext finalWriteContext =
writeSupport.finalizeWrite();
+ Map<String, String> finalMetadata = new
HashMap<>(extraMetaData);
+ String modelName = writeSupport.getName();
+ if (modelName != null) {
+ finalMetadata.put(ParquetWriter.OBJECT_MODEL_NAME_PROP,
modelName);
+ }
+ finalMetadata.putAll(finalWriteContext.getExtraMetaData());
+ parquetFileWriter.end(finalMetadata);
+ } finally {
+ AutoCloseables.uncheckedClose(
+ columnStore, pageStore, bloomFilterWriteStore,
parquetFileWriter);
+ closed = true;
+ }
+ }
+ }
+
+ public void write(T value) throws IOException, InterruptedException {
+ try {
+ writeSupport.write(value);
+ ++recordCount;
+ checkBlockSizeReached();
+ } catch (Throwable t) {
+ aborted = true;
+ throw t;
+ }
+ }
+
+ public long getDataSize() {
+ return lastRowGroupEndPos + columnStore.getBufferedSize();
+ }
+
+ private void checkBlockSizeReached() throws IOException {
+ if (recordCount >= rowGroupRecordCountThreshold) {
+ LOG.debug("record count reaches threshold: flushing {} records to
disk.", recordCount);
+ flushRowGroupToStore();
+ initStore();
+ recordCountForNextMemCheck =
+ min(
+ max(props.getMinRowCountForPageSizeCheck(),
recordCount / 2),
+ props.getMaxRowCountForPageSizeCheck());
+ this.lastRowGroupEndPos = parquetFileWriter.getPos();
+ } else if (recordCount >= recordCountForNextMemCheck) {
+ long memSize = columnStore.getBufferedSize();
+ long recordSize = memSize / recordCount;
+ if (memSize > (nextRowGroupSize - 2 * recordSize)) {
+ LOG.debug(
+ "mem size {} > {}: flushing {} records to disk.",
+ memSize,
+ nextRowGroupSize,
+ recordCount);
+ flushRowGroupToStore();
+ initStore();
+ recordCountForNextMemCheck =
+ min(
+ max(props.getMinRowCountForPageSizeCheck(),
recordCount / 2),
+ props.getMaxRowCountForPageSizeCheck());
+ this.lastRowGroupEndPos = parquetFileWriter.getPos();
+ } else {
+ recordCountForNextMemCheck =
+ min(
+ max(
+ props.getMinRowCountForPageSizeCheck(),
+ (recordCount
+ + (long)
+
(nextRowGroupSize
+ /
((float) recordSize)))
+ / 2),
+ recordCount +
props.getMaxRowCountForPageSizeCheck());
+ LOG.debug(
+ "Checked mem at {} will check again at: {}",
+ recordCount,
+ recordCountForNextMemCheck);
+ }
+ }
+ }
+
+ private void flushRowGroupToStore() throws IOException {
+ try {
+ recordConsumer.flush();
+ LOG.debug(
+ "Flushing mem columnStore to file. allocated memory: {}",
+ columnStore.getAllocatedSize());
+ if (columnStore.getAllocatedSize() > (3 * rowGroupSizeThreshold)) {
+ LOG.warn("Too much memory used: {}",
columnStore.memUsageString());
+ }
+
+ if (recordCount > 0) {
+ rowGroupOrdinal++;
+ parquetFileWriter.startBlock(recordCount);
+ columnStore.flush();
+ pageStore.flushToFileWriter(parquetFileWriter);
+ recordCount = 0;
+ parquetFileWriter.endBlock();
+ this.nextRowGroupSize =
+ Math.min(parquetFileWriter.getNextRowGroupSize(),
rowGroupSizeThreshold);
+ }
+ } finally {
+ AutoCloseables.uncheckedClose(columnStore, pageStore,
bloomFilterWriteStore);
+ columnStore = null;
+ pageStore = null;
+ bloomFilterWriteStore = null;
+ }
+ }
+}
diff --git
a/paimon-format/src/main/java/org/apache/parquet/hadoop/ParquetWriter.java
b/paimon-format/src/main/java/org/apache/parquet/hadoop/ParquetWriter.java
index 3358e84fdc..ee092bbfb9 100644
--- a/paimon-format/src/main/java/org/apache/parquet/hadoop/ParquetWriter.java
+++ b/paimon-format/src/main/java/org/apache/parquet/hadoop/ParquetWriter.java
@@ -26,6 +26,7 @@ import
org.apache.parquet.column.ParquetProperties.WriterVersion;
import org.apache.parquet.compression.CompressionCodecFactory;
import org.apache.parquet.crypto.FileEncryptionProperties;
import org.apache.parquet.hadoop.api.WriteSupport;
+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.io.OutputFile;
@@ -58,7 +59,7 @@ public class ParquetWriter<T> implements Closeable {
// max size (bytes) to write as padding and the min size of a row group
public static final int MAX_PADDING_SIZE_DEFAULT = 8 * 1024 * 1024; // 8MB
- private final InternalParquetRecordWriter<T> writer;
+ private final ColumnCompressionRecordWriter<T> writer;
private final CompressionCodecFactory codecFactory;
ParquetWriter(
@@ -67,6 +68,7 @@ public class ParquetWriter<T> implements Closeable {
WriteSupport<T> writeSupport,
CompressionCodecName compressionCodecName,
CompressionCodecFactory codecFactory,
+ Map<ColumnPath, CompressionCodecName> columnCompressionCodecNames,
long rowGroupSize,
boolean validating,
Configuration conf,
@@ -99,8 +101,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
@@ -130,13 +130,15 @@ public class ParquetWriter<T> implements Closeable {
}
this.writer =
- new InternalParquetRecordWriter<T>(
+ new ColumnCompressionRecordWriter<T>(
fileWriter,
writeSupport,
schema,
extraMetadata,
rowGroupSize,
- compressor,
+ codecFactory,
+ new ColumnCompressionCodecs(
+ compressionCodecName,
columnCompressionCodecNames),
validating,
encodingProps);
}
@@ -188,6 +190,7 @@ public class ParquetWriter<T> implements Closeable {
private ParquetFileWriter.Mode mode;
private CompressionCodecFactory codecFactory = null;
private CompressionCodecName codecName =
DEFAULT_COMPRESSION_CODEC_NAME;
+ private final Map<ColumnPath, CompressionCodecName> columnCodecNames =
new HashMap<>();
private long rowGroupSize = DEFAULT_BLOCK_SIZE;
private int maxPaddingSize = MAX_PADDING_SIZE_DEFAULT;
private boolean enableValidation = DEFAULT_IS_VALIDATING_ENABLED;
@@ -240,6 +243,18 @@ public class ParquetWriter<T> implements Closeable {
return self();
}
+ /**
+ * Set the compression codec used by the specified column for the
constructed writer.
+ *
+ * @param columnPath the path of the column (dot-string)
+ * @param codecName a {@code CompressionCodecName}
+ * @return this builder for method chaining.
+ */
+ public SELF withCompressionCodec(String columnPath,
CompressionCodecName codecName) {
+ this.columnCodecNames.put(ColumnPath.fromDotString(columnPath),
codecName);
+ return self();
+ }
+
/**
* Set the {@link CompressionCodecFactory codec factory} used by the
constructed writer.
*
@@ -658,6 +673,7 @@ public class ParquetWriter<T> implements Closeable {
getWriteSupport(conf),
codecName,
codecFactory,
+ columnCodecNames,
rowGroupSize,
enableValidation,
conf,
diff --git
a/paimon-format/src/test/java/org/apache/paimon/format/parquet/ParquetFormatReadWriteTest.java
b/paimon-format/src/test/java/org/apache/paimon/format/parquet/ParquetFormatReadWriteTest.java
index 0f71fbbcd7..137da29aa6 100644
---
a/paimon-format/src/test/java/org/apache/paimon/format/parquet/ParquetFormatReadWriteTest.java
+++
b/paimon-format/src/test/java/org/apache/paimon/format/parquet/ParquetFormatReadWriteTest.java
@@ -18,6 +18,7 @@
package org.apache.paimon.format.parquet;
+import org.apache.paimon.data.BinaryString;
import org.apache.paimon.data.GenericRow;
import org.apache.paimon.format.FileFormat;
import org.apache.paimon.format.FileFormatFactory;
@@ -32,12 +33,16 @@ import
org.apache.parquet.column.values.bloomfilter.BloomFilter;
import org.apache.parquet.hadoop.ParquetFileReader;
import org.apache.parquet.hadoop.metadata.BlockMetaData;
import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
+import org.apache.parquet.hadoop.metadata.CompressionCodecName;
import org.apache.parquet.hadoop.metadata.ParquetMetadata;
import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
import java.util.concurrent.ThreadLocalRandom;
/** A parquet {@link FormatReadWriteTest}. */
@@ -89,4 +94,42 @@ public class ParquetFormatReadWriteTest extends
FormatReadWriteTest {
}
}
}
+
+ @Test
+ public void testColumnCompressionCodec() throws Exception {
+ Options options = new Options();
+ options.set("parquet.compression#name", "none");
+ ParquetFileFormat format =
+ new ParquetFileFormat(new
FileFormatFactory.FormatContext(options, 1024, 1024));
+
+ RowType rowType =
+ DataTypes.ROW(
+ DataTypes.FIELD(0, "id", DataTypes.INT()),
+ DataTypes.FIELD(1, "name", DataTypes.STRING()));
+
+ PositionOutputStream out = fileIO.newOutputStream(file, false);
+ FormatWriter writer = format.createWriterFactory(rowType).create(out,
"zstd");
+ writer.addElement(GenericRow.of(1, BinaryString.fromString("one")));
+ writer.addElement(GenericRow.of(2, BinaryString.fromString("two")));
+ writer.addElement(GenericRow.of(3, BinaryString.fromString("three")));
+ writer.close();
+ out.close();
+
+ try (ParquetFileReader reader =
+ ParquetUtil.getParquetReader(
+ fileIO, file, fileIO.getFileSize(file), new
Options())) {
+ Map<String, CompressionCodecName> codecs = new HashMap<>();
+ for (BlockMetaData blockMetaData : reader.getFooter().getBlocks())
{
+ for (ColumnChunkMetaData columnChunkMetaData :
blockMetaData.getColumns()) {
+ codecs.put(
+ columnChunkMetaData.getPath().toDotString(),
+ columnChunkMetaData.getCodec());
+ }
+ }
+
+ Assertions.assertThat(codecs)
+ .containsEntry("id", CompressionCodecName.ZSTD)
+ .containsEntry("name", CompressionCodecName.UNCOMPRESSED);
+ }
+ }
}
