allow specifying row group index when reading into Table

aloneguid · Apr 27, 2023 · 0d3e83d · 0d3e83d
1 parent 5ca1d20
commit 0d3e83d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/.github/workflows/full.yml b/.github/workflows/full.yml
@@ -1,7 +1,7 @@
 name: 'Full Workflow'
 
 env:
-  VERSION: 4.10.0
+  VERSION: 4.10.1
   ASM_VERSION: 4.0.0
 
 on:

diff --git a/src/Parquet/ParquetExtensions.cs b/src/Parquet/ParquetExtensions.cs
@@ -94,19 +94,28 @@ public static async Task WriteAsync(this Table table, string path,
         /// </summary>
         /// <param name="reader">Open reader</param>
         /// <param name="progressCallback"></param>
+        /// <param name="rowGroupIndex">When specified, only row group at index will be read.</param>
         /// <returns></returns>
         public static async Task<Table> ReadAsTableAsync(this ParquetReader reader,
-            TableReaderProgressCallback? progressCallback = null) {
+            TableReaderProgressCallback? progressCallback = null,
+            int? rowGroupIndex = null) {
             Table? result = null;
             DataField[] dataFields = reader.Schema!.GetDataFields();
 
-            int stepsTotal = dataFields.Length * reader.RowGroupCount;
+            var rowGroupIndexes = new List<int>();
+            if(rowGroupIndex != null) {
+                rowGroupIndexes.Add(rowGroupIndex.Value);
+            } else {
+                rowGroupIndexes.AddRange(Enumerable.Range(0, reader.RowGroupCount));
+            }
+
+            int stepsTotal = dataFields.Length * rowGroupIndexes.Count;
             int currentStep = 0;
 
             if(reader.RowGroupCount == 0) {
                 result = new Table(reader.Schema, null, 0);
             } else {
-                for(int i = 0; i < reader.RowGroupCount; i++) {
+                foreach(int i in rowGroupIndexes) {
                     using(ParquetRowGroupReader rowGroupReader = reader.OpenRowGroupReader(i)) {
                         DataColumn[] allData = new DataColumn[dataFields.Length];