GH-2976: Parquet CLI compression commands should accept lowercase com…

…pression name (#2977)
apache · Aug 7, 2024 · fd82995 · fd82995
1 parent 26febde
commit fd82995
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 4 deletions.
diff --git a/parquet-cli/src/main/java/org/apache/parquet/cli/commands/RewriteCommand.java b/parquet-cli/src/main/java/org/apache/parquet/cli/commands/RewriteCommand.java
@@ -30,6 +30,7 @@
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.parquet.cli.BaseCommand;
+import org.apache.parquet.cli.util.Codecs;
 import org.apache.parquet.hadoop.metadata.CompressionCodecName;
 import org.apache.parquet.hadoop.rewrite.MaskMode;
 import org.apache.parquet.hadoop.rewrite.ParquetRewriter;
@@ -114,7 +115,7 @@ private RewriteOptions buildOptionsOrFail() throws IOException {
     }
 
     if (codec != null) {
-      CompressionCodecName codecName = CompressionCodecName.valueOf(codec);
+      CompressionCodecName codecName = Codecs.parquetCodec(codec);
       builder.transform(codecName);
     }
 

diff --git a/parquet-cli/src/main/java/org/apache/parquet/cli/commands/TransCompressionCommand.java b/parquet-cli/src/main/java/org/apache/parquet/cli/commands/TransCompressionCommand.java
@@ -30,6 +30,7 @@
 import org.apache.hadoop.fs.Path;
 import org.apache.parquet.HadoopReadOptions;
 import org.apache.parquet.cli.BaseCommand;
+import org.apache.parquet.cli.util.Codecs;
 import org.apache.parquet.hadoop.ParquetFileReader;
 import org.apache.parquet.hadoop.ParquetFileWriter;
 import org.apache.parquet.hadoop.metadata.CompressionCodecName;
@@ -76,7 +77,7 @@ public int run() throws IOException {
 
     Path inPath = new Path(input);
     Path outPath = new Path(output);
-    CompressionCodecName codecName = CompressionCodecName.valueOf(codec);
+    CompressionCodecName codecName = Codecs.parquetCodec(codec);
 
     ParquetMetadata metaData = ParquetFileReader.readFooter(getConf(), inPath, NO_FILTER);
     MessageType schema = metaData.getFileMetaData().getSchema();

diff --git a/parquet-cli/src/test/java/org/apache/parquet/cli/commands/RewriteCommandTest.java b/parquet-cli/src/test/java/org/apache/parquet/cli/commands/RewriteCommandTest.java
@@ -67,4 +67,32 @@ public void testRewriteCommandWithOverwrite() throws IOException {
     Assert.assertEquals(0, command.run());
     Assert.assertTrue(output.exists());
   }
+
+  @Test
+  public void testRewriteCommandWithCompression_GZIP() throws IOException {
+    File file = parquetFile();
+    RewriteCommand command = new RewriteCommand(createLogger());
+    command.inputs = Arrays.asList(file.getAbsolutePath());
+    File output = new File(getTempFolder(), "converted-1.GZIP.parquet");
+    command.output = output.getAbsolutePath();
+    command.codec = "GZIP";
+    command.setConf(new Configuration());
+
+    Assert.assertEquals(0, command.run());
+    Assert.assertTrue(output.exists());
+  }
+
+  @Test
+  public void testRewriteCommandWithCompression_gzip() throws IOException {
+    File file = parquetFile();
+    RewriteCommand command = new RewriteCommand(createLogger());
+    command.inputs = Arrays.asList(file.getAbsolutePath());
+    File output = new File(getTempFolder(), "converted-2.gzip.parquet");
+    command.output = output.getAbsolutePath();
+    command.codec = "gzip";
+    command.setConf(new Configuration());
+
+    Assert.assertEquals(0, command.run());
+    Assert.assertTrue(output.exists());
+  }
 }
diff --git a/parquet-cli/src/test/java/org/apache/parquet/cli/commands/TransCompressionCommandTest.java b/parquet-cli/src/test/java/org/apache/parquet/cli/commands/TransCompressionCommandTest.java
@@ -27,17 +27,32 @@
 public class TransCompressionCommandTest extends ParquetFileTest {
 
   @Test
-  public void testTransCompressionCommand() throws IOException {
+  public void testTransCompressionCommand_ZSTD() throws IOException {
     TransCompressionCommand command = new TransCompressionCommand(createLogger());
 
     command.input = parquetFile().getAbsolutePath();
 
-    File output = new File(getTempFolder(), getClass().getSimpleName() + ".converted.parquet");
+    File output = new File(getTempFolder(), getClass().getSimpleName() + ".converted-1.ZSTD.parquet");
     command.output = output.getAbsolutePath();
     command.codec = "ZSTD";
     command.setConf(new Configuration());
 
     Assert.assertEquals(0, command.run());
     Assert.assertTrue(output.exists());
   }
+
+  @Test
+  public void testTransCompressionCommand_zstd() throws IOException {
+    TransCompressionCommand command = new TransCompressionCommand(createLogger());
+
+    command.input = parquetFile().getAbsolutePath();
+
+    File output = new File(getTempFolder(), getClass().getSimpleName() + ".converted-2.zstd.parquet");
+    command.output = output.getAbsolutePath();
+    command.codec = "zstd";
+    command.setConf(new Configuration());
+
+    Assert.assertEquals(0, command.run());
+    Assert.assertTrue(output.exists());
+  }
 }