[zephyr] Fix PR review feedback: streaming reads, secondary sort, vectorized boundaries

- Stream Arrow run files via iter_batches() instead of whole-table read to
  avoid OOM on large shards that trigger external sort
- Add promote_options="default" to concat_tables in external sort for
  schema-evolved items
- Include _sort_secondary in external sort keys and merge key so within-group
  order is preserved
- Vectorize _find_group_boundaries using pc.not_equal diff instead of
  per-element .as_py() calls — matters for high-cardinality keys
- Store has_sort as a field on _ShardBuffer set at construction instead of
  fragile sorts[0] detection
- Add assertion guard in _arrow_reduce_gen for pickled shards

Author: rjpower

lib/zephyr/src/zephyr/external_sort.py

@@ -193,7 +193,7 @@ def external_sort_merge_arrow(
         batch_tables = list(islice(chunk_tables_gen, EXTERNAL_SORT_FAN_IN))
         if not batch_tables:
             break
-        combined = pa.concat_tables(batch_tables)
+        combined = pa.concat_tables(batch_tables, promote_options="default")
         indices = pc.sort_indices(combined, sort_keys=sort_keys)
         sorted_table = combined.take(indices)
 
@@ -214,8 +214,9 @@ def external_sort_merge_arrow(
         return
 
     def _read_run(path: str) -> Iterator:
-        table = pq.read_table(path)
-        yield from table.to_pylist()
+        pf = pq.ParquetFile(path)
+        for batch in pf.iter_batches():
+            yield from batch.to_pylist()
 
     run_iters = [_read_run(p) for p in run_paths]
     try:

lib/zephyr/src/zephyr/plan.py

@@ -179,16 +179,27 @@ def _flatmap_gen(stream: Iterator, fn: Callable) -> Iterator:
 
 
 def _find_group_boundaries(key_col: pa.ChunkedArray) -> Iterator[tuple[int, int, Any]]:
-    """Yield (start, end, key_value) for each contiguous group in a sorted key column."""
+    """Yield (start, end, key_value) for each contiguous group in a sorted key column.
+
+    Uses Arrow compute to find boundaries vectorized instead of per-element
+    Python scalar extraction, which matters for high-cardinality keys.
+    """
     arr = key_col.combine_chunks()
     n = len(arr)
     if n == 0:
         return
+    if n == 1:
+        yield (0, 1, arr[0].as_py())
+        return
+
+    # Vectorized boundary detection: compare adjacent elements
+    ne_mask = pc.not_equal(arr[:-1], arr[1:])
+    boundary_indices = pc.filter(pa.array(range(1, n), type=pa.int64()), ne_mask).to_pylist()
+
     prev = 0
-    for i in range(1, n):
-        if arr[i].as_py() != arr[prev].as_py():
-            yield (prev, i, arr[prev].as_py())
-            prev = i
+    for idx in boundary_indices:
+        yield (prev, idx, arr[prev].as_py())
+        prev = idx
     yield (prev, n, arr[prev].as_py())
 
 
@@ -223,8 +234,13 @@ def _arrow_reduce_gen(
         ScatterShard,
         _ZEPHYR_SHUFFLE_ITEM_COL,
         _ZEPHYR_SHUFFLE_SORT_KEY_COL,
+        _ZEPHYR_SHUFFLE_SORT_SECONDARY_COL,
     )
 
+    assert not any(
+        it.is_pickled for it in shard.iterators
+    ), "_arrow_reduce_gen requires non-pickled items; use _reduce_gen which routes pickled shards to the Python path"
+
     use_external = (
         external_sort_dir is not None
         and isinstance(shard, ScatterShard)
@@ -233,6 +249,10 @@ def _arrow_reduce_gen(
 
     if use_external:
         sort_keys: list[tuple[str, str]] = [(_ZEPHYR_SHUFFLE_SORT_KEY_COL, "ascending")]
+        # Peek at the first chunk table to check for secondary sort column
+        first_tables = list(islice((t for it in shard.iterators for t in it.get_chunk_tables()), 1))
+        if first_tables and _ZEPHYR_SHUFFLE_SORT_SECONDARY_COL in first_tables[0].column_names:
+            sort_keys.append((_ZEPHYR_SHUFFLE_SORT_SECONDARY_COL, "ascending"))
         logger.info(
             "Arrow external sort triggered for shard with %d iterators, spilling to %s",
             sum(it.chunk_count for it in shard.iterators),
@@ -243,7 +263,11 @@ def _chunk_tables() -> Iterator[pa.Table]:
             for it in shard.iterators:
                 yield from it.get_chunk_tables()
 
+        has_secondary = len(sort_keys) > 1
+
         def _merge_key(row: dict) -> Any:
+            if has_secondary:
+                return (row[_ZEPHYR_SHUFFLE_SORT_KEY_COL], row.get(_ZEPHYR_SHUFFLE_SORT_SECONDARY_COL))
             return row[_ZEPHYR_SHUFFLE_SORT_KEY_COL]
 
         merged_stream = external_sort_merge_arrow(

lib/zephyr/src/zephyr/shuffle.py

@@ -518,6 +518,7 @@ class _ShardBuffer:
 
     shard_idx: int
     pickled: bool = False
+    has_sort: bool = False
     pending: list[tuple[Any, Any, Any | None]] = field(default_factory=list)
     tables: list[pa.RecordBatch] = field(default_factory=list)
     nbytes: int = 0
@@ -534,14 +535,13 @@ def _flush_micro(self) -> None:
         if not self.pending:
             return
         items, keys, sorts = zip(*self.pending, strict=True)
-        has_sort = sorts[0] is not None
         envelope_fn = _make_pickle_envelope if self.pickled else _make_envelope
         enveloped = envelope_fn(
             list(items),
             self.shard_idx,
             self.chunk_idx,
             list(keys),
-            list(sorts) if has_sort else None,
+            list(sorts) if self.has_sort else None,
         )
         batch = pa.RecordBatch.from_pylist(enveloped, schema=self.schema)
         if self.schema is None:
@@ -629,7 +629,7 @@ def _write_parquet_scatter(
 
     def _get_buffer(target: int) -> _ShardBuffer:
         if target not in buffers:
-            buffers[target] = _ShardBuffer(shard_idx=target, pickled=pickled)
+            buffers[target] = _ShardBuffer(shard_idx=target, pickled=pickled, has_sort=sort_fn is not None)
         return buffers[target]
 
     def _flush_pending() -> None:
@@ -694,7 +694,7 @@ def _flush_buffer(buf: _ShardBuffer) -> None:
                 py_items = [pickle.loads(b) for b in py_items]
             combined = _apply_combiner(py_items, key_fn, combiner_fn)
             # Re-create a fresh buffer from combined items
-            combined_buf = _ShardBuffer(shard_idx=buf.shard_idx, pickled=pickled)
+            combined_buf = _ShardBuffer(shard_idx=buf.shard_idx, pickled=pickled, has_sort=sort_fn is not None)
             combined_buf.chunk_idx = buf.chunk_idx
             for item in combined:
                 k = key_fn(item)