duckdb 1.1.2-dev4.0 → 1.1.3

package/src/duckdb/src/execution/operator/persistent/csv_buffer.cpp

@@ -0,0 +1,72 @@
+#include "duckdb/execution/operator/persistent/csv_buffer.hpp"
+#include "duckdb/common/string_util.hpp"
+
+namespace duckdb {
+
+CSVBuffer::CSVBuffer(ClientContext &context, idx_t buffer_size_p, CSVFileHandle &file_handle,
+                     idx_t &global_csv_current_position, idx_t file_number_p)
+    : context(context), first_buffer(true), file_number(file_number_p) {
+	this->handle = AllocateBuffer(buffer_size_p);
+
+	auto buffer = Ptr();
+	actual_size = file_handle.Read(buffer, buffer_size_p);
+	global_csv_start = global_csv_current_position;
+	global_csv_current_position += actual_size;
+	if (actual_size >= 3 && buffer[0] == '\xEF' && buffer[1] == '\xBB' && buffer[2] == '\xBF') {
+		start_position += 3;
+	}
+	last_buffer = file_handle.FinishedReading();
+}
+
+CSVBuffer::CSVBuffer(ClientContext &context, BufferHandle buffer_p, idx_t buffer_size_p, idx_t actual_size_p,
+                     bool final_buffer, idx_t global_csv_current_position, idx_t file_number_p)
+    : context(context), handle(std::move(buffer_p)), actual_size(actual_size_p), last_buffer(final_buffer),
+      global_csv_start(global_csv_current_position), file_number(file_number_p) {
+}
+
+unique_ptr<CSVBuffer> CSVBuffer::Next(CSVFileHandle &file_handle, idx_t buffer_size, idx_t &global_csv_current_position,
+                                      idx_t file_number_p) {
+	auto next_buffer = AllocateBuffer(buffer_size);
+	idx_t next_buffer_actual_size = file_handle.Read(next_buffer.Ptr(), buffer_size);
+	if (next_buffer_actual_size == 0) {
+		// We are done reading
+		return nullptr;
+	}
+
+	auto next_csv_buffer =
+	    make_uniq<CSVBuffer>(context, std::move(next_buffer), buffer_size, next_buffer_actual_size,
+	                         file_handle.FinishedReading(), global_csv_current_position, file_number_p);
+	global_csv_current_position += next_buffer_actual_size;
+	return next_csv_buffer;
+}
+
+BufferHandle CSVBuffer::AllocateBuffer(idx_t buffer_size) {
+	auto &buffer_manager = BufferManager::GetBufferManager(context);
+	return buffer_manager.Allocate(MaxValue<idx_t>(Storage::BLOCK_SIZE, buffer_size));
+}
+
+idx_t CSVBuffer::GetBufferSize() {
+	return actual_size;
+}
+
+idx_t CSVBuffer::GetStart() {
+	return start_position;
+}
+
+bool CSVBuffer::IsCSVFileLastBuffer() {
+	return last_buffer;
+}
+
+bool CSVBuffer::IsCSVFileFirstBuffer() {
+	return first_buffer;
+}
+
+idx_t CSVBuffer::GetCSVGlobalStart() {
+	return global_csv_start;
+}
+
+idx_t CSVBuffer::GetFileNumber() {
+	return file_number;
+}
+
+} // namespace duckdb

package/src/duckdb/src/execution/operator/persistent/csv_file_handle.cpp

@@ -0,0 +1,158 @@
+#include "duckdb/execution/operator/persistent/csv_file_handle.hpp"
+
+namespace duckdb {
+
+CSVFileHandle::CSVFileHandle(FileSystem &fs, Allocator &allocator, unique_ptr<FileHandle> file_handle_p,
+                             const string &path_p, FileCompressionType compression, bool enable_reset)
+    : fs(fs), allocator(allocator), file_handle(std::move(file_handle_p)), path(path_p), compression(compression),
+      reset_enabled(enable_reset) {
+	can_seek = file_handle->CanSeek();
+	on_disk_file = file_handle->OnDiskFile();
+	file_size = file_handle->GetFileSize();
+}
+
+unique_ptr<FileHandle> CSVFileHandle::OpenFileHandle(FileSystem &fs, Allocator &allocator, const string &path,
+                                                     FileCompressionType compression) {
+	auto file_handle = fs.OpenFile(path, FileFlags::FILE_FLAGS_READ, FileLockType::NO_LOCK, compression);
+	if (file_handle->CanSeek()) {
+		file_handle->Reset();
+	}
+	return file_handle;
+}
+
+unique_ptr<CSVFileHandle> CSVFileHandle::OpenFile(FileSystem &fs, Allocator &allocator, const string &path,
+                                                  FileCompressionType compression, bool enable_reset) {
+	auto file_handle = CSVFileHandle::OpenFileHandle(fs, allocator, path, compression);
+	return make_uniq<CSVFileHandle>(fs, allocator, std::move(file_handle), path, compression, enable_reset);
+}
+
+bool CSVFileHandle::CanSeek() {
+	return can_seek;
+}
+
+void CSVFileHandle::Seek(idx_t position) {
+	if (!can_seek) {
+		throw InternalException("Cannot seek in this file");
+	}
+	file_handle->Seek(position);
+}
+
+idx_t CSVFileHandle::SeekPosition() {
+	if (!can_seek) {
+		throw InternalException("Cannot seek in this file");
+	}
+	return file_handle->SeekPosition();
+}
+
+void CSVFileHandle::Reset() {
+	requested_bytes = 0;
+	read_position = 0;
+	if (can_seek) {
+		// we can seek - reset the file handle
+		file_handle->Reset();
+	} else if (on_disk_file) {
+		// we cannot seek but it is an on-disk file - re-open the file
+		file_handle = CSVFileHandle::OpenFileHandle(fs, allocator, path, compression);
+	} else {
+		if (!reset_enabled) {
+			throw InternalException("Reset called but reset is not enabled for this CSV Handle");
+		}
+		read_position = 0;
+	}
+}
+bool CSVFileHandle::OnDiskFile() {
+	return on_disk_file;
+}
+
+idx_t CSVFileHandle::FileSize() {
+	return file_size;
+}
+
+bool CSVFileHandle::FinishedReading() {
+	return requested_bytes >= file_size;
+}
+
+idx_t CSVFileHandle::Read(void *buffer, idx_t nr_bytes) {
+	requested_bytes += nr_bytes;
+	if (on_disk_file || can_seek) {
+		// if this is a plain file source OR we can seek we are not caching anything
+		return file_handle->Read(buffer, nr_bytes);
+	}
+	// not a plain file source: we need to do some bookkeeping around the reset functionality
+	idx_t result_offset = 0;
+	if (read_position < buffer_size) {
+		// we need to read from our cached buffer
+		auto buffer_read_count = MinValue<idx_t>(nr_bytes, buffer_size - read_position);
+		memcpy(buffer, cached_buffer.get() + read_position, buffer_read_count);
+		result_offset += buffer_read_count;
+		read_position += buffer_read_count;
+		if (result_offset == nr_bytes) {
+			return nr_bytes;
+		}
+	} else if (!reset_enabled && cached_buffer.IsSet()) {
+		// reset is disabled, but we still have cached data
+		// we can remove any cached data
+		cached_buffer.Reset();
+		buffer_size = 0;
+		buffer_capacity = 0;
+		read_position = 0;
+	}
+	// we have data left to read from the file
+	// read directly into the buffer
+	auto bytes_read = file_handle->Read(char_ptr_cast(buffer) + result_offset, nr_bytes - result_offset);
+	file_size = file_handle->GetFileSize();
+	read_position += bytes_read;
+	if (reset_enabled) {
+		// if reset caching is enabled, we need to cache the bytes that we have read
+		if (buffer_size + bytes_read >= buffer_capacity) {
+			// no space; first enlarge the buffer
+			buffer_capacity = MaxValue<idx_t>(NextPowerOfTwo(buffer_size + bytes_read), buffer_capacity * 2);
+
+			auto new_buffer = allocator.Allocate(buffer_capacity);
+			if (buffer_size > 0) {
+				memcpy(new_buffer.get(), cached_buffer.get(), buffer_size);
+			}
+			cached_buffer = std::move(new_buffer);
+		}
+		memcpy(cached_buffer.get() + buffer_size, char_ptr_cast(buffer) + result_offset, bytes_read);
+		buffer_size += bytes_read;
+	}
+
+	return result_offset + bytes_read;
+}
+
+string CSVFileHandle::ReadLine() {
+	bool carriage_return = false;
+	string result;
+	char buffer[1];
+	while (true) {
+		idx_t bytes_read = Read(buffer, 1);
+		if (bytes_read == 0) {
+			return result;
+		}
+		if (carriage_return) {
+			if (buffer[0] != '\n') {
+				if (!file_handle->CanSeek()) {
+					throw BinderException(
+					    "Carriage return newlines not supported when reading CSV files in which we cannot seek");
+				}
+				file_handle->Seek(file_handle->SeekPosition() - 1);
+				return result;
+			}
+		}
+		if (buffer[0] == '\n') {
+			return result;
+		}
+		if (buffer[0] != '\r') {
+			result += buffer[0];
+		} else {
+			carriage_return = true;
+		}
+	}
+}
+
+void CSVFileHandle::DisableReset() {
+	this->reset_enabled = false;
+}
+
+} // namespace duckdb

package/src/duckdb/src/execution/operator/persistent/csv_reader_options.cpp

@@ -0,0 +1,280 @@
+#include "duckdb/execution/operator/persistent/csv_reader_options.hpp"
+#include "duckdb/common/bind_helpers.hpp"
+#include "duckdb/common/vector_size.hpp"
+#include "duckdb/common/string_util.hpp"
+
+namespace duckdb {
+
+static bool ParseBoolean(const Value &value, const string &loption);
+
+static bool ParseBoolean(const vector<Value> &set, const string &loption) {
+	if (set.empty()) {
+		// no option specified: default to true
+		return true;
+	}
+	if (set.size() > 1) {
+		throw BinderException("\"%s\" expects a single argument as a boolean value (e.g. TRUE or 1)", loption);
+	}
+	return ParseBoolean(set[0], loption);
+}
+
+static bool ParseBoolean(const Value &value, const string &loption) {
+
+	if (value.type().id() == LogicalTypeId::LIST) {
+		auto &children = ListValue::GetChildren(value);
+		return ParseBoolean(children, loption);
+	}
+	if (value.type() == LogicalType::FLOAT || value.type() == LogicalType::DOUBLE ||
+	    value.type().id() == LogicalTypeId::DECIMAL) {
+		throw BinderException("\"%s\" expects a boolean value (e.g. TRUE or 1)", loption);
+	}
+	return BooleanValue::Get(value.DefaultCastAs(LogicalType::BOOLEAN));
+}
+
+static string ParseString(const Value &value, const string &loption) {
+	if (value.IsNull()) {
+		return string();
+	}
+	if (value.type().id() == LogicalTypeId::LIST) {
+		auto &children = ListValue::GetChildren(value);
+		if (children.size() != 1) {
+			throw BinderException("\"%s\" expects a single argument as a string value", loption);
+		}
+		return ParseString(children[0], loption);
+	}
+	if (value.type().id() != LogicalTypeId::VARCHAR) {
+		throw BinderException("\"%s\" expects a string argument!", loption);
+	}
+	return value.GetValue<string>();
+}
+
+static int64_t ParseInteger(const Value &value, const string &loption) {
+	if (value.type().id() == LogicalTypeId::LIST) {
+		auto &children = ListValue::GetChildren(value);
+		if (children.size() != 1) {
+			// no option specified or multiple options specified
+			throw BinderException("\"%s\" expects a single argument as an integer value", loption);
+		}
+		return ParseInteger(children[0], loption);
+	}
+	return value.GetValue<int64_t>();
+}
+
+void BufferedCSVReaderOptions::SetHeader(bool input) {
+	this->header = input;
+	this->has_header = true;
+}
+
+void BufferedCSVReaderOptions::SetCompression(const string &compression_p) {
+	this->compression = FileCompressionTypeFromString(compression_p);
+}
+
+void BufferedCSVReaderOptions::SetEscape(const string &input) {
+	this->escape = input;
+	this->has_escape = true;
+}
+
+void BufferedCSVReaderOptions::SetDelimiter(const string &input) {
+	this->delimiter = StringUtil::Replace(input, "\\t", "\t");
+	this->has_delimiter = true;
+	if (input.empty()) {
+		this->delimiter = string("\0", 1);
+	}
+}
+
+void BufferedCSVReaderOptions::SetQuote(const string &quote_p) {
+	this->quote = quote_p;
+	this->has_quote = true;
+}
+
+void BufferedCSVReaderOptions::SetNewline(const string &input) {
+	if (input == "\\n" || input == "\\r") {
+		new_line = NewLineIdentifier::SINGLE;
+	} else if (input == "\\r\\n") {
+		new_line = NewLineIdentifier::CARRY_ON;
+	} else {
+		throw InvalidInputException("This is not accepted as a newline: " + input);
+	}
+	has_newline = true;
+}
+
+void BufferedCSVReaderOptions::SetDateFormat(LogicalTypeId type, const string &format, bool read_format) {
+	string error;
+	if (read_format) {
+		error = StrTimeFormat::ParseFormatSpecifier(format, date_format[type]);
+		date_format[type].format_specifier = format;
+	} else {
+		error = StrTimeFormat::ParseFormatSpecifier(format, write_date_format[type]);
+	}
+	if (!error.empty()) {
+		throw InvalidInputException("Could not parse DATEFORMAT: %s", error.c_str());
+	}
+	has_format[type] = true;
+}
+
+void BufferedCSVReaderOptions::SetReadOption(const string &loption, const Value &value,
+                                             vector<string> &expected_names) {
+	if (SetBaseOption(loption, value)) {
+		return;
+	}
+	if (loption == "auto_detect") {
+		auto_detect = ParseBoolean(value, loption);
+	} else if (loption == "sample_size") {
+		int64_t sample_size = ParseInteger(value, loption);
+		if (sample_size < 1 && sample_size != -1) {
+			throw BinderException("Unsupported parameter for SAMPLE_SIZE: cannot be smaller than 1");
+		}
+		if (sample_size == -1) {
+			sample_chunks = std::numeric_limits<uint64_t>::max();
+			sample_chunk_size = STANDARD_VECTOR_SIZE;
+		} else if (sample_size <= STANDARD_VECTOR_SIZE) {
+			sample_chunk_size = sample_size;
+			sample_chunks = 1;
+		} else {
+			sample_chunk_size = STANDARD_VECTOR_SIZE;
+			sample_chunks = sample_size / STANDARD_VECTOR_SIZE + 1;
+		}
+	} else if (loption == "skip") {
+		skip_rows = ParseInteger(value, loption);
+		skip_rows_set = true;
+	} else if (loption == "max_line_size" || loption == "maximum_line_size") {
+		maximum_line_size = ParseInteger(value, loption);
+	} else if (loption == "sample_chunk_size") {
+		sample_chunk_size = ParseInteger(value, loption);
+		if (sample_chunk_size > STANDARD_VECTOR_SIZE) {
+			throw BinderException(
+			    "Unsupported parameter for SAMPLE_CHUNK_SIZE: cannot be bigger than STANDARD_VECTOR_SIZE %d",
+			    STANDARD_VECTOR_SIZE);
+		} else if (sample_chunk_size < 1) {
+			throw BinderException("Unsupported parameter for SAMPLE_CHUNK_SIZE: cannot be smaller than 1");
+		}
+	} else if (loption == "sample_chunks") {
+		sample_chunks = ParseInteger(value, loption);
+		if (sample_chunks < 1) {
+			throw BinderException("Unsupported parameter for SAMPLE_CHUNKS: cannot be smaller than 1");
+		}
+	} else if (loption == "force_not_null") {
+		force_not_null = ParseColumnList(value, expected_names, loption);
+	} else if (loption == "date_format" || loption == "dateformat") {
+		string format = ParseString(value, loption);
+		SetDateFormat(LogicalTypeId::DATE, format, true);
+	} else if (loption == "timestamp_format" || loption == "timestampformat") {
+		string format = ParseString(value, loption);
+		SetDateFormat(LogicalTypeId::TIMESTAMP, format, true);
+	} else if (loption == "ignore_errors") {
+		ignore_errors = ParseBoolean(value, loption);
+	} else if (loption == "buffer_size") {
+		buffer_size = ParseInteger(value, loption);
+		if (buffer_size == 0) {
+			throw InvalidInputException("Buffer Size option must be higher than 0");
+		}
+	} else if (loption == "decimal_separator") {
+		decimal_separator = ParseString(value, loption);
+		if (decimal_separator != "." && decimal_separator != ",") {
+			throw BinderException("Unsupported parameter for DECIMAL_SEPARATOR: should be '.' or ','");
+		}
+	} else if (loption == "null_padding") {
+		null_padding = ParseBoolean(value, loption);
+	} else if (loption == "allow_quoted_nulls") {
+		allow_quoted_nulls = ParseBoolean(value, loption);
+	} else if (loption == "parallel") {
+		parallel_mode = ParseBoolean(value, loption) ? ParallelMode::PARALLEL : ParallelMode::SINGLE_THREADED;
+	} else if (loption == "rejects_table") {
+		// skip, handled in SetRejectsOptions
+		auto table_name = ParseString(value, loption);
+		if (table_name.empty()) {
+			throw BinderException("REJECTS_TABLE option cannot be empty");
+		}
+		rejects_table_name = table_name;
+	} else if (loption == "rejects_recovery_columns") {
+		// Get the list of columns to use as a recovery key
+		auto &children = ListValue::GetChildren(value);
+		for (auto &child : children) {
+			auto col_name = child.GetValue<string>();
+			rejects_recovery_columns.push_back(col_name);
+		}
+	} else if (loption == "rejects_limit") {
+		int64_t limit = ParseInteger(value, loption);
+		if (limit < 0) {
+			throw BinderException("Unsupported parameter for REJECTS_LIMIT: cannot be negative");
+		}
+		rejects_limit = limit;
+	} else {
+		throw BinderException("Unrecognized option for CSV reader \"%s\"", loption);
+	}
+}
+
+void BufferedCSVReaderOptions::SetWriteOption(const string &loption, const Value &value) {
+	if (loption == "new_line") {
+		// Steal this from SetBaseOption so we can write different newlines (e.g., format JSON ARRAY)
+		write_newline = ParseString(value, loption);
+		return;
+	}
+
+	if (SetBaseOption(loption, value)) {
+		return;
+	}
+
+	if (loption == "force_quote") {
+		force_quote = ParseColumnList(value, name_list, loption);
+	} else if (loption == "date_format" || loption == "dateformat") {
+		string format = ParseString(value, loption);
+		SetDateFormat(LogicalTypeId::DATE, format, false);
+	} else if (loption == "timestamp_format" || loption == "timestampformat") {
+		string format = ParseString(value, loption);
+		if (StringUtil::Lower(format) == "iso") {
+			format = "%Y-%m-%dT%H:%M:%S.%fZ";
+		}
+		SetDateFormat(LogicalTypeId::TIMESTAMP, format, false);
+		SetDateFormat(LogicalTypeId::TIMESTAMP_TZ, format, false);
+	} else if (loption == "prefix") {
+		prefix = ParseString(value, loption);
+	} else if (loption == "suffix") {
+		suffix = ParseString(value, loption);
+	} else {
+		throw BinderException("Unrecognized option CSV writer \"%s\"", loption);
+	}
+}
+
+bool BufferedCSVReaderOptions::SetBaseOption(const string &loption, const Value &value) {
+	// Make sure this function was only called after the option was turned into lowercase
+	D_ASSERT(!std::any_of(loption.begin(), loption.end(), ::isupper));
+
+	if (StringUtil::StartsWith(loption, "delim") || StringUtil::StartsWith(loption, "sep")) {
+		SetDelimiter(ParseString(value, loption));
+	} else if (loption == "quote") {
+		SetQuote(ParseString(value, loption));
+	} else if (loption == "new_line") {
+		SetNewline(ParseString(value, loption));
+	} else if (loption == "escape") {
+		SetEscape(ParseString(value, loption));
+	} else if (loption == "header") {
+		SetHeader(ParseBoolean(value, loption));
+	} else if (loption == "null" || loption == "nullstr") {
+		null_str = ParseString(value, loption);
+	} else if (loption == "encoding") {
+		auto encoding = StringUtil::Lower(ParseString(value, loption));
+		if (encoding != "utf8" && encoding != "utf-8") {
+			throw BinderException("Copy is only supported for UTF-8 encoded files, ENCODING 'UTF-8'");
+		}
+	} else if (loption == "compression") {
+		SetCompression(ParseString(value, loption));
+	} else {
+		// unrecognized option in base CSV
+		return false;
+	}
+	return true;
+}
+
+std::string BufferedCSVReaderOptions::ToString() const {
+	return "  file=" + file_path + "\n  delimiter='" + delimiter +
+	       (has_delimiter ? "'" : (auto_detect ? "' (auto detected)" : "' (default)")) + "\n  quote='" + quote +
+	       (has_quote ? "'" : (auto_detect ? "' (auto detected)" : "' (default)")) + "\n  escape='" + escape +
+	       (has_escape ? "'" : (auto_detect ? "' (auto detected)" : "' (default)")) +
+	       "\n  header=" + std::to_string(header) +
+	       (has_header ? "" : (auto_detect ? " (auto detected)" : "' (default)")) +
+	       "\n  sample_size=" + std::to_string(sample_chunk_size * sample_chunks) +
+	       "\n  ignore_errors=" + std::to_string(ignore_errors) + "\n  all_varchar=" + std::to_string(all_varchar);
+}
+
+} // namespace duckdb