HyukjinKwon commented on code in PR #58990:
URL: https://github.com/apache/spark/pull/58990#discussion_r4139551222
##########
python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py:
##########
@@ -322,6 +366,168 @@ def sum_udf(col):
self.assertEqual(out[0].column("_0").to_pylist(), [10, 30])
[email protected](not have_pyarrow, pyarrow_requirement_message)
+class ArrowGroupedAggIncrementalPartialUDFHandlerTests(unittest.TestCase):
+ # Input columns are ``[k, v]``: the leading grouping key, then the
aggregator input. Output
+ # is one row per key: ``k_0`` followed by the ``_0`` buffer struct.
+ _KEY_SCHEMA = StructType([StructField("k", LongType())])
+
+ def _handler(self, cap=10000):
+ return ArrowGroupedAggIncrementalPartialUDFHandler(
+ udfs=[(_SumAggregator(), [1], {}, LongType())],
+
runner_conf=RunnerConf({"spark.sql.execution.arrow.maxRecordsPerBatch":
str(cap)}),
+ eval_conf=EvalConf({"grouping_key_schema":
self._KEY_SCHEMA.json()}),
+ )
+
+ @staticmethod
+ def _rows(out):
+ return [
+ (k, buf["s"], buf["c"])
+ for b in out
+ for k, buf in zip(b.column("k_0").to_pylist(),
b.column("_0").to_pylist())
+ ]
+
+ def test_combines_rows_per_key_across_batches(self):
+ out = list(
+ self._handler().run(0, iter([_batch(k=[1, 2, 1], v=[1, 2, 3]),
_batch(k=[2], v=[4])]))
+ )
+ self.assertEqual(len(out), 1)
+ self.assertEqual(sorted(self._rows(out)), [(1, 4, 2), (2, 6, 2)])
+
+ def test_empty_partition_emits_nothing(self):
+ self.assertEqual(list(self._handler().run(0, iter([]))), [])
+
+ def test_flushes_when_key_count_reaches_cap(self):
+ # cap=2: the check runs after each input batch. The first batch brings
the map to exactly
+ # 2 keys, so it flushes (>= cap); key 1 then reappears in a fresh map,
and the FINAL stage
+ # merges the duplicate downstream. A mid-stream flush is not chunked,
so the second batch
+ # flushes all 3 of its keys at once.
+ out = list(
+ self._handler(cap=2).run(
+ 0, iter([_batch(k=[1, 2], v=[1, 2]), _batch(k=[1, 3, 4],
v=[10, 3, 4])])
+ )
+ )
+ self.assertEqual([b.num_rows for b in out], [2, 3])
+ self.assertEqual(sorted(self._rows(out[:1])), [(1, 1, 1), (2, 2, 1)])
+ self.assertEqual(sorted(self._rows(out[1:])), [(1, 10, 1), (3, 3, 1),
(4, 4, 1)])
+
+ def test_below_cap_is_not_flushed(self):
+ out = list(
+ self._handler(cap=3).run(0, iter([_batch(k=[1, 2], v=[1, 2]),
_batch(k=[1], v=[5])]))
+ )
+ self.assertEqual([b.num_rows for b in out], [2])
+ self.assertEqual(sorted(self._rows(out)), [(1, 6, 2), (2, 2, 1)])
+
+ def test_non_positive_cap_is_unbounded(self):
+ out = list(self._handler(cap=0).run(0, iter([_batch(k=[1, 2, 3], v=[1,
2, 3])])))
+ self.assertEqual([b.num_rows for b in out], [3])
+
+ def test_nan_keys_are_combined(self):
+ batch = pa.RecordBatch.from_arrays(
+ [pa.array([float("nan"), float("nan")]), pa.array([1, 2],
type=pa.int64())],
+ ["k", "v"],
+ )
+ handler = ArrowGroupedAggIncrementalPartialUDFHandler(
+ udfs=[(_SumAggregator(), [1], {}, LongType())],
+ runner_conf=RunnerConf({}),
+ eval_conf=EvalConf(
+ {"grouping_key_schema": StructType([StructField("k",
DoubleType())]).json()}
+ ),
+ )
+ out = list(handler.run(0, iter([batch])))
+ self.assertEqual(out[0].num_rows, 1)
+ self.assertEqual(out[0].column("_0").to_pylist(), [{"s": 3, "c": 2}])
+
+
[email protected](not have_pyarrow, pyarrow_requirement_message)
+class ArrowGroupedAggIncrementalFinalUDFHandlerTests(unittest.TestCase):
+ # Each group's rows carry one partial-buffer struct column; the handler
merges them and
+ # emits one ``finish`` value per group.
+ _BUFFER_TYPE = pa.struct([("s", pa.int64()), ("c", pa.int64())]) if
have_pyarrow else None
+
+ def _buffers(self, buffers):
+ return pa.RecordBatch.from_arrays([pa.array(buffers,
type=self._BUFFER_TYPE)], ["buf"])
+
+ def _handler(self):
+ return ArrowGroupedAggIncrementalFinalUDFHandler(
+ udfs=[(_SumAggregator(), [0], {}, LongType())],
+ runner_conf=RunnerConf({}),
+ eval_conf=None,
+ )
+
+ def test_merges_partial_buffers_per_group(self):
+ groups = iter(
+ [
+ iter([self._buffers([{"s": 1, "c": 1}]), self._buffers([{"s":
2, "c": 1}])]),
+ iter([self._buffers([{"s": 10, "c": 2}])]),
+ ]
+ )
+ out = list(self._handler().run(0, groups))
+ self.assertEqual([b.column("_0").to_pylist() for b in out], [[3],
[10]])
+
+ def test_null_buffers_are_skipped(self):
+ # An all-null buffer group (the operator's empty global aggregation)
finishes ``zero``.
+ out = list(self._handler().run(0, _one_group(self._buffers([None]))))
+ self.assertEqual(out[0].column("_0").to_pylist(), [0])
+ out = list(self._handler().run(0, _one_group(self._buffers([None,
{"s": 5, "c": 1}]))))
+ self.assertEqual(out[0].column("_0").to_pylist(), [5])
+
+
[email protected](not have_pyarrow, pyarrow_requirement_message)
+class ArrowWindowAggIncrementalUDFHandlerTests(unittest.TestCase):
+ def _handler(self, agg, bound_type, args_offsets):
+ return ArrowWindowAggIncrementalUDFHandler(
+ udfs=[(agg, args_offsets, {}, LongType())],
+ runner_conf=RunnerConf({"window_bound_types": bound_type}),
+ eval_conf=None,
+ )
+
+ def test_unbounded_frame_repeats_one_value(self):
+ handler = self._handler(_SumAggregator(), "unbounded", [0])
+ out = list(handler.run(0, _one_group(_batch(v=[1, 2]), _batch(v=[3]))))
+ self.assertEqual(out[0].column("_0").to_pylist(), [6, 6, 6])
+
+ def test_growing_frame_extends_running_buffer(self):
+ # rowsBetween(unboundedPreceding, currentRow): same lower bound,
growing upper bound,
+ # so each row folds only its newly-included row.
+ agg = _SumAggregator()
+ handler = self._handler(agg, "bounded", [0, 1, 2])
+ batch = _batch(begin=[0, 0, 0, 0], end=[1, 2, 3, 4], v=[1, 2, 3, 4])
+ out = list(handler.run(0, _one_group(batch)))
+ self.assertEqual(out[0].column("_0").to_pylist(), [1, 3, 6, 10])
+ self.assertEqual(agg.reduce_calls, 4)
+
+ def test_sliding_frame_refolds_from_zero(self):
+ # rowsBetween(-1, currentRow): the lower bound advances, so each frame
refolds.
Review Comment:
**Nit (P3):** For `rowsBetween(-1, currentRow)` on this fixture, `begin=[0,
0, 1, 2]`. Row 1 still starts at 0 and only grows on the right, so it takes the
extend branch. Only rows 2 and 3 refold from `zero`, which is why the expected
count is `1 + 1 + 2 + 2` rather than 7. The assertion is right, but the comment
says every frame refolds. Maybe something like:
```suggestion
# rowsBetween(-1, currentRow): row 1 still starts at 0 and extends;
rows 2-3 refold.
```
**Verification:**
- **Inspection:** The corrected comment matches the fixture's lower bounds
and the asserted reduce-call count.
- **Regression:** The incremental window handler unit tests still pass with
unchanged assertions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
