diff --git a/gluten-flink/planner/src/main/java/org/apache/gluten/velox/HiveSourceSinkFactory.java b/gluten-flink/planner/src/main/java/org/apache/gluten/velox/HiveSourceSinkFactory.java index 56d356f237a..defc1025df6 100644 --- a/gluten-flink/planner/src/main/java/org/apache/gluten/velox/HiveSourceSinkFactory.java +++ b/gluten-flink/planner/src/main/java/org/apache/gluten/velox/HiveSourceSinkFactory.java @@ -25,8 +25,13 @@ import java.util.HashMap; import java.util.Map; +import java.util.Properties; public class HiveSourceSinkFactory extends FileSystemSinkFactory { + private static final String COMPRESSION_KIND = "sink.file.compression"; + private static final String[] SUPPORTED_COMPRESSION_TABLE_KEYS = { + "orc.compress", "parquet.compression", "parquet.compression.codec", "parquet.compression-codec" + }; @Override public boolean match(Transformation transformation) { @@ -42,8 +47,10 @@ protected Map buildTableParams( Configuration tableOptions = getTableOptions(partitionCommitter, fileWriterOperator); Map tableParams = new HashMap<>(tableOptions.toMap()); tableParams.put("path", getLocationPath(partitionCommitter, fileWriterOperator)); - tableParams.putIfAbsent("format", resolveWriteFormat(fileWriterOperator)); + Object bucketsBuilder = getBucketsBuilder(fileWriterOperator); + tableParams.putIfAbsent("format", resolveWriteFormat(bucketsBuilder)); tableParams.put("connector", "hive"); + addHiveCompressionParams(bucketsBuilder, tableParams); return tableParams; } @@ -57,15 +64,21 @@ protected String getDefaultFormat() { return "hive"; } + private String resolveWriteFormat(Object bucketsBuilder) { + String format = resolveFormatFromBucketsBuilder(bucketsBuilder); + if (format != null) { + return format; + } + return getDefaultFormat(); + } + @Override - protected String resolveFormatFromHadoopBulkWriterFactory(Object writerFactory) { - Class factoryClass = writerFactory.getClass(); - if (factoryClass.getName().contains("HiveBulkWriterFactory")) { - Object hiveWriterFactory = - ReflectUtils.getObjectField(factoryClass, writerFactory, "factory"); + protected String resolveFormatFromBucketsBuilder(Object bucketsBuilder) { + Object hiveWriterFactory = getHiveWriterFactoryFromBucketsBuilder(bucketsBuilder); + if (hiveWriterFactory != null) { return resolveFormatFromHiveWriterFactory(hiveWriterFactory); } - return super.resolveFormatFromHadoopBulkWriterFactory(writerFactory); + return super.resolveFormatFromBucketsBuilder(bucketsBuilder); } private String resolveFormatFromHiveWriterFactory(Object hiveWriterFactory) { @@ -96,4 +109,87 @@ private String resolveFormatFromHiveWriterFactory(Object hiveWriterFactory) { ReflectUtils.getObjectField(factoryClass, hiveWriterFactory, "hiveOutputFormatClz"); return inferFormatFromClassName(outputFormatClz.getName()); } + + private void addHiveCompressionParams(Object bucketsBuilder, Map tableParams) { + Object hiveWriterFactory = getHiveWriterFactoryFromBucketsBuilder(bucketsBuilder); + if (hiveWriterFactory == null) { + return; + } + Properties tableProperties = + (Properties) ReflectUtils.tryGetObjectField(hiveWriterFactory, "tableProperties"); + addNativeCompressionParamFromTableProperties(tableProperties, tableParams); + } + + private Object getBucketsBuilder(OneInputStreamOperator fileWriterOperator) { + return ReflectUtils.getObjectField( + ABSTRACT_STREAMING_WRITER_CLASS, fileWriterOperator, "bucketsBuilder"); + } + + private Object getHiveWriterFactoryFromBucketsBuilder(Object bucketsBuilder) { + Object writerFactory = ReflectUtils.tryGetObjectField(bucketsBuilder, "writerFactory"); + if (writerFactory == null + || !writerFactory.getClass().getName().contains("HiveBulkWriterFactory")) { + return null; + } + return ReflectUtils.getObjectField(writerFactory.getClass(), writerFactory, "factory"); + } + + static void addNativeCompressionParamFromTableProperties( + Properties tableProperties, Map tableParams) { + if (!isParquetFormat(tableParams.get("format"))) { + tableParams.remove(COMPRESSION_KIND); + return; + } + if (tableProperties == null) { + return; + } + + String compressionKind = resolveCompressionKind(tableProperties); + if (compressionKind != null) { + tableParams.put(COMPRESSION_KIND, compressionKind); + } + } + + private static boolean isParquetFormat(String format) { + return format != null && "parquet".equalsIgnoreCase(format.trim()); + } + + private static String resolveCompressionKind(Properties tableProperties) { + for (String key : SUPPORTED_COMPRESSION_TABLE_KEYS) { + String compressionKind = normalizeCompressionKind(tableProperties.getProperty(key)); + if (compressionKind != null) { + return compressionKind; + } + } + return null; + } + + static String normalizeCompressionKind(String compression) { + if (compression == null) { + return null; + } + switch (compression.trim().toLowerCase()) { + case "snappy": + return "snappy"; + case "gzip": + return "gzip"; + case "zstd": + case "zstandard": + return "zstd"; + case "lz4": + return "lz4"; + case "lzo": + return "lzo"; + case "zlib": + case "deflate": + return "zlib"; + case "": + case "none": + case "no": + case "false": + case "uncompressed": + default: + return null; + } + } } diff --git a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperator.java b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperator.java index 833b2c452a4..0f4eee8411d 100644 --- a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperator.java +++ b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperator.java @@ -95,11 +95,15 @@ protected void initializeNativeState(StateInitializationContext context) throws } restoredCheckpointRecords = records.toArray(new String[0]); } - LOG.info( - "Restore native file writer state for operator {}, restored {}, records {}", - getDescription(), - context.isRestored(), - Arrays.toString(restoredCheckpointRecords)); + if (restoredCheckpointRecords != null) { + LOG.info( + "Restore native file writer state for operator {}, restored {}, records {}, contents {}", + getDescription(), + context.isRestored(), + restoredCheckpointRecords.length, + Arrays.toString( + restoredCheckpointRecords != null ? restoredCheckpointRecords : new String[0])); + } if (task == null) { initSession(); } diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/velox/HiveSourceSinkFactoryTest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/velox/HiveSourceSinkFactoryTest.java new file mode 100644 index 00000000000..1b1dabfbe77 --- /dev/null +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/velox/HiveSourceSinkFactoryTest.java @@ -0,0 +1,118 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.gluten.velox; + +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.Map; +import java.util.Properties; + +import static org.assertj.core.api.Assertions.assertThat; + +class HiveSourceSinkFactoryTest { + + @Test + void addNativeCompressionParamMapsSupportedParquetCodecs() { + String[][] compressionCodecs = { + {"SNAPPY", "snappy"}, + {"GZIP", "gzip"}, + {"zstandard", "zstd"}, + {"LZ4", "lz4"}, + {"LZO", "lzo"}, + {"deflate", "zlib"} + }; + + for (String[] compressionCodec : compressionCodecs) { + Properties tableProperties = new Properties(); + tableProperties.setProperty("parquet.compression", compressionCodec[0]); + + Map tableParams = new HashMap<>(); + tableParams.put("format", "parquet"); + HiveSourceSinkFactory.addNativeCompressionParamFromTableProperties( + tableProperties, tableParams); + + assertThat(tableParams) + .containsEntry("sink.file.compression", compressionCodec[1]) + .doesNotContainKey("parquet.compression"); + } + } + + @Test + void addNativeCompressionParamReadsSupportedParquetCompressionKeys() { + Properties tableProperties = new Properties(); + tableProperties.setProperty("parquet.compression.codec", "SNAPPY"); + + Map tableParams = new HashMap<>(); + tableParams.put("format", "parquet"); + HiveSourceSinkFactory.addNativeCompressionParamFromTableProperties( + tableProperties, tableParams); + + assertThat(tableParams).containsEntry("sink.file.compression", "snappy"); + } + + @Test + void addNativeCompressionParamDoesNotProduceConfigForUnsupportedFormats() { + for (String format : new String[] {"orc", "json", "csv", "hive"}) { + Properties tableProperties = new Properties(); + tableProperties.setProperty("parquet.compression", "SNAPPY"); + + Map tableParams = new HashMap<>(); + tableParams.put("format", format); + tableParams.put("sink.file.compression", "snappy"); + HiveSourceSinkFactory.addNativeCompressionParamFromTableProperties( + tableProperties, tableParams); + + assertThat(tableParams) + .containsEntry("format", format) + .doesNotContainKey("sink.file.compression"); + } + } + + @Test + void addNativeCompressionParamDoesNotProduceConfigForUnsupportedParquetCodecs() { + for (String compressionCodec : new String[] {"brotli", "org.example.SnappyCodec"}) { + Properties tableProperties = new Properties(); + tableProperties.setProperty("parquet.compression", compressionCodec); + + Map tableParams = new HashMap<>(); + tableParams.put("format", "parquet"); + HiveSourceSinkFactory.addNativeCompressionParamFromTableProperties( + tableProperties, tableParams); + + assertThat(tableParams) + .containsEntry("format", "parquet") + .doesNotContainKey("sink.file.compression"); + } + } + + @Test + void addNativeCompressionParamDoesNotProduceConfigForUnsupportedCompressionKeys() { + Properties tableProperties = new Properties(); + tableProperties.setProperty("custom.compress", "SNAPPY"); + tableProperties.setProperty("custom.codec", "GZIP"); + + Map tableParams = new HashMap<>(); + tableParams.put("format", "parquet"); + HiveSourceSinkFactory.addNativeCompressionParamFromTableProperties( + tableProperties, tableParams); + + assertThat(tableParams) + .containsEntry("format", "parquet") + .doesNotContainKey("sink.file.compression"); + } +}