From 844133d2f599d3f2f93427facf34aa8ef1fe08c3 Mon Sep 17 00:00:00 2001 From: siddiqui irshad Date: Thu, 29 Oct 2026 01:34:29 +0530 Subject: [PATCH] GH-55295: [C++][Dataset] Reject negative num_rows in Scanner::Head Fixes #55295 Signed-off-by: siddiqui irshad --- cpp/src/arrow/dataset/scanner.cc | 3 +++ cpp/src/arrow/dataset/scanner_test.cc | 1 + python/pyarrow/tests/test_dataset.py | 6 ++++++ 3 files changed, 10 insertions(+) diff --git a/cpp/src/arrow/dataset/scanner.cc b/cpp/src/arrow/dataset/scanner.cc index 222e1323d4a8..f743b0d2a87b 100644 --- a/cpp/src/arrow/dataset/scanner.cc +++ b/cpp/src/arrow/dataset/scanner.cc @@ -593,6 +593,9 @@ Result> AsyncScanner::TakeRows(const Array& indices) { } Result> AsyncScanner::Head(int64_t num_rows) { + if (num_rows < 0) { + return Status::Invalid("num_rows must be non-negative"); + } if (num_rows == 0) { return Table::FromRecordBatches(options()->projected_schema, {}); } diff --git a/cpp/src/arrow/dataset/scanner_test.cc b/cpp/src/arrow/dataset/scanner_test.cc index a45847b49def..7ab5ded5df01 100644 --- a/cpp/src/arrow/dataset/scanner_test.cc +++ b/cpp/src/arrow/dataset/scanner_test.cc @@ -1634,6 +1634,7 @@ TEST_P(TestScanner, Head) { ASSERT_OK_AND_ASSIGN(expected, Table::FromRecordBatches(schema_, {})); ASSERT_OK_AND_ASSIGN(actual, scanner->Head(0)); AssertTablesEqual(*expected, *actual, /*same_chunk_layout=*/false); + ASSERT_RAISES(Invalid, scanner->Head(-1)); ASSERT_OK_AND_ASSIGN(expected, Table::FromRecordBatches(schema_, {batch})); ASSERT_OK_AND_ASSIGN(actual, scanner->Head(batch_size)); diff --git a/python/pyarrow/tests/test_dataset.py b/python/pyarrow/tests/test_dataset.py index 9ddfdfdc9c49..b272e3ce3c66 100644 --- a/python/pyarrow/tests/test_dataset.py +++ b/python/pyarrow/tests/test_dataset.py @@ -648,7 +648,13 @@ def test_head(dataset, dataset_reader): result = dataset_reader.head(dataset, 1024, columns=['i64']).to_pydict() assert result == {'i64': list(range(5)) * 2} + with pytest.raises(ValueError, match="num_rows must be non-negative"): + dataset_reader.head(dataset, -1) + fragment = next(dataset.get_fragments()) + with pytest.raises(ValueError, match="num_rows must be non-negative"): + fragment.head(-1) + result = fragment.head(1, columns=['i64']).to_pydict() assert result == {'i64': [0]}