From 885795c45c371db1faa24b36742889df352c5af6 Mon Sep 17 00:00:00 2001 From: zouyunhe Date: Mon, 13 Jul 2026 10:09:59 +0000 Subject: [PATCH 1/3] [GLUTEN-12206][FLINK] Support filesystem writer bucket state snapshot and restore --- .../operators/GlutenOneInputOperator.java | 13 +- .../GlutenStreamingFileWriterOperator.java | 54 ++++++++ ...GlutenStreamingFileWriterOperatorTest.java | 123 ++++++++++++++++++ 3 files changed, 187 insertions(+), 3 deletions(-) create mode 100644 gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperatorTest.java diff --git a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java index ba9329e67b8..eedc7c0e5c7 100644 --- a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java +++ b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java @@ -331,20 +331,27 @@ public void prepareSnapshotPreBarrier(long checkpointId) throws Exception { @Override public void snapshotState(StateSnapshotContext context) throws Exception { - // TODO: implement it - task.snapshotState(context.getCheckpointId()); + snapshotNativeState(context); super.snapshotState(context); } + protected void snapshotNativeState(StateSnapshotContext context) throws Exception { + task.snapshotState(context.getCheckpointId()); + } + @Override public void initializeState(StateInitializationContext context) throws Exception { + initializeNativeState(context); + super.initializeState(context); + } + + protected void initializeNativeState(StateInitializationContext context) throws Exception { if (task == null) { initSession(); } if (!(getKeyedStateBackend() instanceof RocksDBKeyedStateBackend)) { task.initializeState(0, null); } - super.initializeState(context); } @Override 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 e508fb80c32..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 @@ -22,15 +22,31 @@ import io.github.zhztheplayer.velox4j.type.RowType; import org.apache.flink.api.common.TaskInfo; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; import org.apache.flink.connector.file.table.stream.PartitionCommitInfo; +import org.apache.flink.runtime.state.StateInitializationContext; +import org.apache.flink.runtime.state.StateSnapshotContext; import org.apache.flink.streaming.api.watermark.Watermark; import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; import java.util.Map; public class GlutenStreamingFileWriterOperator extends GlutenOneInputOperator { + private static final Logger LOG = + LoggerFactory.getLogger(GlutenStreamingFileWriterOperator.class); + + private transient ListState checkpointState; + private transient String[] restoredCheckpointRecords = new String[0]; + public GlutenStreamingFileWriterOperator( StatefulPlanNode plan, String id, @@ -56,6 +72,44 @@ public void processWatermark(Watermark mark) throws Exception { output.emitWatermark(mark); } + @Override + protected void snapshotNativeState(StateSnapshotContext context) throws Exception { + checkpointState.clear(); + String[] checkpointRecords = task.snapshotState(context.getCheckpointId()); + for (String checkpointRecord : checkpointRecords) { + checkpointState.add(checkpointRecord); + } + } + + @Override + protected void initializeNativeState(StateInitializationContext context) throws Exception { + checkpointState = + context + .getOperatorStateStore() + .getListState( + new ListStateDescriptor<>("gluten-native-file-writer-checkpoint", String.class)); + if (context.isRestored()) { + List records = new ArrayList<>(); + for (String checkpointRecord : checkpointState.get()) { + records.add(checkpointRecord); + } + restoredCheckpointRecords = records.toArray(new String[0]); + } + 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(); + } + task.initializeState(0, null, restoredCheckpointRecords); + } + @Override public void notifyCheckpointComplete(long checkpointId) throws Exception { String[] committed = task.notifyCheckpointComplete(checkpointId); diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperatorTest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperatorTest.java new file mode 100644 index 00000000000..6f3b0145c11 --- /dev/null +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenStreamingFileWriterOperatorTest.java @@ -0,0 +1,123 @@ +/* + * 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.table.runtime.operators; + +import io.github.zhztheplayer.velox4j.plan.EmptyNode; +import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode; +import io.github.zhztheplayer.velox4j.query.SerialTask; +import io.github.zhztheplayer.velox4j.stateful.KeyedStateBackendParameters; +import io.github.zhztheplayer.velox4j.type.IntegerType; +import io.github.zhztheplayer.velox4j.type.RowType; + +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; +import org.apache.flink.runtime.state.StateInitializationContext; +import org.apache.flink.runtime.state.StateInitializationContextImpl; +import org.apache.flink.runtime.state.StateSnapshotContextSynchronousImpl; +import org.apache.flink.streaming.api.operators.collect.utils.MockOperatorStateStore; +import org.apache.flink.table.data.RowData; + +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +class GlutenStreamingFileWriterOperatorTest { + private static final ListStateDescriptor CHECKPOINT_STATE_DESCRIPTOR = + new ListStateDescriptor<>("gluten-native-file-writer-checkpoint", String.class); + + @Test + void snapshotNativeStateStoresNativeCheckpointRecords() throws Exception { + MockOperatorStateStore stateStore = new MockOperatorStateStore(); + TestSerialTask task = new TestSerialTask(new String[] {"bucket-a", "bucket-b"}); + GlutenStreamingFileWriterOperator operator = createOperator(task); + + operator.initializeNativeState(createInitializationContext(null, stateStore)); + operator.snapshotNativeState(new StateSnapshotContextSynchronousImpl(42L, 0L)); + + ListState checkpointState = stateStore.getListState(CHECKPOINT_STATE_DESCRIPTOR); + assertThat(task.snapshotCheckpointId).isEqualTo(42L); + assertThat(checkpointState.get()).containsExactly("bucket-a", "bucket-b"); + } + + @Test + void initializeNativeStateRestoresCheckpointRecordsToNativeTask() throws Exception { + MockOperatorStateStore stateStore = new MockOperatorStateStore(); + ListState checkpointState = stateStore.getListState(CHECKPOINT_STATE_DESCRIPTOR); + checkpointState.add("bucket-a"); + checkpointState.add("bucket-b"); + TestSerialTask task = new TestSerialTask(new String[0]); + GlutenStreamingFileWriterOperator operator = createOperator(task); + + operator.initializeNativeState(createInitializationContext(1L, stateStore)); + + assertThat(task.initializeContext).isZero(); + assertThat(task.initializeParameters).isNull(); + assertThat(task.restoredCheckpointRecords).containsExactly("bucket-a", "bucket-b"); + } + + private static StateInitializationContext createInitializationContext( + Long restoredCheckpointId, MockOperatorStateStore stateStore) { + return new StateInitializationContextImpl( + restoredCheckpointId, stateStore, null, Collections.emptyList(), Collections.emptyList()); + } + + private static GlutenStreamingFileWriterOperator createOperator(TestSerialTask task) { + RowType rowType = new RowType(List.of("id"), List.of(new IntegerType())); + StatefulPlanNode plan = new StatefulPlanNode("scan", new EmptyNode(rowType)); + GlutenStreamingFileWriterOperator operator = + new GlutenStreamingFileWriterOperator<>( + plan, + "scan", + rowType, + Map.of("scan", rowType), + RowData.class, + "test native file writer"); + operator.task = task; + return operator; + } + + private static class TestSerialTask extends SerialTask { + private final String[] checkpointRecords; + private long snapshotCheckpointId = -1L; + private long initializeContext = -1L; + private KeyedStateBackendParameters initializeParameters; + private String[] restoredCheckpointRecords; + + private TestSerialTask(String[] checkpointRecords) { + super(null, 0L); + this.checkpointRecords = checkpointRecords; + } + + @Override + public String[] snapshotState(long checkpointId) { + snapshotCheckpointId = checkpointId; + return checkpointRecords; + } + + @Override + public void initializeState( + long context, KeyedStateBackendParameters parameters, String[] checkpointRecords) { + initializeContext = context; + initializeParameters = parameters; + restoredCheckpointRecords = checkpointRecords; + } + } +} From 40a3616700519317ef3c195b9dc5494b58f31714 Mon Sep 17 00:00:00 2001 From: zouyunhe Date: Tue, 14 Jul 2026 10:19:15 +0000 Subject: [PATCH 2/3] Support compressin type of hive table --- .../gluten/velox/HiveSourceSinkFactory.java | 96 +++++++++++++++++++ .../velox/HiveSourceSinkFactoryTest.java | 49 ++++++++++ 2 files changed, 145 insertions(+) create mode 100644 gluten-flink/ut/src/test/java/org/apache/gluten/velox/HiveSourceSinkFactoryTest.java 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..b030dd4f519 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[] TABLE_COMPRESSION_KEYS = { + "orc.compress", "parquet.compression", "parquet.compression.codec", "parquet.compression-codec" + }; @Override public boolean match(Transformation transformation) { @@ -44,6 +49,7 @@ protected Map buildTableParams( tableParams.put("path", getLocationPath(partitionCommitter, fileWriterOperator)); tableParams.putIfAbsent("format", resolveWriteFormat(fileWriterOperator)); tableParams.put("connector", "hive"); + addHiveCompressionParams(fileWriterOperator, tableParams); return tableParams; } @@ -96,4 +102,94 @@ private String resolveFormatFromHiveWriterFactory(Object hiveWriterFactory) { ReflectUtils.getObjectField(factoryClass, hiveWriterFactory, "hiveOutputFormatClz"); return inferFormatFromClassName(outputFormatClz.getName()); } + + private void addHiveCompressionParams( + OneInputStreamOperator fileWriterOperator, Map tableParams) { + Object hiveWriterFactory = getHiveWriterFactory(fileWriterOperator); + if (hiveWriterFactory == null) { + return; + } + Properties tableProperties = + (Properties) ReflectUtils.tryGetObjectField(hiveWriterFactory, "tableProperties"); + addCompressionParamsFromTableProperties(tableProperties, tableParams); + } + + private Object getHiveWriterFactory(OneInputStreamOperator fileWriterOperator) { + Object bucketsBuilder = + ReflectUtils.getObjectField( + ABSTRACT_STREAMING_WRITER_CLASS, fileWriterOperator, "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 addCompressionParamsFromTableProperties( + Properties tableProperties, Map tableParams) { + if (tableProperties == null) { + return; + } + for (String key : tableProperties.stringPropertyNames()) { + String value = tableProperties.getProperty(key); + if (isCompressionProperty(key, value)) { + tableParams.putIfAbsent(key, value); + } + } + + String compressionKind = resolveCompressionKind(tableProperties); + if (compressionKind != null) { + tableParams.put(COMPRESSION_KIND, compressionKind); + } + } + + private static String resolveCompressionKind(Properties tableProperties) { + for (String key : TABLE_COMPRESSION_KEYS) { + String compressionKind = normalizeCompressionKind(tableProperties.getProperty(key)); + if (compressionKind != null) { + return compressionKind; + } + } + return null; + } + + private static boolean isCompressionProperty(String key, String value) { + return key != null + && value != null + && (key.toLowerCase().contains("compress") || key.toLowerCase().contains("codec")); + } + + static String normalizeCompressionKind(String compression) { + if (compression == null) { + return null; + } + String normalized = compression.trim().toLowerCase(); + if (normalized.isEmpty() + || "none".equals(normalized) + || "no".equals(normalized) + || "false".equals(normalized) + || "uncompressed".equals(normalized)) { + return null; + } + if (normalized.contains("snappy")) { + return "snappy"; + } + if (normalized.contains("gzip")) { + return "gzip"; + } + if (normalized.contains("zstd") || normalized.contains("zstandard")) { + return "zstd"; + } + if (normalized.contains("lz4")) { + return "lz4"; + } + if (normalized.contains("lzo")) { + return "lzo"; + } + if (normalized.contains("zlib") || normalized.contains("deflate")) { + return "zlib"; + } + return normalized; + } } 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..e5b5b071413 --- /dev/null +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/velox/HiveSourceSinkFactoryTest.java @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.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 addCompressionParamsReadsHiveTableProperties() { + Properties tableProperties = new Properties(); + tableProperties.setProperty("orc.compress", "SNAPPY"); + + Map tableParams = new HashMap<>(); + HiveSourceSinkFactory.addCompressionParamsFromTableProperties(tableProperties, tableParams); + + assertThat(tableParams) + .containsEntry("orc.compress", "SNAPPY") + .containsEntry("sink.file.compression", "snappy"); + } + + @Test + void normalizeCompressionKindMapsHiveTablePropertyValues() { + assertThat(HiveSourceSinkFactory.normalizeCompressionKind("SNAPPY")).isEqualTo("snappy"); + assertThat(HiveSourceSinkFactory.normalizeCompressionKind("GZIP")).isEqualTo("gzip"); + assertThat(HiveSourceSinkFactory.normalizeCompressionKind("zstandard")).isEqualTo("zstd"); + assertThat(HiveSourceSinkFactory.normalizeCompressionKind("UNCOMPRESSED")).isNull(); + } +} From f5128fe22074bbf0c31c76c3c75db6684b3eb855 Mon Sep 17 00:00:00 2001 From: zouyunhe Date: Wed, 22 Jul 2026 02:59:37 +0000 Subject: [PATCH 3/3] fix reviews --- .../gluten/velox/HiveSourceSinkFactory.java | 114 +++++++++--------- .../velox/HiveSourceSinkFactoryTest.java | 91 ++++++++++++-- 2 files changed, 137 insertions(+), 68 deletions(-) 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 b030dd4f519..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 @@ -29,7 +29,7 @@ public class HiveSourceSinkFactory extends FileSystemSinkFactory { private static final String COMPRESSION_KIND = "sink.file.compression"; - private static final String[] TABLE_COMPRESSION_KEYS = { + private static final String[] SUPPORTED_COMPRESSION_TABLE_KEYS = { "orc.compress", "parquet.compression", "parquet.compression.codec", "parquet.compression-codec" }; @@ -47,9 +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(fileWriterOperator, tableParams); + addHiveCompressionParams(bucketsBuilder, tableParams); return tableParams; } @@ -63,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) { @@ -103,21 +110,22 @@ private String resolveFormatFromHiveWriterFactory(Object hiveWriterFactory) { return inferFormatFromClassName(outputFormatClz.getName()); } - private void addHiveCompressionParams( - OneInputStreamOperator fileWriterOperator, Map tableParams) { - Object hiveWriterFactory = getHiveWriterFactory(fileWriterOperator); + private void addHiveCompressionParams(Object bucketsBuilder, Map tableParams) { + Object hiveWriterFactory = getHiveWriterFactoryFromBucketsBuilder(bucketsBuilder); if (hiveWriterFactory == null) { return; } Properties tableProperties = (Properties) ReflectUtils.tryGetObjectField(hiveWriterFactory, "tableProperties"); - addCompressionParamsFromTableProperties(tableProperties, tableParams); + addNativeCompressionParamFromTableProperties(tableProperties, tableParams); } - private Object getHiveWriterFactory(OneInputStreamOperator fileWriterOperator) { - Object bucketsBuilder = - ReflectUtils.getObjectField( - ABSTRACT_STREAMING_WRITER_CLASS, fileWriterOperator, "bucketsBuilder"); + 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")) { @@ -126,16 +134,14 @@ private Object getHiveWriterFactory(OneInputStreamOperator fileWriterOpera return ReflectUtils.getObjectField(writerFactory.getClass(), writerFactory, "factory"); } - static void addCompressionParamsFromTableProperties( + static void addNativeCompressionParamFromTableProperties( Properties tableProperties, Map tableParams) { - if (tableProperties == null) { + if (!isParquetFormat(tableParams.get("format"))) { + tableParams.remove(COMPRESSION_KIND); return; } - for (String key : tableProperties.stringPropertyNames()) { - String value = tableProperties.getProperty(key); - if (isCompressionProperty(key, value)) { - tableParams.putIfAbsent(key, value); - } + if (tableProperties == null) { + return; } String compressionKind = resolveCompressionKind(tableProperties); @@ -144,8 +150,12 @@ static void addCompressionParamsFromTableProperties( } } + private static boolean isParquetFormat(String format) { + return format != null && "parquet".equalsIgnoreCase(format.trim()); + } + private static String resolveCompressionKind(Properties tableProperties) { - for (String key : TABLE_COMPRESSION_KEYS) { + for (String key : SUPPORTED_COMPRESSION_TABLE_KEYS) { String compressionKind = normalizeCompressionKind(tableProperties.getProperty(key)); if (compressionKind != null) { return compressionKind; @@ -154,42 +164,32 @@ private static String resolveCompressionKind(Properties tableProperties) { return null; } - private static boolean isCompressionProperty(String key, String value) { - return key != null - && value != null - && (key.toLowerCase().contains("compress") || key.toLowerCase().contains("codec")); - } - static String normalizeCompressionKind(String compression) { if (compression == null) { return null; } - String normalized = compression.trim().toLowerCase(); - if (normalized.isEmpty() - || "none".equals(normalized) - || "no".equals(normalized) - || "false".equals(normalized) - || "uncompressed".equals(normalized)) { - return null; - } - if (normalized.contains("snappy")) { - return "snappy"; - } - if (normalized.contains("gzip")) { - return "gzip"; - } - if (normalized.contains("zstd") || normalized.contains("zstandard")) { - return "zstd"; - } - if (normalized.contains("lz4")) { - return "lz4"; - } - if (normalized.contains("lzo")) { - return "lzo"; - } - if (normalized.contains("zlib") || normalized.contains("deflate")) { - return "zlib"; + 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; } - return normalized; } } 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 index e5b5b071413..1b1dabfbe77 100644 --- 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 @@ -27,23 +27,92 @@ class HiveSourceSinkFactoryTest { @Test - void addCompressionParamsReadsHiveTableProperties() { + 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("orc.compress", "SNAPPY"); + tableProperties.setProperty("parquet.compression.codec", "SNAPPY"); Map tableParams = new HashMap<>(); - HiveSourceSinkFactory.addCompressionParamsFromTableProperties(tableProperties, tableParams); + tableParams.put("format", "parquet"); + HiveSourceSinkFactory.addNativeCompressionParamFromTableProperties( + tableProperties, tableParams); - assertThat(tableParams) - .containsEntry("orc.compress", "SNAPPY") - .containsEntry("sink.file.compression", "snappy"); + 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 normalizeCompressionKindMapsHiveTablePropertyValues() { - assertThat(HiveSourceSinkFactory.normalizeCompressionKind("SNAPPY")).isEqualTo("snappy"); - assertThat(HiveSourceSinkFactory.normalizeCompressionKind("GZIP")).isEqualTo("gzip"); - assertThat(HiveSourceSinkFactory.normalizeCompressionKind("zstandard")).isEqualTo("zstd"); - assertThat(HiveSourceSinkFactory.normalizeCompressionKind("UNCOMPRESSED")).isNull(); + 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"); } }