(tempdir)
| 4339 | @pytest.mark.parquet |
| 4340 | @pytest.mark.threading |
| 4341 | def test_write_dataset_with_backpressure(tempdir): |
| 4342 | consumer_gate = threading.Event() |
| 4343 | |
| 4344 | # A filesystem that blocks all writes so that we can build |
| 4345 | # up backpressure. The writes are released at the end of |
| 4346 | # the test. |
| 4347 | class GatingFs(ProxyHandler): |
| 4348 | def open_output_stream(self, path, metadata): |
| 4349 | # Block until the end of the test |
| 4350 | consumer_gate.wait() |
| 4351 | return self._fs.open_output_stream(path, metadata=metadata) |
| 4352 | gating_fs = fs.PyFileSystem(GatingFs(fs.LocalFileSystem())) |
| 4353 | |
| 4354 | schema = pa.schema([pa.field('data', pa.int32())]) |
| 4355 | # The scanner should queue ~ 8Mi rows (~8 batches) but due to ARROW-16258 |
| 4356 | # it always queues 32 batches. |
| 4357 | batch = pa.record_batch([pa.array(list(range(1_000_000)))], schema=schema) |
| 4358 | batches_read = 0 |
| 4359 | min_backpressure = 32 |
| 4360 | end = 200 |
| 4361 | keep_going = True |
| 4362 | |
| 4363 | def counting_generator(): |
| 4364 | nonlocal batches_read |
| 4365 | while batches_read < end: |
| 4366 | if not keep_going: |
| 4367 | return |
| 4368 | time.sleep(0.01) |
| 4369 | batches_read += 1 |
| 4370 | yield batch |
| 4371 | |
| 4372 | scanner = ds.Scanner.from_batches( |
| 4373 | counting_generator(), schema=schema, use_threads=True) |
| 4374 | |
| 4375 | write_thread = threading.Thread( |
| 4376 | target=lambda: ds.write_dataset( |
| 4377 | scanner, str(tempdir), format='parquet', filesystem=gating_fs)) |
| 4378 | write_thread.start() |
| 4379 | |
| 4380 | try: |
| 4381 | start = time.time() |
| 4382 | |
| 4383 | def duration(): |
| 4384 | return time.time() - start |
| 4385 | |
| 4386 | # This test is timing dependent. There is no signal from the C++ |
| 4387 | # when backpressure has been hit. We don't know exactly when |
| 4388 | # backpressure will be hit because it may take some time for the |
| 4389 | # signal to get from the sink to the scanner. |
| 4390 | # |
| 4391 | # The test may emit false positives on slow systems. It could |
| 4392 | # theoretically emit a false negative if the scanner managed to read |
| 4393 | # and emit all 200 batches before the backpressure signal had a chance |
| 4394 | # to propagate but the 0.01s delay in the generator should make that |
| 4395 | # scenario unlikely. |
| 4396 | last_value = 0 |
| 4397 | backpressure_probably_hit = False |
| 4398 | while duration() < 10: |
nothing calls this directly
no test coverage detected