tests and fixes

Author: Damian Stewart

src/__init__.py

@@ -10,7 +10,7 @@
 from tqdm.auto import tqdm
 
 
-def are_parquet_file_row_groups_sorted(pf: pq.ParquetFile, column_name: str) -> bool:
+def are_parquet_file_row_groups_sorted(pf: pq.ParquetFile, column_name: str) -> tuple[bool, str, str]:
     sort_column_index = next(i for i, name in enumerate(pf.schema.names)
                              if name == column_name)
 
@@ -27,6 +27,7 @@ def are_parquet_file_row_groups_sorted(pf: pq.ParquetFile, column_name: str) ->
             return False, None, None
         whole_min = column.statistics.min if whole_min is None else min(column.statistics.min, whole_min)
         whole_max = column.statistics.max if whole_max is None else max(column.statistics.max, whole_max)
+        prev_max = column.statistics.max
     return True, whole_min, whole_max
 
 

src/script/__init__.py

@@ -1,21 +1,15 @@
+import pytest
+import random
 from unittest.mock import MagicMock
 import os
 import sys
 
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'script'))
-from is_table_sorted import are_parquet_file_row_groups_sorted
+from src.script.is_table_sorted import are_parquet_file_row_groups_sorted
 
 
-def _create_mock_parquet_file(column_name: str, row_groups_stats: list[tuple]):
-    """
-    Helper to create a mock ParquetFile with specified row group statistics.
-
-    Args:
-        column_name: Name of the column to sort by
-        row_groups_stats: List of (min, max) tuples for each row group
-    """
+def _create_mock_parquet_file(column_name: str, row_groups_stats: list[tuple[str, str]]):
     mock_pf = MagicMock()
-    mock_pf.schema.names = [column_name, 'data']
+    mock_pf.schema.names = [column_name]
     mock_pf.num_row_groups = len(row_groups_stats)
 
     mock_row_groups = []
@@ -31,180 +25,46 @@ def _create_mock_parquet_file(column_name: str, row_groups_stats: list[tuple]):
     return mock_pf
 
 
-# Tests for sorted row groups
-
 def test_single_row_group_sorted():
-    """Test with a single row group (trivially sorted)"""
-    mock_pf = _create_mock_parquet_file(
-        'url_surtkey',
-        [('com,example)/page1', 'com,example)/page3')]
-    )
-
-    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, 'url_surtkey')
-
-    assert is_sorted is True
-    assert min_val is not None
-    assert max_val is not None
-
-
-def test_multiple_row_groups_strictly_increasing():
-    """Test with multiple row groups in strictly increasing order"""
-    mock_pf = _create_mock_parquet_file(
-        'url_surtkey',
-        [
-            ('com,aaa)/', 'com,bbb)/'),
-            ('com,ccc)/', 'com,ddd)/'),
-            ('com,eee)/', 'com,fff)/')
-        ]
-    )
-
-    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, 'url_surtkey')
-
-    assert is_sorted is True
-    assert min_val is not None
-    assert max_val is not None
-
-
-def test_boundary_case_adjacent_values():
-    """Test with row groups that have adjacent but non-overlapping values"""
-    mock_pf = _create_mock_parquet_file(
-        'url',
-        [
-            ('com,example)/a', 'com,example)/z'),
-            ('com,example,aaa)/', 'com,example,zzz)/')
-        ]
-    )
-
-    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, 'url')
-
-    assert is_sorted is True
-    assert min_val is not None
-    assert max_val is not None
-
-
-def test_two_row_groups_strictly_increasing_strings():
-    """Test with two row groups with string values in strictly increasing order"""
-    mock_pf = _create_mock_parquet_file(
-        'url_surtkey',
-        [
-            ('com,apple)/', 'com,banana)/'),
-            ('com,cherry)/', 'com,date)/')
-        ]
-    )
-
-    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, 'url_surtkey')
-
-    assert is_sorted is True
-    assert min_val is not None
-    assert max_val is not None
-
-
-def test_many_row_groups_strictly_increasing():
-    """Test with many row groups, all strictly increasing"""
-    row_groups = [
-        ('com,aaa)/', 'com,aaa,zzz)/'),
-        ('com,bbb)/', 'com,bbb,zzz)/'),
-        ('com,ccc)/', 'com,ccc,zzz)/'),
-        ('com,ddd)/', 'com,ddd,zzz)/'),
-        ('com,eee)/', 'com,eee,zzz)/'),
-    ]
+    mock_pf = _create_mock_parquet_file('url_surtkey', [('a', 'b')])
+    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, column_name='url_surtkey')
+    assert is_sorted
+    assert min_val == 'a'
+    assert max_val == 'b'
+
+
+def test_row_groups_sorted():
+    all_row_groups = [('a', 'b'), ('c', 'd'), ('e', 'f'), ('g', 'h')]
+    for n in range(1, len(all_row_groups)):
+        row_groups = all_row_groups[:n]
+        mock_pf = _create_mock_parquet_file('url_surtkey', row_groups)
+        is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, column_name='url_surtkey')
+        assert is_sorted
+        assert is_sorted
+        assert min_val == row_groups[0][0]
+        assert max_val == row_groups[-1][1]
+
+
+def test_row_groups_unsorted():
+    all_row_groups = [('a', 'b'), ('c', 'd'), ('e', 'f'), ('g', 'h')]
+    count = 0
+    while count < 100:
+        for n in range(2, len(all_row_groups)):
+            row_groups = all_row_groups[:n].copy()
+            random.shuffle(row_groups)
+            if row_groups == all_row_groups[:n]:
+                # shuffle resulted in same order, try again
+                continue
+
+            mock_pf = _create_mock_parquet_file('url_surtkey', row_groups)
+            is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, column_name='url_surtkey')
+            assert not is_sorted
+
+        count += 1
+
+
+def test_row_groups_overlapping():
+    row_groups = [('a', 'c'), ('b', 'd')]
     mock_pf = _create_mock_parquet_file('url_surtkey', row_groups)
-
-    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, 'url_surtkey')
-
-    assert is_sorted is True
-    assert min_val is not None
-    assert max_val is not None
-
-
-# Tests for non-sorted row groups
-
-def test_two_row_groups_overlapping():
-    """Test with two row groups where second min is less than first max (overlapping)"""
-    mock_pf = _create_mock_parquet_file(
-        'url_surtkey',
-        [
-            ('a', 'd'),
-            ('b', 'e')
-        ]
-    )
-
-    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, 'url_surtkey')
-
-    assert is_sorted is False
-    assert min_val is None
-    assert max_val is None
-
-
-def test_row_groups_completely_out_of_order():
-    """Test with row groups in descending order"""
-    mock_pf = _create_mock_parquet_file(
-        'url_surtkey',
-        [
-            ('z', 'zz'),
-            ('a', 'b')  # completely before the first group
-        ]
-    )
-
-    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, 'url_surtkey')
-
-    assert is_sorted is False
-    assert min_val is None
-    assert max_val is None
-
-
-def test_multiple_row_groups_with_middle_unsorted():
-    """Test with multiple row groups where the middle one breaks the sort order"""
-    mock_pf = _create_mock_parquet_file(
-        'url_surtkey',
-        [
-            ('a', 'b'),
-            ('z', 'zz'),  # correctly sorted so far
-            ('c', 'd')  # breaks ordering (min 'c' < previous max 'zz')
-        ]
-    )
-
-    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, 'url_surtkey')
-
-    assert is_sorted is False
-    assert min_val is None
-    assert max_val is None
-
-
-def test_row_groups_equal_boundary_allowed():
-    """Test that row groups where second min equals first max are allowed (>= not >)"""
-    mock_pf = _create_mock_parquet_file(
-        'url_surtkey',
-        [
-            ('a', 'b'),
-            ('b', 'c')  # min equals prev_max - this is allowed
-        ]
-    )
-
-    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, 'url_surtkey')
-
-    assert is_sorted is True
-    assert min_val is not None
-    assert max_val is not None
-
-
-def test_slight_overlap_in_middle():
-    """Test detecting overlap in the middle of many row groups"""
-    mock_pf = _create_mock_parquet_file(
-        'url_surtkey',
-        [
-            ('a', 'az'),
-            ('b', 'bz'),
-            ('c', 'cz'),
-            ('ba', 'baz'),  # overlaps with previous ('ba' < 'c')
-            ('d', 'dz'),
-        ]
-    )
-
-    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, 'url_surtkey')
-
-    assert is_sorted is False
-    assert min_val is None
-    assert max_val is None
-
-
+    is_sorted, min_val, max_val = are_parquet_file_row_groups_sorted(mock_pf, column_name='url_surtkey')
+    assert not is_sorted