diff --git a/src/azure_blob_filesystem.cpp b/src/azure_blob_filesystem.cpp index 9ef49e7..3ce690b 100644 --- a/src/azure_blob_filesystem.cpp +++ b/src/azure_blob_filesystem.cpp @@ -62,8 +62,8 @@ static bool Match(vector::const_iterator key, vector::const_iter //////// AzureBlobContextState //////// AzureBlobContextState::AzureBlobContextState(Azure::Storage::Blobs::BlobServiceClient client, - const AzureReadOptions &azure_read_options) - : AzureContextState(azure_read_options), service_client(std::move(client)) { + const AzureOptions &options) + : AzureContextState(options), service_client(std::move(client)) { } Azure::Storage::Blobs::BlobContainerClient @@ -73,27 +73,29 @@ AzureBlobContextState::GetBlobContainerClient(const std::string &blobContainerNa //////// AzureBlobStorageFileHandle //////// AzureBlobStorageFileHandle::AzureBlobStorageFileHandle(AzureBlobStorageFileSystem &fs, const OpenFileInfo &info, - FileOpenFlags flags, const AzureReadOptions &read_options, + FileOpenFlags flags, const AzureOptions &opts, Azure::Storage::Blobs::BlockBlobClient blob_client) - : AzureFileHandle(fs, info, flags, FileType::FILE_TYPE_INVALID, read_options), blob_client(std::move(blob_client)) { + : AzureFileHandle(fs, info, flags, FileType::FILE_TYPE_INVALID, opts), blob_client(std::move(blob_client)) { } void AzureBlobStorageFileHandle::Sync() { - D_ASSERT((flags.OpenForWriting() || flags.OpenForAppending()) || pending_block_ids.empty()); - if (pending_block_ids.empty()) { + if (!(flags.OpenForWriting() || flags.OpenForAppending())) { return; } try { + if (staged_block_count == 0) { + return; + } std::vector all_ids; - all_ids.reserve(committed_block_count + pending_block_ids.size()); - for (uint32_t i = 0; i < committed_block_count; i++) { + const auto total = static_cast(committed_block_count) + staged_block_count; + all_ids.reserve(total); + for (uint32_t i = 0; i < total; i++) { all_ids.push_back(MakeBlockId(i)); } - all_ids.insert(all_ids.end(), pending_block_ids.begin(), pending_block_ids.end()); auto res = blob_client.CommitBlockList(all_ids); last_modified = AzureBlobStorageFileSystem::ToTimestamp(res.Value.LastModified); - committed_block_count += pending_block_ids.size(); - pending_block_ids.clear(); + committed_block_count += staged_block_count; + staged_block_count = 0; } catch (const Azure::Storage::StorageException &e) { throw IOException("AzureBlobStorageFileSystem FileSync of '%s' failed with %s Reason Phrase: %s", GetPath(), e.ErrorCode, e.ReasonPhrase); @@ -126,8 +128,8 @@ unique_ptr AzureBlobStorageFileSystem::CreateHandle(const OpenF auto container = storage_context->As().GetBlobContainerClient(parsed_url.container); auto blob_client = container.GetBlockBlobClient(parsed_url.path); - auto handle = make_uniq(*this, info, flags, storage_context->read_options, - std::move(blob_client)); + auto handle = + make_uniq(*this, info, flags, storage_context->options, std::move(blob_client)); if (!handle->PostConstruct()) { return nullptr; } @@ -392,9 +394,9 @@ void AzureBlobStorageFileSystem::ReadRange(AzureFileHandle &handle, idx_t file_o range.Length = buffer_out_len; Azure::Storage::Blobs::DownloadBlobToOptions options; options.Range = range; - options.TransferOptions.Concurrency = afh.read_options.transfer_concurrency; - options.TransferOptions.InitialChunkSize = afh.read_options.transfer_chunk_size; - options.TransferOptions.ChunkSize = afh.read_options.transfer_chunk_size; + options.TransferOptions.Concurrency = afh.options.read_transfer_concurrency; + options.TransferOptions.InitialChunkSize = afh.options.read_transfer_chunk_size; + options.TransferOptions.ChunkSize = afh.options.read_transfer_chunk_size; auto res = afh.blob_client.DownloadTo((uint8_t *)buffer_out, buffer_out_len, options); } catch (const Azure::Storage::StorageException &e) { @@ -406,10 +408,9 @@ void AzureBlobStorageFileSystem::ReadRange(AzureFileHandle &handle, idx_t file_o shared_ptr AzureBlobStorageFileSystem::CreateStorageContext(optional_ptr opener, const string &path, const AzureParsedUrl &parsed_url) { - auto azure_read_options = ParseAzureReadOptions(opener); + auto azure_options = ParseAzureOptions(opener); - return make_shared_ptr(ConnectToBlobStorageAccount(opener, path, parsed_url), - azure_read_options); + return make_shared_ptr(ConnectToBlobStorageAccount(opener, path, parsed_url), azure_options); } int64_t AzureBlobStorageFileSystem::Write(FileHandle &handle, void *buffer, int64_t nr_bytes) { @@ -427,12 +428,16 @@ void AzureBlobStorageFileSystem::Write(FileHandle &handle, void *buffer, int64_t throw InternalException("Write called on file opened in read mode"); } - D_ASSERT(location == afh.file_offset || location == 0); + D_ASSERT(location == afh.file_offset); - auto block_id = MakeBlockId(afh.committed_block_count + static_cast(afh.pending_block_ids.size())); + auto block_id = MakeBlockId(static_cast(afh.committed_block_count) + afh.staged_block_count); auto body_stream = Azure::Core::IO::MemoryBodyStream(static_cast(buffer), nr_bytes); afh.blob_client.StageBlock(block_id, body_stream); // throws on error - afh.pending_block_ids.push_back(block_id); + afh.staged_block_count++; + if (afh.staged_block_count >= afh.options.write_staged_blocks_max) { + afh.Sync(); + } + afh.file_offset += nr_bytes; afh.length += nr_bytes; DUCKDB_LOG_FILE_SYSTEM_WRITE(handle, nr_bytes, afh.file_offset); diff --git a/src/azure_dfs_filesystem.cpp b/src/azure_dfs_filesystem.cpp index 2071257..1849eb6 100644 --- a/src/azure_dfs_filesystem.cpp +++ b/src/azure_dfs_filesystem.cpp @@ -90,8 +90,8 @@ static void Walk(const Azure::Storage::Files::DataLake::DataLakeFileSystemClient //////// AzureDfsContextState //////// AzureDfsContextState::AzureDfsContextState(Azure::Storage::Files::DataLake::DataLakeServiceClient client, - const AzureReadOptions &azure_read_options) - : AzureContextState(azure_read_options), service_client(std::move(client)) { + const AzureOptions &options) + : AzureContextState(options), service_client(std::move(client)) { } Azure::Storage::Files::DataLake::DataLakeFileSystemClient @@ -101,9 +101,9 @@ AzureDfsContextState::GetDfsFileSystemClient(const std::string &file_system_name //////// AzureDfsContextState //////// AzureDfsStorageFileHandle::AzureDfsStorageFileHandle(AzureDfsStorageFileSystem &fs, const OpenFileInfo &info, - FileOpenFlags flags, const AzureReadOptions &read_options, + FileOpenFlags flags, const AzureOptions &options, Azure::Storage::Files::DataLake::DataLakeFileClient client) - : AzureFileHandle(fs, info, flags, FileType::FILE_TYPE_INVALID, read_options), file_client(std::move(client)) { + : AzureFileHandle(fs, info, flags, FileType::FILE_TYPE_INVALID, options), file_client(std::move(client)) { } void AzureDfsStorageFileHandle::Sync(bool close) { @@ -139,7 +139,7 @@ unique_ptr AzureDfsStorageFileSystem::CreateHandle(const OpenFi auto storage_context = GetOrCreateStorageContext(opener, info.path, parsed_url); auto file_system_client = storage_context->As().GetDfsFileSystemClient(parsed_url.container); - auto handle = make_uniq(*this, info, flags, storage_context->read_options, + auto handle = make_uniq(*this, info, flags, storage_context->options, file_system_client.GetFileClient(parsed_url.path)); if (!handle->PostConstruct()) { return nullptr; @@ -335,9 +335,9 @@ void AzureDfsStorageFileSystem::ReadRange(AzureFileHandle &handle, idx_t file_of range.Length = buffer_out_len; Azure::Storage::Files::DataLake::DownloadFileToOptions options; options.Range = range; - options.TransferOptions.Concurrency = afh.read_options.transfer_concurrency; - options.TransferOptions.InitialChunkSize = afh.read_options.transfer_chunk_size; - options.TransferOptions.ChunkSize = afh.read_options.transfer_chunk_size; + options.TransferOptions.Concurrency = afh.options.read_transfer_concurrency; + options.TransferOptions.InitialChunkSize = afh.options.read_transfer_chunk_size; + options.TransferOptions.ChunkSize = afh.options.read_transfer_chunk_size; auto res = afh.file_client.DownloadTo((uint8_t *)buffer_out, buffer_out_len, options); } catch (const Azure::Storage::StorageException &e) { @@ -349,10 +349,9 @@ void AzureDfsStorageFileSystem::ReadRange(AzureFileHandle &handle, idx_t file_of shared_ptr AzureDfsStorageFileSystem::CreateStorageContext(optional_ptr opener, const string &path, const AzureParsedUrl &parsed_url) { - auto azure_read_options = ParseAzureReadOptions(opener); + auto azure_options = ParseAzureOptions(opener); - return make_shared_ptr(ConnectToDfsStorageAccount(opener, path, parsed_url), - azure_read_options); + return make_shared_ptr(ConnectToDfsStorageAccount(opener, path, parsed_url), azure_options); } int64_t AzureDfsStorageFileSystem::Write(FileHandle &handle, void *buffer, int64_t nr_bytes) { diff --git a/src/azure_extension.cpp b/src/azure_extension.cpp index 01812ed..28c7a4a 100644 --- a/src/azure_extension.cpp +++ b/src/azure_extension.cpp @@ -60,22 +60,28 @@ static void LoadInternal(ExtensionLoader &loader) { "values are: default, curl", LogicalType::VARCHAR, "default"); - AzureReadOptions default_read_options; + AzureOptions default_options; config.AddExtensionOption("azure_read_transfer_concurrency", "Maximum number of threads the Azure client can use for a single parallel read. " "If azure_read_transfer_chunk_size is less than azure_read_buffer_size then setting " "this > 1 will allow the Azure client to do concurrent requests to fill the buffer.", - LogicalType::INTEGER, Value::INTEGER(default_read_options.transfer_concurrency)); + LogicalType::INTEGER, Value::INTEGER(default_options.read_transfer_concurrency)); config.AddExtensionOption("azure_read_transfer_chunk_size", "Maximum size in bytes that the Azure client will read in a single request. " "It is recommended that this is a factor of azure_read_buffer_size.", - LogicalType::BIGINT, Value::BIGINT(default_read_options.transfer_chunk_size)); + LogicalType::BIGINT, Value::BIGINT(default_options.read_transfer_chunk_size)); config.AddExtensionOption("azure_read_buffer_size", "Size of the read buffer. It is recommended that this is evenly divisible by " "azure_read_transfer_chunk_size.", - LogicalType::UBIGINT, Value::UBIGINT(default_read_options.buffer_size)); + LogicalType::UBIGINT, Value::UBIGINT(default_options.read_buffer_size)); + + config.AddExtensionOption("azure_write_staged_blocks_max", + "Maximum number of staged (uncommitted) blocks before an automatic mid-write commit. " + "Azure limits blobs to 100,000 total blocks; this safety valve prevents hitting that " + "limit by committing staged blocks early. Default: 10000.", + LogicalType::UBIGINT, Value::UBIGINT(default_options.write_staged_blocks_max)); auto *http_proxy = std::getenv("HTTP_PROXY"); Value default_http_value = http_proxy ? Value(http_proxy) : Value(nullptr); diff --git a/src/azure_filesystem.cpp b/src/azure_filesystem.cpp index 976d53d..65572c0 100644 --- a/src/azure_filesystem.cpp +++ b/src/azure_filesystem.cpp @@ -12,8 +12,7 @@ namespace duckdb { -AzureContextState::AzureContextState(const AzureReadOptions &read_options) - : read_options(read_options), is_valid(true) { +AzureContextState::AzureContextState(const AzureOptions &options) : options(options), is_valid(true) { } bool AzureContextState::IsValid() const { @@ -25,16 +24,16 @@ void AzureContextState::QueryEnd() { } AzureFileHandle::AzureFileHandle(AzureStorageFileSystem &fs, const OpenFileInfo &info, FileOpenFlags flags, - FileType file_type, const AzureReadOptions &read_options) + FileType file_type, const AzureOptions &options) : FileHandle(fs, info.path, flags), flags(flags), // File info is_remote_loaded(false), file_type(file_type), length(0), last_modified(0), // Read info buffer_available(0), buffer_idx(0), file_offset(0), buffer_start(0), buffer_end(0), // Options - read_options(read_options) { + options(options) { if (!flags.RequireParallelAccess() && !flags.DirectIO()) { - read_buffer = duckdb::unique_ptr(new data_t[read_options.buffer_size]); + read_buffer = duckdb::unique_ptr(new data_t[options.read_buffer_size]); } // Set metadata of file when available, it avoids to invoke to the storage to get them. @@ -158,7 +157,7 @@ void AzureStorageFileSystem::Read(FileHandle &handle, void *buffer, int64_t nr_b } if (to_read > 0 && hfh.buffer_available == 0) { - auto new_buffer_available = MinValue(hfh.read_options.buffer_size, hfh.length - hfh.file_offset); + auto new_buffer_available = MinValue(hfh.options.read_buffer_size, hfh.length - hfh.file_offset); // Bypass buffer if we read more than buffer size if (to_read > new_buffer_available) { @@ -256,22 +255,27 @@ shared_ptr AzureStorageFileSystem::GetOrCreateStorageContext( return CreateStorageContext(opener, path, parsed_url); } -AzureReadOptions AzureStorageFileSystem::ParseAzureReadOptions(optional_ptr opener) { - AzureReadOptions options; +AzureOptions AzureStorageFileSystem::ParseAzureOptions(optional_ptr opener) { + AzureOptions options; Value concurrency_val; if (FileOpener::TryGetCurrentSetting(opener, "azure_read_transfer_concurrency", concurrency_val)) { - options.transfer_concurrency = concurrency_val.GetValue(); + options.read_transfer_concurrency = concurrency_val.GetValue(); } Value chunk_size_val; if (FileOpener::TryGetCurrentSetting(opener, "azure_read_transfer_chunk_size", chunk_size_val)) { - options.transfer_chunk_size = chunk_size_val.GetValue(); + options.read_transfer_chunk_size = chunk_size_val.GetValue(); } Value buffer_size_val; if (FileOpener::TryGetCurrentSetting(opener, "azure_read_buffer_size", buffer_size_val)) { - options.buffer_size = buffer_size_val.GetValue(); + options.read_buffer_size = buffer_size_val.GetValue(); + } + + Value write_staged_blocks_max_val; + if (FileOpener::TryGetCurrentSetting(opener, "azure_write_staged_blocks_max", write_staged_blocks_max_val)) { + options.write_staged_blocks_max = write_staged_blocks_max_val.GetValue(); } return options; diff --git a/src/include/azure_blob_filesystem.hpp b/src/include/azure_blob_filesystem.hpp index 88b377c..38c95f7 100644 --- a/src/include/azure_blob_filesystem.hpp +++ b/src/include/azure_blob_filesystem.hpp @@ -15,7 +15,7 @@ namespace duckdb { class AzureBlobContextState : public AzureContextState { public: - AzureBlobContextState(Azure::Storage::Blobs::BlobServiceClient client, const AzureReadOptions &azure_read_options); + AzureBlobContextState(Azure::Storage::Blobs::BlobServiceClient client, const AzureOptions &options); Azure::Storage::Blobs::BlobContainerClient GetBlobContainerClient(const std::string &blobContainerName) const; ~AzureBlobContextState() override = default; @@ -28,8 +28,7 @@ class AzureBlobStorageFileSystem; class AzureBlobStorageFileHandle : public AzureFileHandle { public: AzureBlobStorageFileHandle(AzureBlobStorageFileSystem &fs, const OpenFileInfo &info, FileOpenFlags flags, - const AzureReadOptions &read_options, - Azure::Storage::Blobs::BlockBlobClient blob_client); + const AzureOptions &options, Azure::Storage::Blobs::BlockBlobClient blob_client); ~AzureBlobStorageFileHandle() override = default; void Sync(); @@ -38,7 +37,7 @@ class AzureBlobStorageFileHandle : public AzureFileHandle { public: Azure::Storage::Blobs::BlockBlobClient blob_client; size_t committed_block_count = 0; // maintain position while open; only used if Sync() called before Close() - std::vector pending_block_ids; + uint32_t staged_block_count = 0; }; class AzureBlobStorageFileSystem : public AzureStorageFileSystem { diff --git a/src/include/azure_dfs_filesystem.hpp b/src/include/azure_dfs_filesystem.hpp index b25a73a..6ef17e0 100644 --- a/src/include/azure_dfs_filesystem.hpp +++ b/src/include/azure_dfs_filesystem.hpp @@ -15,8 +15,7 @@ namespace duckdb { class AzureDfsContextState : public AzureContextState { public: - AzureDfsContextState(Azure::Storage::Files::DataLake::DataLakeServiceClient client, - const AzureReadOptions &azure_read_options); + AzureDfsContextState(Azure::Storage::Files::DataLake::DataLakeServiceClient client, const AzureOptions &options); Azure::Storage::Files::DataLake::DataLakeFileSystemClient GetDfsFileSystemClient(const std::string &file_system_name) const; @@ -29,8 +28,7 @@ class AzureDfsStorageFileSystem; class AzureDfsStorageFileHandle : public AzureFileHandle { public: AzureDfsStorageFileHandle(AzureDfsStorageFileSystem &fs, const OpenFileInfo &info, FileOpenFlags flags, - const AzureReadOptions &read_options, - Azure::Storage::Files::DataLake::DataLakeFileClient client); + const AzureOptions &options, Azure::Storage::Files::DataLake::DataLakeFileClient client); ~AzureDfsStorageFileHandle() override = default; void Sync(bool close = false); diff --git a/src/include/azure_filesystem.hpp b/src/include/azure_filesystem.hpp index 0d0d75d..6a5e645 100644 --- a/src/include/azure_filesystem.hpp +++ b/src/include/azure_filesystem.hpp @@ -14,15 +14,16 @@ namespace duckdb { -struct AzureReadOptions { - int32_t transfer_concurrency = 5; - int64_t transfer_chunk_size = 1 * 1024 * 1024; - idx_t buffer_size = 1 * 1024 * 1024; +struct AzureOptions { + int32_t read_transfer_concurrency = 5; + int64_t read_transfer_chunk_size = (int64_t)8 * 1024 * 1024; + idx_t read_buffer_size = (idx_t)8 * 1024 * 1024; + idx_t write_staged_blocks_max = 10000; }; class AzureContextState : public ClientContextState { public: - const AzureReadOptions read_options; + const AzureOptions options; public: virtual bool IsValid() const; @@ -40,7 +41,7 @@ class AzureContextState : public ClientContextState { } protected: - AzureContextState(const AzureReadOptions &read_options); + explicit AzureContextState(const AzureOptions &options); protected: bool is_valid; @@ -62,7 +63,7 @@ class AzureFileHandle : public FileHandle { protected: AzureFileHandle(AzureStorageFileSystem &fs, const OpenFileInfo &info, FileOpenFlags flags, FileType file_type, - const AzureReadOptions &read_options); + const AzureOptions &options); public: FileOpenFlags flags; @@ -82,7 +83,7 @@ class AzureFileHandle : public FileHandle { idx_t buffer_start; idx_t buffer_end; - const AzureReadOptions read_options; + const AzureOptions options; }; class AzureStorageFileSystem : public FileSystem { @@ -132,7 +133,7 @@ class AzureStorageFileSystem : public FileSystem { const AzureParsedUrl &parsed_url) = 0; virtual void LoadRemoteFileInfo(AzureFileHandle &handle) = 0; - static AzureReadOptions ParseAzureReadOptions(optional_ptr opener); + static AzureOptions ParseAzureOptions(optional_ptr opener); public: static timestamp_t ToTimestamp(const Azure::DateTime &dt);