Move RowGroup size calculation from TUI to Reader logic and fix bug on uncompressed_size computation from pyarrow

Author: raulcd

src/datanomy/reader/parquet.py

@@ -7,6 +7,90 @@
 from pyarrow.lib import ArrowInvalid
 
 
+class RowGroup:
+    """Class to represent a Parquet row group."""
+
+    def __init__(self, row_group_metadata: Any) -> None:
+        """
+        Initialize the RowGroup.
+
+        Parameters
+        ----------
+            row_group_metadata: Metadata for the row group
+        """
+        self._metadata = row_group_metadata
+
+    @property
+    def num_columns(self) -> int:
+        """
+        Get number of columns in the row group.
+
+        Returns
+        -------
+            Number of columns
+        """
+        return int(self._metadata.num_columns)
+
+    @property
+    def num_rows(self) -> int:
+        """
+        Get number of rows in the row group.
+
+        Returns
+        -------
+            Number of rows
+        """
+        return int(self._metadata.num_rows)
+
+    def column(self, index: int) -> Any:
+        """
+        Get metadata for a specific column in the row group.
+
+        Parameters
+        ----------
+            index: Column index
+
+        Returns
+        -------
+            Column metadata
+        """
+        return self._metadata.column(index)
+
+    @property
+    def has_compression(self) -> bool:
+        """
+        Check if any column in the row group uses compression.
+
+        Returns
+        -------
+            True if any column is compressed, False otherwise
+        """
+        for j in range(self.num_columns):
+            if self.column(j).compression != "UNCOMPRESSED":
+                return True
+        return False
+
+    @property
+    def total_sizes(self) -> tuple[int, int]:
+        """
+        Get total compressed and uncompressed size of the row group in bytes.
+
+        Returns
+        -------
+            Total compressed and uncompressed size of the row group in bytes.
+            Tuple of (compressed_size, uncompressed_size)
+        """
+        compressed_sum = sum(
+            self.column(j).total_compressed_size for j in range(self.num_columns)
+        )
+        # This should be self._metadata.total_byte_size but there's a bug on
+        # pyarrow 22.0.0, see: https://github.com/apache/arrow/issues/48138
+        uncompressed_sum = sum(
+            self.column(j).total_uncompressed_size for j in range(self.num_columns)
+        )
+        return compressed_sum, uncompressed_sum
+
+
 class ParquetReader:
     """Main class to read and inspect Parquet files."""
 
@@ -114,6 +198,20 @@ def get_row_group_info(self, index: int) -> Any:
         """
         return self.parquet_file.metadata.row_group(index)
 
+    def get_row_group(self, index: int) -> RowGroup:
+        """
+        Get a specific RowGroup object.
+
+        Parameters
+        ----------
+            index: Row group index
+
+        Returns
+        -------
+            RowGroup object
+        """
+        return RowGroup(self.parquet_file.metadata.row_group(index))
+
     @property
     def metadata_size(self) -> int:
         """

src/datanomy/tui/parquet.py

@@ -92,27 +92,16 @@ def _row_groups(self) -> list[Panel]:
         # Row groups
         row_group_panels: list[Panel] = []
         for i in range(self.reader.num_row_groups):
-            rg = self.reader.get_row_group_info(i)
-
-            # Calculate compressed and uncompressed sizes
-            compressed_sum = sum(
-                rg.column(j).total_compressed_size for j in range(rg.num_columns)
-            )
-            uncompressed_sum = rg.total_byte_size  # This is the uncompressed size
-
-            # Check if any column is compressed
-            has_compression = any(
-                rg.column(j).compression != "UNCOMPRESSED"
-                for j in range(rg.num_columns)
-            )
+            rg = self.reader.get_row_group(i)
+            compressed_sum, uncompressed_sum = rg.total_sizes
 
             compressed_str = format_size(compressed_sum)
             uncompressed_str = format_size(uncompressed_sum)
 
             # Summary info
             rg_summary = Text()
             rg_summary.append(f"Rows: {rg.num_rows:,}\n")
-            if has_compression:
+            if rg.has_compression:
                 rg_summary.append(f"Compressed: {compressed_str}\n")
                 rg_summary.append(f"Uncompressed: {uncompressed_str}\n")
                 # Calculate compression ratio

tests/test_reader.py

@@ -86,3 +86,15 @@ def test_reader_accepts_file_without_parquet_extension(
     reader = ParquetReader(parquet_without_extension)
     assert reader.num_rows == 3
     assert len(reader.schema_arrow) == 2
+
+
+def test_row_group_total_sizes(simple_parquet: Path) -> None:
+    """Test that row_group_total_sizes returns correct compressed and uncompressed sizes."""
+    reader = ParquetReader(simple_parquet)
+
+    row_group = reader.get_row_group(0)
+
+    assert row_group.has_compression
+    compressed_size, uncompressed_size = row_group.total_sizes
+    assert compressed_size == 474
+    assert uncompressed_size == 487