Add missing stream synchronizations to various tests (#21122)

I noticed these issues when investigating #21094. I suspect that this was the underlying issue behind #19900.

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Matthew Murray (https://github.com/Matt711)

Approvers:
  - Matthew Murray (https://github.com/Matt711)

URL: #21122

Author: vyasr

python/pylibcudf/tests/common/utils.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
@@ -16,6 +16,20 @@
 from pylibcudf.io.types import CompressionType
 
 
+def synchronize_stream(stream=None):
+    """Synchronize a stream, handling both explicit streams and None (default stream).
+
+    Parameters
+    ----------
+    stream : Stream or None
+        The stream to synchronize. If None, synchronizes the default stream.
+    """
+    if stream is None:
+        plc.utils.DEFAULT_STREAM.synchronize()
+    else:
+        stream.synchronize()
+
+
 def metadata_from_arrow_type(
     pa_type: pa.Array,
     name: str = "",

python/pylibcudf/tests/io/test_avro.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 import io
@@ -7,7 +7,7 @@
 import fastavro
 import pyarrow as pa
 import pytest
-from utils import assert_table_and_meta_eq
+from utils import assert_table_and_meta_eq, synchronize_stream
 
 from rmm.pylibrmm.device_buffer import DeviceBuffer
 from rmm.pylibrmm.stream import Stream
@@ -156,6 +156,8 @@ def test_read_avro_from_device_buffers(avro_dtypes, avro_dtype_data, stream):
     buf = buffer.getbuffer()
     device_buf = DeviceBuffer.to_device(buf, plc.utils._get_stream(stream))
 
+    synchronize_stream(stream)
+
     options = plc.io.avro.AvroReaderOptions.builder(
         plc.io.types.SourceInfo([device_buf])
     ).build()

python/pylibcudf/tests/io/test_csv.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 import io
 import os
@@ -11,6 +11,7 @@
     assert_table_and_meta_eq,
     make_source,
     sink_to_str,
+    synchronize_stream,
     write_source_str,
 )
 
@@ -314,6 +315,8 @@ def test_read_csv_from_device_buffers(csv_table_data, stream):
         csv_string.encode("utf-8"), plc.utils._get_stream(stream)
     )
 
+    synchronize_stream(stream)
+
     options = plc.io.csv.CsvReaderOptions.builder(
         plc.io.SourceInfo([buf])
     ).build()
@@ -379,6 +382,8 @@ def test_write_csv(
         stream,
     )
 
+    synchronize_stream(stream)
+
     # Convert everything to string to make comparisons easier
     str_result = sink_to_str(sink)
 
@@ -423,6 +428,8 @@ def test_write_csv_na_rep(na_rep):
         )
     )
 
+    synchronize_stream()
+
     # Convert everything to string to make comparisons easier
     str_result = sink_to_str(sink)
 

python/pylibcudf/tests/io/test_experimental_hybrid_scan.py

@@ -1,10 +1,11 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 import io
 
 import pyarrow as pa
 import pyarrow.parquet as pq
 import pytest
+from utils import synchronize_stream
 
 from rmm import DeviceBuffer
 from rmm.pylibrmm.stream import Stream
@@ -327,6 +328,8 @@ def test_hybrid_scan_materialize_columns(
         for r in filter_ranges
     ]
 
+    synchronize_stream(stream)
+
     # Materialize filter columns (mr is optional, defaults to None)
     filter_result = simple_hybrid_scan_reader.materialize_filter_columns(
         filtered_row_groups,
@@ -337,6 +340,8 @@ def test_hybrid_scan_materialize_columns(
         stream,
     )
 
+    synchronize_stream(stream)
+
     # Filter column should have 1 column, with rows passing the filter
     expected_result_rows = num_rows - filter_threshold
     assert filter_result.tbl.num_columns() == 1
@@ -356,6 +361,8 @@ def test_hybrid_scan_materialize_columns(
         for r in payload_ranges
     ]
 
+    synchronize_stream(stream)
+
     # Materialize payload columns (mr is optional, defaults to None)
     payload_result = simple_hybrid_scan_reader.materialize_payload_columns(
         filtered_row_groups,
@@ -366,6 +373,8 @@ def test_hybrid_scan_materialize_columns(
         stream,
     )
 
+    synchronize_stream(stream)
+
     assert payload_result.tbl.num_columns() == 2
     assert payload_result.tbl.num_rows() == expected_result_rows
 
@@ -379,6 +388,8 @@ def test_hybrid_scan_materialize_columns(
     comparison_options.set_filter(filter_expression)
     expected_result = plc.io.parquet.read_parquet(comparison_options, stream)
 
+    synchronize_stream(stream)
+
     # Combine hybrid scan results
     hybrid_columns = filter_result.tbl.columns() + payload_result.tbl.columns()
     hybrid_table = plc.Table(hybrid_columns)
@@ -437,6 +448,8 @@ def test_hybrid_scan_has_next_table_chunk(
         for r in filter_ranges
     ]
 
+    synchronize_stream()
+
     # Setup chunking first
     simple_hybrid_scan_reader.setup_chunking_for_filter_columns(
         512,  # chunk_read_limit
@@ -503,6 +516,8 @@ def test_hybrid_scan_chunked_reading(
         for r in filter_ranges
     ]
 
+    synchronize_stream(stream)
+
     # Setup chunking for filter columns with small chunk size
     chunk_read_limit = 512  # Small limit to force multiple chunks
     pass_read_limit = 0  # No limit

python/pylibcudf/tests/io/test_json.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 import io
 
@@ -9,6 +9,7 @@
     assert_table_and_meta_eq,
     make_source,
     sink_to_str,
+    synchronize_stream,
     write_source_str,
 )
 
@@ -44,6 +45,8 @@ def test_write_json_basic(
 
     plc.io.json.write_json(options, stream)
 
+    synchronize_stream(stream)
+
     exp = pa_table.to_pandas()
 
     # Convert everything to string to make
@@ -82,6 +85,8 @@ def test_write_json_nulls(na_rep, include_nulls):
 
     plc.io.json.write_json(options)
 
+    synchronize_stream()
+
     exp = pa_tbl.to_pandas()
 
     # Convert everything to string to make
@@ -133,6 +138,8 @@ def test_write_json_bool_opts(true_value, false_value):
 
     plc.io.json.write_json(options)
 
+    synchronize_stream()
+
     exp = pa_tbl.to_pandas()
 
     # Convert everything to string to make
@@ -428,6 +435,8 @@ def test_read_json_from_device_buffers(table_data, num_buffers, stream):
         json_str.encode("utf-8"), plc.utils._get_stream(stream)
     )
 
+    synchronize_stream(stream)
+
     options = (
         plc.io.json.JsonReaderOptions.builder(
             plc.io.SourceInfo([buf] * num_buffers)
@@ -471,6 +480,8 @@ def test_utf8_escaped_json_writer(tmp_path):
     )
     plc.io.json.write_json(options)
 
+    synchronize_stream()
+
     output_string = path.read_text(encoding="utf-8").strip()
 
     assert output_string == '[{"0":"C𝞵𝓓𝒻"}]'

python/pylibcudf/tests/io/test_orc.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 import pyarrow as pa
@@ -8,6 +8,7 @@
     assert_table_and_meta_eq,
     get_bytes_from_source,
     make_source,
+    synchronize_stream,
 )
 
 from rmm.pylibrmm.device_buffer import DeviceBuffer
@@ -107,6 +108,8 @@ def test_read_orc_from_device_buffers(
         get_bytes_from_source(source), plc.utils._get_stream(stream)
     )
 
+    synchronize_stream(stream)
+
     options = plc.io.orc.OrcReaderOptions.builder(
         plc.io.types.SourceInfo([buf] * num_buffers)
     ).build()
@@ -179,6 +182,8 @@ def test_roundtrip_pa_table(
 
     plc.io.orc.write_orc(options, stream)
 
+    synchronize_stream(stream)
+
     read_table = pa.orc.read_table(str(tmpfile_name))
 
     res = plc.io.types.TableWithMetadata(

python/pylibcudf/tests/io/test_parquet.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 import io
 
@@ -10,6 +10,7 @@
     assert_table_and_meta_eq,
     get_bytes_from_source,
     make_source,
+    synchronize_stream,
 )
 
 from rmm.pylibrmm.device_buffer import DeviceBuffer
@@ -192,6 +193,8 @@ def test_read_parquet_from_device_buffers(
         get_bytes_from_source(source), plc.utils._get_stream(stream)
     )
 
+    synchronize_stream(stream)
+
     options = plc.io.parquet.ParquetReaderOptions.builder(
         plc.io.SourceInfo([buf] * num_buffers)
     ).build()
@@ -289,6 +292,9 @@ def test_write_parquet(
         options.set_max_dictionary_size(max_dictionary_size)
 
     result = plc.io.parquet.write_parquet(options, stream)
+
+    synchronize_stream(stream)
+
     assert isinstance(result, memoryview)