Whenever uploading a single buffer: single PUT instead of multipart upload

Author: carlopi

extension/httpfs/include/s3fs.hpp

@@ -215,6 +215,8 @@ class S3FileSystem : public HTTPFileSystem {
 	// Uploads the contents of write_buffer to S3.
 	// Note: caller is responsible to not call this method twice on the same buffer
 	static void UploadBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer);
+	static void UploadSingleBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer);
+	static void UploadBufferImplementation(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer, string query_param, bool direct_throw);
 
 	vector<OpenFileInfo> Glob(const string &glob_pattern, FileOpener *opener = nullptr) override;
 	bool ListFiles(const string &directory, const std::function<void(const string &, bool)> &callback,

extension/httpfs/s3fs.cpp

@@ -355,10 +355,21 @@ void S3FileSystem::NotifyUploadsInProgress(S3FileHandle &file_handle) {
 }
 
 void S3FileSystem::UploadBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer) {
-	auto &s3fs = (S3FileSystem &)file_handle.file_system;
-
 	string query_param = "partNumber=" + to_string(write_buffer->part_no + 1) + "&" +
 	                     "uploadId=" + S3FileSystem::UrlEncode(file_handle.multipart_upload_id, true);
+
+	UploadBufferImplementation(file_handle, write_buffer, query_param, false);
+
+	NotifyUploadsInProgress(file_handle);
+}
+
+void S3FileSystem::UploadSingleBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer) {
+	UploadBufferImplementation(file_handle, write_buffer, "", true);
+}
+
+void S3FileSystem::UploadBufferImplementation(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer, string query_param, bool single_upload) {
+	auto &s3fs = (S3FileSystem &)file_handle.file_system;
+
 	unique_ptr<HTTPResponse> res;
 	string etag;
 
@@ -376,6 +387,9 @@ void S3FileSystem::UploadBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuf
 		}
 		etag = res->headers.GetHeaderValue("ETag");
 	} catch (std::exception &ex) {
+		if (single_upload) {
+			throw;
+		}
 		ErrorData error(ex);
 		if (error.Type() != ExceptionType::IO && error.Type() != ExceptionType::HTTP) {
 			throw;
@@ -387,6 +401,7 @@ void S3FileSystem::UploadBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuf
 			file_handle.upload_exception = std::current_exception();
 		}
 
+		D_ASSERT(!single_upload); // If we are here we are in the multi-buffer situation
 		NotifyUploadsInProgress(file_handle);
 		return;
 	}
@@ -401,8 +416,6 @@ void S3FileSystem::UploadBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuf
 
 	// Free up space for another thread to acquire an S3WriteBuffer
 	write_buffer.reset();
-
-	NotifyUploadsInProgress(file_handle);
 }
 
 void S3FileSystem::FlushBuffer(S3FileHandle &file_handle, shared_ptr<S3WriteBuffer> write_buffer) {
@@ -465,7 +478,13 @@ void S3FileSystem::FlushAllBuffers(S3FileHandle &file_handle) {
 	file_handle.write_buffers_lock.unlock();
 
 	if (file_handle.initialized_multipart_upload == false) {
-		file_handle.multipart_upload_id = InitializeMultipartUpload(file_handle);
+		if (to_flush.size() == 1) {
+			UploadSingleBuffer(file_handle, to_flush[0]);
+			file_handle.upload_finalized= true;
+			return;
+		} else {
+			file_handle.multipart_upload_id = InitializeMultipartUpload(file_handle);
+		}
 	}
 	// Flush all buffers that aren't already uploading
 	for (auto &write_buffer : to_flush) {
@@ -483,6 +502,10 @@ void S3FileSystem::FlushAllBuffers(S3FileHandle &file_handle) {
 
 void S3FileSystem::FinalizeMultipartUpload(S3FileHandle &file_handle) {
 	auto &s3fs = (S3FileSystem &)file_handle.file_system;
+	if (file_handle.upload_finalized) {
+		return;
+	}
+
 	file_handle.upload_finalized = true;
 
 	std::stringstream ss;

test/configs/duckdb-tests.json

@@ -11,6 +11,12 @@
         "test/sql/secrets/secret_autoloading_errors.test",
         "test/sql/copy/csv/zstd_crash.test"
       ]
+    },
+    {
+      "reason": "Improved from 1 PUT + 2 POST to 1 PUT",
+      "paths": [
+        "test/sql/copy/s3/metadata_cache.test"
+      ]
     }
   ]
 }

test/sql/copy/no_head_on_write.test

@@ -48,19 +48,18 @@ copy (select 1 as a) to 's3://test-bucket/test-file.parquet'
 query I
 select request.type FROM duckdb_logs_parsed('HTTP')
 ----
-POST
 PUT
-POST
 
 statement ok
 CALL truncate_duckdb_logs();
 
 statement ok
-copy (select 1 as a) to 's3://test-bucket/test-file.csv'
+copy (select random() as a FROM range(8000000)) to 's3://test-bucket/test-file2.csv'
 
 query I
 select request.type FROM duckdb_logs_parsed('HTTP')
 ----
 POST
 PUT
+PUT
 POST