update

Author: ArnavBalyan

parquet-cli/src/main/java/org/apache/parquet/cli/commands/ShowSizeStatisticsCommand.java

@@ -27,6 +27,8 @@
 import com.google.common.collect.Lists;
 import java.io.IOException;
 import java.util.List;
+import java.util.Set;
+import java.util.HashSet;
 import org.apache.commons.text.TextStringBuilder;
 import org.apache.parquet.cli.BaseCommand;
 import org.apache.parquet.column.statistics.SizeStatistics;
@@ -47,6 +49,16 @@ public ShowSizeStatisticsCommand(Logger console) {
   @Parameter(description = "<parquet path>")
   List<String> targets;
 
+  @Parameter(
+      names = {"-c", "--column", "--columns"},
+      description = "List of columns (dot paths) to include")
+  List<String> columns;
+
+  @Parameter(
+      names = {"-r", "--row-group", "--row-groups"},
+      description = "List of row-group indexes to include (0-based)")
+  List<Integer> rowGroups;
+
   @Override
   @SuppressWarnings("unchecked")
   public int run() throws IOException {
@@ -60,9 +72,13 @@ public int run() throws IOException {
 
       console.info("\nFile path: {}", source);
 
-      List<BlockMetaData> rowGroups = footer.getBlocks();
-      for (int index = 0, n = rowGroups.size(); index < n; index++) {
-        printRowGroupSizeStats(console, index, rowGroups.get(index), schema);
+      List<BlockMetaData> blocks = footer.getBlocks();
+      Set<Integer> allowedRowGroups = rowGroups == null ? null : new HashSet<>(rowGroups);
+      for (int index = 0, n = blocks.size(); index < n; index++) {
+        if (allowedRowGroups != null && !allowedRowGroups.contains(index)) {
+          continue;
+        }
+        printRowGroupSizeStats(console, index, blocks.get(index), schema);
         console.info("");
       }
     }
@@ -84,7 +100,16 @@ private void printRowGroupSizeStats(Logger console, int index, BlockMetaData row
     console.info(
         String.format(formatString, "column", "unencoded bytes", "rep level histogram", "def level histogram"));
 
+    Set<String> allowedColumns = null;
+    if (columns != null && !columns.isEmpty()) {
+      allowedColumns = new HashSet<>(columns);
+    }
+
     for (ColumnChunkMetaData column : rowGroup.getColumns()) {
+      String dotPath = column.getPath().toDotString();
+      if (allowedColumns != null && !allowedColumns.contains(dotPath)) {
+        continue;
+      }
       printColumnSizeStats(console, column, schema, maxColumnWidth);
     }
   }
@@ -111,6 +136,12 @@ private void printColumnSizeStats(Logger console, ColumnChunkMetaData column, Me
 
   @Override
   public List<String> getExamples() {
-    return Lists.newArrayList("# Show size statistics for a Parquet file", "sample.parquet");
+    return Lists.newArrayList(
+        "# Show size statistics for a Parquet file",
+        "sample.parquet",
+        "# Show size statistics for selected columns",
+        "sample.parquet -c name,tags",
+        "# Show size statistics for a specific row-group",
+        "sample.parquet -r 0");
   }
 }

parquet-cli/src/test/java/org/apache/parquet/cli/commands/ShowSizeStatisticsCommandTest.java

@@ -34,4 +34,24 @@ public void testShowSizeStatisticsCommand() throws IOException {
     command.setConf(new Configuration());
     Assert.assertEquals(0, command.run());
   }
+
+  @Test
+  public void testShowSizeStatisticsWithColumnFilter() throws IOException {
+    File file = parquetFile();
+    ShowSizeStatisticsCommand command = new ShowSizeStatisticsCommand(createLogger());
+    command.targets = Arrays.asList(file.getAbsolutePath());
+    command.columns = Arrays.asList(INT32_FIELD, INT64_FIELD);
+    command.setConf(new Configuration());
+    Assert.assertEquals(0, command.run());
+  }
+
+  @Test
+  public void testShowSizeStatisticsWithRowGroupFilter() throws IOException {
+    File file = parquetFile();
+    ShowSizeStatisticsCommand command = new ShowSizeStatisticsCommand(createLogger());
+    command.targets = Arrays.asList(file.getAbsolutePath());
+    command.rowGroups = Arrays.asList(0);
+    command.setConf(new Configuration());
+    Assert.assertEquals(0, command.run());
+  }
 }