diff --git a/cpp/src/arrow/dataset/scanner.cc b/cpp/src/arrow/dataset/scanner.cc index 222e1323d4a8..d6a686e9c597 100644 --- a/cpp/src/arrow/dataset/scanner.cc +++ b/cpp/src/arrow/dataset/scanner.cc @@ -1039,8 +1039,9 @@ Result MakeScanNode(acero::ExecPlan* plan, merged_batch_gen = MakeConcatenatedGenerator(std::move(batch_gen_gen)); } } else { - merged_batch_gen = - MakeMergedGenerator(std::move(batch_gen_gen), scan_options->fragment_readahead); + // Disabling readahead still requires one active fragment. + merged_batch_gen = MakeMergedGenerator(std::move(batch_gen_gen), + std::max(1, scan_options->fragment_readahead)); } AsyncGenerator batch_gen; diff --git a/cpp/src/arrow/dataset/scanner_test.cc b/cpp/src/arrow/dataset/scanner_test.cc index a45847b49def..07fb757314e3 100644 --- a/cpp/src/arrow/dataset/scanner_test.cc +++ b/cpp/src/arrow/dataset/scanner_test.cc @@ -2491,6 +2491,37 @@ TEST(ScanNode, Trivial) { ASSERT_THAT(plan.Run(), Finishes(ResultWith(UnorderedElementsAreArray(expected)))); } +TEST(ScanNode, NoFragmentReadahead) { + for (bool empty : {false, true}) { + for (bool require_sequenced_output : {false, true}) { + SCOPED_TRACE("empty=" + std::to_string(empty) + + ", sequenced=" + std::to_string(require_sequenced_output)); + TestPlan plan; + auto basic = MakeBasicDataset(); + if (empty) { + basic.dataset = + std::make_shared(basic.dataset->schema(), FragmentVector{}); + basic.batches.clear(); + } + auto options = std::make_shared(); + options->projection = Materialize({"a", "b", "c"}, /*include_aug_fields=*/true); + options->fragment_readahead = 0; + + ASSERT_OK(acero::Declaration::Sequence( + {{"scan", + ScanNodeOptions{basic.dataset, options, require_sequenced_output}}, + {"sink", acero::SinkNodeOptions{&plan.sink_gen}}}) + .AddToPlan(plan.get())); + ASSERT_FINISHES_OK_AND_ASSIGN(auto batches, plan.Run()); + if (require_sequenced_output) { + ASSERT_THAT(batches, ::testing::ElementsAreArray(basic.batches)); + } else { + ASSERT_THAT(batches, UnorderedElementsAreArray(basic.batches)); + } + } + } +} + TEST(ScanNode, FilteredOnVirtualColumn) { TestPlan plan; diff --git a/python/pyarrow/tests/test_dataset.py b/python/pyarrow/tests/test_dataset.py index 0a94c0bd9875..d3cec3ef9c58 100644 --- a/python/pyarrow/tests/test_dataset.py +++ b/python/pyarrow/tests/test_dataset.py @@ -565,6 +565,21 @@ def test_scanner_options(dataset): assert batch.num_columns == 7 +@pytest.mark.parquet +@pytest.mark.parametrize("use_threads", [False, True]) +@pytest.mark.parametrize("num_fragments", [0, 3]) +def test_scanner_no_fragment_readahead(tmp_path, use_threads, num_fragments): + schema = pa.schema([("value", pa.int64())]) + for i in range(num_fragments): + pq.write_table(pa.table({"value": [i]}, schema=schema), + tmp_path / f"{i}.parquet") + dataset = ds.dataset(tmp_path, format="parquet", schema=schema) + result = dataset.scanner(fragment_readahead=0, + use_threads=use_threads).to_table() + expected = pa.table({"value": list(range(num_fragments))}, schema=schema) + assert result.equals(expected) + + @pytest.mark.parquet def test_scanner(dataset, dataset_reader): scanner = dataset_reader.scanner(