[bulk_executor] fix: add --persegment support to scancount (#92)

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

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: "<segment_number>: <count>" or "Segment <n>: <count>"
+_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)"
+            )