Small fixes, all tests passing

duckdb · Mytherin · Jun 3, 2024 · May 21, 2024 · May 22, 2024 · May 22, 2024
commit 0e3e7658fbf3d5b6d1da13cf739712c6c1ec559d
diff --git a/src/execution/operator/csv_scanner/scanner/base_scanner.cpp b/src/execution/operator/csv_scanner/scanner/base_scanner.cpp
@@ -13,7 +13,7 @@ BaseScanner::BaseScanner(shared_ptr<CSVBufferManager> buffer_manager_p, shared_p
                          shared_ptr<CSVErrorHandler> error_handler_p, bool sniffing_p,
                          shared_ptr<CSVFileScan> csv_file_scan_p, CSVIterator iterator_p)
     : csv_file_scan(std::move(csv_file_scan_p)), sniffing(sniffing_p), error_handler(std::move(error_handler_p)),
-      state_machine(std::move(state_machine_p)), iterator(iterator_p), buffer_manager(std::move(buffer_manager_p)) {
+      state_machine(std::move(state_machine_p)), buffer_manager(std::move(buffer_manager_p)), iterator(iterator_p) {
 	D_ASSERT(buffer_manager);
 	D_ASSERT(state_machine);
 	// Initialize current buffer handle
@@ -41,14 +41,15 @@ bool BaseScanner::FinishedFile() {
 	return iterator.pos.buffer_pos + 1 == cur_buffer_handle->actual_size;
 }
 
-void BaseScanner::SkipCSVRows(idx_t rows_to_skip) {
+CSVIterator BaseScanner::SkipCSVRows(shared_ptr<CSVBufferManager> buffer_manager,
+                                     const shared_ptr<CSVStateMachine> &state_machine, idx_t rows_to_skip) {
 	if (rows_to_skip == 0) {
-		return;
+		return {};
 	}
-	SkipScanner row_skipper(buffer_manager, state_machine, error_handler, rows_to_skip);
+	auto error_handler = make_shared_ptr<CSVErrorHandler>();
+	SkipScanner row_skipper(std::move(buffer_manager), state_machine, error_handler, rows_to_skip);
 	row_skipper.ParseChunk();
-	iterator.pos = row_skipper.GetIteratorPosition();
-	lines_read += row_skipper.GetLinesRead();
+	return row_skipper.GetIterator();
 }
 
 CSVIterator &BaseScanner::GetIterator() {

diff --git a/src/execution/operator/csv_scanner/scanner/column_count_scanner.cpp b/src/execution/operator/csv_scanner/scanner/column_count_scanner.cpp
@@ -50,9 +50,11 @@ ColumnCountScanner::ColumnCountScanner(shared_ptr<CSVBufferManager> buffer_manag
 }
 
 unique_ptr<StringValueScanner> ColumnCountScanner::UpgradeToStringValueScanner() {
-	auto scanner = make_uniq<StringValueScanner>(0U, buffer_manager, state_machine, error_handler, nullptr, true);
-	scanner->SkipCSVRows(scanner->state_machine->dialect_options.skip_rows.GetValue());
-	return scanner;
+	auto iterator = SkipCSVRows(buffer_manager, state_machine, state_machine->dialect_options.skip_rows.GetValue());
+	if (iterator.done) {
+		return make_uniq<StringValueScanner>(0U, buffer_manager, state_machine, error_handler, nullptr, true);
+	}
+	return make_uniq<StringValueScanner>(0U, buffer_manager, state_machine, error_handler, nullptr, true, iterator);
 }
 
 ColumnCountResult &ColumnCountScanner::ParseChunk() {

diff --git a/src/execution/operator/csv_scanner/scanner/scanner_boundary.cpp b/src/execution/operator/csv_scanner/scanner/scanner_boundary.cpp
@@ -92,7 +92,7 @@ idx_t CSVIterator::GetFileIdx() const {
 }
 
 idx_t CSVIterator::GetBufferIdx() const {
-	return boundary.buffer_idx;
+	return pos.buffer_idx;
 }
 
 idx_t CSVIterator::GetBoundaryIdx() const {

diff --git a/src/execution/operator/csv_scanner/scanner/string_value_scanner.cpp b/src/execution/operator/csv_scanner/scanner/string_value_scanner.cpp
@@ -111,7 +111,7 @@ StringValueResult::StringValueResult(CSVStates &states, CSVStateMachine &state_m
 	if (iterator.first_one) {
 		lines_read +=
 		    state_machine.dialect_options.skip_rows.GetValue() + state_machine.dialect_options.header.GetValue();
-		if (lines_read == 0){
+		if (lines_read == 0) {
 			SkipBOM();
 		}
 	}
@@ -948,7 +948,6 @@ void StringValueScanner::Flush(DataChunk &insert_chunk) {
 
 void StringValueScanner::Initialize() {
 	states.Initialize();
-
 	if (result.result_size != 1 && !(sniffing && state_machine->options.null_padding &&
 	                                 !state_machine->options.dialect_options.skip_rows.IsSetByUser())) {
 		SetStart();
@@ -1233,8 +1232,8 @@ bool StringValueScanner::MoveToNextBuffer() {
 }
 
 void StringValueResult::SkipBOM() {
-	if (buffer_size >= 3 && buffer_ptr[0] == '\xEF' && buffer_ptr[1] == '\xBB' &&
-	    buffer_ptr[2] == '\xBF' && iterator.pos.buffer_pos == 0) {
+	if (buffer_size >= 3 && buffer_ptr[0] == '\xEF' && buffer_ptr[1] == '\xBB' && buffer_ptr[2] == '\xBF' &&
+	    iterator.pos.buffer_pos == 0) {
 		iterator.pos.buffer_pos = 3;
 	}
 }

diff --git a/src/execution/operator/csv_scanner/sniffer/dialect_detection.cpp b/src/execution/operator/csv_scanner/sniffer/dialect_detection.cpp
@@ -326,10 +326,10 @@ void CSVSniffer::SkipLines(vector<unique_ptr<ColumnCountScanner>> &csv_state_mac
 	auto &first_scanner = *csv_state_machines[0];
 	// We figure out the iterator position for the first scanner
 	if (options.dialect_options.skip_rows.IsSetByUser()) {
-		first_scanner.SkipCSVRows(options.dialect_options.skip_rows.GetValue());
 		// The iterator position is the same regardless of the scanner configuration, hence we apply the same iterator
 		// To the remaining scanners
-		const auto first_iterator = first_scanner.GetIterator();
+		const auto first_iterator = BaseScanner::SkipCSVRows(first_scanner.buffer_manager, first_scanner.state_machine,
+		                                                     options.dialect_options.skip_rows.GetValue());
 		for (idx_t i = 1; i < csv_state_machines.size(); i++) {
 			auto &cur_scanner = *csv_state_machines[i];
 			cur_scanner.SetIterator(first_iterator);

diff --git a/src/execution/operator/csv_scanner/sniffer/type_detection.cpp b/src/execution/operator/csv_scanner/sniffer/type_detection.cpp
@@ -336,7 +336,6 @@ void CSVSniffer::DetectTypes() {
 
 		// Reset candidate for parsing
 		auto candidate = candidate_cc->UpgradeToStringValueScanner();
-
 		// Parse chunk and read csv with info candidate
 		auto &data_chunk = candidate->ParseChunk().ToChunk();
 		idx_t row_idx = 0;

diff --git a/src/include/duckdb/execution/operator/csv_scanner/base_scanner.hpp b/src/include/duckdb/execution/operator/csv_scanner/base_scanner.hpp
@@ -95,9 +95,13 @@ class BaseScanner {
 
 	bool ever_quoted = false;
 
+	//! Shared pointer to the buffer_manager, this is shared across multiple scanners
+	shared_ptr<CSVBufferManager> buffer_manager;
+
 	//! Skips Notes and/or parts of the data, starting from the top.
 	//! notes are dirty lines on top of the file, before the actual data
-	void SkipCSVRows(idx_t rows_to_skip);
+	static CSVIterator SkipCSVRows(shared_ptr<CSVBufferManager> buffer_manager,
+	                               const shared_ptr<CSVStateMachine> &state_machine, idx_t rows_to_skip);
 
 protected:
 	//! Boundaries of this scanner
@@ -110,9 +114,6 @@ class BaseScanner {
 	//! Hold the current buffer ptr
 	char *buffer_handle_ptr = nullptr;
 
-	//! Shared pointer to the buffer_manager, this is shared across multiple scanners
-	shared_ptr<CSVBufferManager> buffer_manager;
-
 	//! If this scanner has been initialized
 	bool initialized = false;
 	//! How many lines were read by this scanner

diff --git a/test/sql/copy/csv/test_skip.test b/test/sql/copy/csv/test_skip.test
@@ -68,7 +68,7 @@ SELECT * EXCLUDE (prompt) from sniff_csv('__TEST_DIR__/skip.csv',skip=3000)
 query IIIIIIIIII
 SELECT * EXCLUDE (prompt) from sniff_csv('__TEST_DIR__/skip.csv',skip=11000)
 ----
-,	"	"	\n	11000	0	[{'name': column0, 'type': VARCHAR}]	NULL	NULL	skip=11000
+,	"	"	\n	11000	0	[{'name': column0, 'type': BIGINT}]	NULL	NULL	skip=11000
 
 # Test with different buffer sizes