From 164b6e156facfdb8c5e634a566a6b69685686ded Mon Sep 17 00:00:00 2001 From: Colin Leek Date: Thu, 2 Jul 2026 17:58:48 +0000 Subject: [PATCH 1/2] test: failing tests for #92 scancount --persegment Unit tests assert: - DictAccumulator (AccumulatorParam subclass) exists for per-segment count propagation via Spark Accumulators (not plain dicts) - _count_data adds segment count to per_segment_accumulator - run() creates DictAccumulator when --persegment is set - run() prints per-segment counts when --persegment is set - Backward compat: no per-segment output without the flag E2e test (tests/e2e/commands/test_scancount_persegment_smoke.py): - Runs real Glue scancount --persegment against a transient table - Verifies per-segment output lines appear and sum to total --- .../test_scancount_persegment_smoke.py | 133 +++++++++ .../tests/server/test_scancount_persegment.py | 263 ++++++++++++++++++ 2 files changed, 396 insertions(+) create mode 100644 tools/bulk_executor/tests/e2e/commands/test_scancount_persegment_smoke.py create mode 100644 tools/bulk_executor/tests/server/test_scancount_persegment.py diff --git a/tools/bulk_executor/tests/e2e/commands/test_scancount_persegment_smoke.py b/tools/bulk_executor/tests/e2e/commands/test_scancount_persegment_smoke.py new file mode 100644 index 00000000..fb6e6374 --- /dev/null +++ b/tools/bulk_executor/tests/e2e/commands/test_scancount_persegment_smoke.py @@ -0,0 +1,133 @@ +"""Command smoke: `bulk scancount --persegment`. + +Exercises the --persegment flag against a real Glue environment. Creates a +transient table, seeds it with items distributed across segments, then runs +scancount with --persegment and verifies: + +1. The command succeeds (Glue job reaches SUCCEEDED) +2. Per-segment output lines appear in stdout (format: "N: COUNT") +3. The sum of per-segment counts equals the total item count +""" +from __future__ import annotations + +import re + +import pytest + +from tests.e2e.helpers.assertions import assert_glue_succeeded, table_item_count +from tests.e2e.helpers.transient_table import transient_table +from tests.e2e.helpers.command_runner import run_command +from tests.e2e.connector.conftest import PerfRow + + +# Per-segment output lines: ": " or "Segment : " +_SEGMENT_LINE = re.compile(r"(?:Segment\s+)?(\d+)\s*:\s*([\d,]+)") + +# Total line (existing behavior) +_TOTAL_LINE = re.compile(r"Total records counted:\s*([\d,]+)") + + +def _parse_segment_counts(stdout: str) -> dict[int, int]: + """Parse per-segment output lines from scancount --persegment stdout.""" + counts = {} + for match in _SEGMENT_LINE.finditer(stdout): + seg_id = int(match.group(1)) + count = int(match.group(2).replace(",", "")) + counts[seg_id] = count + return counts + + +def _parse_total(stdout: str) -> int | None: + match = _TOTAL_LINE.search(stdout) + if match: + return int(match.group(1).replace(",", "")) + return None + + +@pytest.mark.e2e +class TestScancountPerSegmentSmoke: + """Real-AWS smoke test for scancount --persegment (issue #92).""" + + def test_persegment_outputs_per_segment_counts(self, e2e_config, cmd_perf_collector): + """Run scancount --persegment on a seeded table; verify per-segment output.""" + with transient_table(e2e_config.aws_region, label="scancount-ps") as table: + # Seed the table with enough items to span multiple segments + seed = run_command( + "fill", + table=table, + extra_args=["--numitems", "200", "--generator", "default"], + ) + assert_glue_succeeded("scancount-persegment setup (fill)", seed, e2e_config.aws_region) + + # Verify seed worked + seeded_count = table_item_count(e2e_config.aws_region, table) + assert seeded_count > 0, "Table must have items for scancount test" + + # Run scancount with --persegment flag + result = run_command( + "scancount", + table=table, + extra_args=["--persegment", "--segments", "5"], + ) + perf = assert_glue_succeeded("scancount --persegment", result, e2e_config.aws_region) + + # Parse per-segment output + segment_counts = _parse_segment_counts(result.stdout) + assert len(segment_counts) > 0, ( + f"--persegment must produce per-segment output lines. " + f"Got stdout:\n{result.stdout[-2000:]}" + ) + + # Verify each segment has a non-negative count + for seg_id, count in segment_counts.items(): + assert count >= 0, f"Segment {seg_id} has negative count: {count}" + + # Sum of per-segment counts should equal total + segment_sum = sum(segment_counts.values()) + total = _parse_total(result.stdout) + if total is not None: + assert segment_sum == total, ( + f"Sum of per-segment counts ({segment_sum}) must equal " + f"reported total ({total})" + ) + + # Segment sum should match the actual table item count + assert segment_sum == seeded_count, ( + f"Sum of per-segment counts ({segment_sum}) must equal " + f"actual table item count ({seeded_count})" + ) + + cmd_perf_collector.add(PerfRow( + command="scancount --persegment", + wall_seconds=result.wall_seconds, + dpu_seconds=perf.dpu_seconds if perf else None, + items=segment_sum, + )) + + def test_persegment_without_flag_shows_only_total(self, e2e_config, cmd_perf_collector): + """Without --persegment, scancount still shows only the total (backward compat).""" + with transient_table(e2e_config.aws_region, label="scancount-nops") as table: + seed = run_command( + "fill", + table=table, + extra_args=["--numitems", "50", "--generator", "default"], + ) + assert_glue_succeeded("scancount-nopersegment setup (fill)", seed, e2e_config.aws_region) + + result = run_command("scancount", table=table) + assert_glue_succeeded("scancount (no --persegment)", result, e2e_config.aws_region) + + total = _parse_total(result.stdout) + assert total is not None, ( + f"scancount without --persegment must still print total. " + f"Got:\n{result.stdout[-1000:]}" + ) + + # Should NOT have per-segment breakdown lines + segment_counts = _parse_segment_counts(result.stdout) + # Filter out the total line which might match the pattern + # A real per-segment output would have multiple numbered lines + assert len(segment_counts) <= 1, ( + f"scancount without --persegment should NOT show per-segment " + f"breakdown (got {len(segment_counts)} segment lines)" + ) diff --git a/tools/bulk_executor/tests/server/test_scancount_persegment.py b/tools/bulk_executor/tests/server/test_scancount_persegment.py new file mode 100644 index 00000000..cc8b5a57 --- /dev/null +++ b/tools/bulk_executor/tests/server/test_scancount_persegment.py @@ -0,0 +1,263 @@ +"""Unit tests for scancount --persegment feature (issue #92). + +The --persegment flag should output per-segment item counts so users can +detect data skew across parallel scan segments. The implementation MUST use +Spark Accumulators (not plain Python dicts) to propagate per-segment counts +from executors back to the driver. + +These tests verify: +1. A DictAccumulator (AccumulatorParam subclass) is used for per-segment counts +2. _count_data adds its local_count to the per-segment accumulator keyed by segment id +3. run() prints per-segment output when --persegment is set +4. run() still prints only the total when --persegment is not set (backward compat) +""" + +import json +from unittest.mock import MagicMock, patch + +import pytest + +from python_modules import scancount as sc_module + +sc_module.get_error_message = lambda e: str(e) + + +# --- DictAccumulator ----------------------------------------------------------- + + +class TestDictAccumulator: + """Issue #92 requires a DictAccumulator (AccumulatorParam subclass) that + merges per-segment counts from executors to the driver. + + Plain Python dicts do NOT propagate from Spark executors to the driver — + only Accumulators do. The fix must define a DictAccumulator class.""" + + def test_dict_accumulator_class_exists(self): + """The module must define a DictAccumulator for per-segment count propagation.""" + assert hasattr(sc_module, 'DictAccumulator'), ( + "scancount must define a DictAccumulator (AccumulatorParam subclass) " + "to propagate per-segment counts from executors to driver. " + "Plain dicts do NOT work in Spark's distributed execution model." + ) + + def test_dict_accumulator_is_accumulator_param(self): + """DictAccumulator must subclass AccumulatorParam for Spark compatibility.""" + from pyspark import AccumulatorParam + acc = sc_module.DictAccumulator() + assert isinstance(acc, AccumulatorParam), ( + "DictAccumulator must inherit from AccumulatorParam so Spark " + "can serialize and merge it across executors." + ) + + def test_dict_accumulator_zero_returns_empty_dict(self): + """zero() must return {} — the identity for merging segment counts.""" + acc = sc_module.DictAccumulator() + assert acc.zero({}) == {} + assert acc.zero(None) == {} + + def test_dict_accumulator_addInPlace_merges_counts(self): + """addInPlace must sum values for matching keys and add new keys.""" + acc = sc_module.DictAccumulator() + v1 = {0: 100, 1: 200} + v2 = {1: 50, 2: 300} + result = acc.addInPlace(v1, v2) + assert result == {0: 100, 1: 250, 2: 300}, ( + "addInPlace should sum counts for same segment ids and add new ones" + ) + assert result is v1, "should mutate v1 in place (Spark contract)" + + def test_dict_accumulator_addInPlace_empty_right(self): + """Merging an empty dict leaves v1 unchanged.""" + acc = sc_module.DictAccumulator() + v1 = {3: 500} + result = acc.addInPlace(v1, {}) + assert result == {3: 500} + + def test_dict_accumulator_addInPlace_empty_left(self): + """Merging into an empty dict yields v2's entries.""" + acc = sc_module.DictAccumulator() + result = acc.addInPlace({}, {7: 42}) + assert result == {7: 42} + + +# --- _count_data per-segment accumulator usage --------------------------------- + + +class TestCountDataPerSegmentAccumulator: + """_count_data must add its local_count to the per-segment accumulator + (keyed by segment id) when a per_segment_accumulator is provided.""" + + def _run_count_data_with_persegment(self, monkeypatch, segment=5, + total_segments=200, scan_count=77): + """Helper: run _count_data with a per-segment accumulator and return it.""" + session = MagicMock() + table = MagicMock() + table.scan = MagicMock(return_value={'Count': scan_count}) + session.resource.return_value.Table.return_value = table + + rl = MagicMock() + rl.get_session.return_value = session + monkeypatch.setattr(sc_module, 'RateLimiterWorker', MagicMock(return_value=rl)) + + total_acc = MagicMock() + error_acc = MagicMock() + per_segment_acc = MagicMock() + + sc_module._count_data( + {}, 'tbl', None, None, None, None, + segment, total_segments, + total_acc, error_acc, MagicMock(), + per_segment_acc, + ) + + return per_segment_acc, total_acc + + def test_per_segment_accumulator_receives_segment_count(self, monkeypatch): + """_count_data must call per_segment_accumulator.add({segment: local_count}).""" + per_seg_acc, _ = self._run_count_data_with_persegment( + monkeypatch, segment=5, scan_count=77 + ) + per_seg_acc.add.assert_called_once_with({5: 77}) + + def test_per_segment_accumulator_different_segment_ids(self, monkeypatch): + """Each segment id maps its own count into the accumulator.""" + per_seg_acc, _ = self._run_count_data_with_persegment( + monkeypatch, segment=199, scan_count=12345 + ) + per_seg_acc.add.assert_called_once_with({199: 12345}) + + def test_total_accumulator_still_updated(self, monkeypatch): + """Per-segment accumulator does not replace the total accumulator.""" + _, total_acc = self._run_count_data_with_persegment( + monkeypatch, segment=0, scan_count=42 + ) + total_acc.add.assert_called_once_with(42) + + +# --- run() --persegment output ------------------------------------------------- + + +class TestRunPerSegmentOutput: + """When parsed_args includes persegment=True, run() must: + 1. Create a per-segment accumulator using DictAccumulator + 2. Pass it to each _count_data worker + 3. Print per-segment counts after all workers finish""" + + @pytest.fixture + def base_args(self): + return { + 'table': 'my-table', + 'index': None, + 'filter_expression': None, + 'expression_values': None, + 'expression_names': None, + 's3-bucket-name': 'rate-bucket', + 'JOB_RUN_ID': 'jr-001', + 'persegment': True, + } + + @pytest.fixture + def mock_env(self, monkeypatch): + """Set up the mocked environment for run().""" + monkeypatch.setattr(sc_module.boto3, 'Session', + MagicMock(return_value=MagicMock(region_name='us-east-1'))) + monkeypatch.setattr(sc_module, 'RateLimiterSharedConfig', MagicMock()) + monkeypatch.setattr(sc_module, 'RateLimiterAggregator', MagicMock()) + monkeypatch.setattr(sc_module, 'get_and_print_dynamodb_table_info', MagicMock( + return_value={'item_count': 500, 'size_bytes': 4096, 'region_name': 'us-east-1'} + )) + monkeypatch.setattr(sc_module, 'get_and_print_table_scan_cost', MagicMock(return_value=0.75)) + monkeypatch.setattr(sc_module, 'get_dynamodb_throughput_configs', MagicMock(return_value={})) + + def test_persegment_creates_dict_accumulator(self, monkeypatch, mock_env, base_args): + """run() must create a per-segment accumulator using DictAccumulator when persegment=True.""" + sc = MagicMock() + # Track accumulator calls to verify DictAccumulator is used + accumulator_calls = [] + + def track_accumulator(init, *args): + accumulator_calls.append((init, args)) + return MagicMock(value=init) + + sc.accumulator = MagicMock(side_effect=track_accumulator) + sc.parallelize.return_value.foreach = MagicMock() + sc.parallelize.return_value.count = MagicMock() + + sc_module.run(MagicMock(), sc, MagicMock(), base_args) + + # Should have 3 accumulators: total (int), errors (list), per-segment (dict) + assert len(accumulator_calls) >= 3, ( + f"Expected at least 3 accumulator() calls (total, errors, per-segment), " + f"got {len(accumulator_calls)}" + ) + # The per-segment accumulator should be initialized with {} and DictAccumulator + per_seg_call = accumulator_calls[2] + assert per_seg_call[0] == {}, ( + "Per-segment accumulator should be initialized with empty dict {}" + ) + assert isinstance(per_seg_call[1][0], sc_module.DictAccumulator), ( + "Per-segment accumulator must use DictAccumulator (not a plain dict)" + ) + + def test_persegment_prints_per_segment_counts(self, monkeypatch, mock_env, base_args, capsys): + """run() must print per-segment counts when persegment=True.""" + sc = MagicMock() + # Simulate per-segment results: segments 0-4 with varying counts + per_segment_data = {0: 1000, 1: 900, 2: 1000, 3: 123456789, 4: 1000} + total = sum(per_segment_data.values()) + + accs = [ + MagicMock(value=total), # total_matched_accumulator + MagicMock(value=[]), # error_accumulator + MagicMock(value=per_segment_data), # per_segment_accumulator + ] + sc.accumulator = MagicMock(side_effect=accs) + sc.parallelize.return_value.foreach = MagicMock() + sc.parallelize.return_value.count = MagicMock() + + sc_module.run(MagicMock(), sc, MagicMock(), base_args) + out = capsys.readouterr().out + + # Each segment's count must appear in the output + assert '0:' in out or '0 :' in out or 'Segment 0' in out, ( + f"Per-segment output must include segment 0's count. Got:\n{out}" + ) + assert '123,456,789' in out or '123456789' in out, ( + f"Per-segment output must show segment 3's large count. Got:\n{out}" + ) + + def test_no_persegment_flag_omits_per_segment_output(self, monkeypatch, mock_env, base_args, capsys): + """Without --persegment, run() prints only the total (backward compat).""" + base_args['persegment'] = False + + sc = MagicMock() + accs = [MagicMock(value=5000), MagicMock(value=[])] + sc.accumulator = MagicMock(side_effect=accs) + sc.parallelize.return_value.foreach = MagicMock() + sc.parallelize.return_value.count = MagicMock() + + sc_module.run(MagicMock(), sc, MagicMock(), base_args) + out = capsys.readouterr().out + + # Should show total but NOT per-segment breakdown + assert '5,000' in out, "Total should still print" + # Should not have segment-numbered lines + for i in range(200): + assert f'{i}: ' not in out or f'Segment {i}' not in out, ( + "Per-segment output should NOT appear without --persegment" + ) + + def test_persegment_missing_defaults_to_false(self, monkeypatch, mock_env, base_args, capsys): + """If 'persegment' key is missing from args, behave as if False.""" + del base_args['persegment'] + + sc = MagicMock() + accs = [MagicMock(value=100), MagicMock(value=[])] + sc.accumulator = MagicMock(side_effect=accs) + sc.parallelize.return_value.foreach = MagicMock() + sc.parallelize.return_value.count = MagicMock() + + # Should not raise — persegment defaults to off + sc_module.run(MagicMock(), sc, MagicMock(), base_args) + out = capsys.readouterr().out + assert '100' in out From bce36b88fcb677fd4ef175cdd63951213e896b90 Mon Sep 17 00:00:00 2001 From: Colin Leek Date: Thu, 2 Jul 2026 18:14:45 +0000 Subject: [PATCH 2/2] fix: add --persegment support to scancount using DictAccumulator (#92) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements per-segment item count reporting for scancount using a Spark Accumulator (DictAccumulator) to correctly propagate counts from executors to the driver. Plain dicts don't work in Spark's distributed model — only Accumulators do. Changes: - Add DictAccumulator(AccumulatorParam) that merges segment count dicts - Add per_segment_accumulator parameter to _count_data - Create per-segment accumulator in run() when persegment=True - Print per-segment breakdown after total count --- .../src/python_modules/scancount/__init__.py | 25 +++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/tools/bulk_executor/server/src/python_modules/scancount/__init__.py b/tools/bulk_executor/server/src/python_modules/scancount/__init__.py index eb437c27..15d9608c 100644 --- a/tools/bulk_executor/server/src/python_modules/scancount/__init__.py +++ b/tools/bulk_executor/server/src/python_modules/scancount/__init__.py @@ -41,6 +41,15 @@ def addInPlace(self, v1, v2): v1.extend(v2) return v1 +class DictAccumulator(AccumulatorParam): + def zero(self, initialValue): + return {} + + def addInPlace(self, v1, v2): + for k, v in v2.items(): + v1[k] = v1.get(k, 0) + v + return v1 + DYNAMO_DB_THROTTLE_EXCEPTION = 'ProvisionedThroughputExceededException' DYNAMO_DB_VALIDATION_EXCEPTION = 'ValidationException' @@ -77,11 +86,16 @@ def run(job, spark_context, glue_context, parsed_args): # Since each task might generate errors, let's accumulate them and report intelligently error_accumulator = spark_context.accumulator([], ListAccumulator()) + persegment = parsed_args.get('persegment', False) + per_segment_accumulator = None + if persegment: + per_segment_accumulator = spark_context.accumulator({}, DictAccumulator()) + # Distribute work among partitions, each knowing what segment it's to handle try: parallelize_count = 200 rdd = spark_context.parallelize(range(parallelize_count), parallelize_count) - rdd.foreach(lambda worker_id: _count_data(monitor_options, table_name, index_name, filter_expression, expression_values, expression_names, worker_id, parallelize_count, total_matched_accumulator, error_accumulator, rate_limiter_shared_config)) + rdd.foreach(lambda worker_id: _count_data(monitor_options, table_name, index_name, filter_expression, expression_values, expression_names, worker_id, parallelize_count, total_matched_accumulator, error_accumulator, rate_limiter_shared_config, per_segment_accumulator)) rdd.count() except Exception as e: raise Exception(f"Error in parallel execution: {get_error_message(e)}") from None @@ -94,7 +108,12 @@ def run(job, spark_context, glue_context, parsed_args): # Print the total records inserted using the accumulator after all tasks complete print(f"Total records counted: {total_matched_accumulator.value:,}") -def _count_data(monitor_options, table_name, index_name, filter_expression, expression_values, expression_names, segment, total_segments, total_matched_accumulator, error_accumulator, rate_limiter_shared_config): + if persegment and per_segment_accumulator is not None: + print("\nPer-segment counts:") + for seg in sorted(per_segment_accumulator.value.keys()): + print(f" Segment {seg}: {per_segment_accumulator.value[seg]:,}") + +def _count_data(monitor_options, table_name, index_name, filter_expression, expression_values, expression_names, segment, total_segments, total_matched_accumulator, error_accumulator, rate_limiter_shared_config, per_segment_accumulator=None): rate_limiter_worker = RateLimiterWorker( shared_config=rate_limiter_shared_config, @@ -145,4 +164,6 @@ def _count_data(monitor_options, table_name, index_name, filter_expression, expr print(f"Worker {segment}/{total_segments} counted {local_count} records.") total_matched_accumulator.add(local_count) + if per_segment_accumulator is not None: + per_segment_accumulator.add({segment: local_count}) return local_count