Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 27 additions & 22 deletions src/azure_blob_filesystem.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -62,8 +62,8 @@ static bool Match(vector<string>::const_iterator key, vector<string>::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
Expand All @@ -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<std::string> 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<uint32_t>(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);
Expand Down Expand Up @@ -126,8 +128,8 @@ unique_ptr<AzureFileHandle> AzureBlobStorageFileSystem::CreateHandle(const OpenF
auto container = storage_context->As<AzureBlobContextState>().GetBlobContainerClient(parsed_url.container);
auto blob_client = container.GetBlockBlobClient(parsed_url.path);

auto handle = make_uniq<AzureBlobStorageFileHandle>(*this, info, flags, storage_context->read_options,
std::move(blob_client));
auto handle =
make_uniq<AzureBlobStorageFileHandle>(*this, info, flags, storage_context->options, std::move(blob_client));
if (!handle->PostConstruct()) {
return nullptr;
}
Expand Down Expand Up @@ -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) {
Expand All @@ -406,10 +408,9 @@ void AzureBlobStorageFileSystem::ReadRange(AzureFileHandle &handle, idx_t file_o
shared_ptr<AzureContextState> AzureBlobStorageFileSystem::CreateStorageContext(optional_ptr<FileOpener> opener,
const string &path,
const AzureParsedUrl &parsed_url) {
auto azure_read_options = ParseAzureReadOptions(opener);
auto azure_options = ParseAzureOptions(opener);

return make_shared_ptr<AzureBlobContextState>(ConnectToBlobStorageAccount(opener, path, parsed_url),
azure_read_options);
return make_shared_ptr<AzureBlobContextState>(ConnectToBlobStorageAccount(opener, path, parsed_url), azure_options);
}

int64_t AzureBlobStorageFileSystem::Write(FileHandle &handle, void *buffer, int64_t nr_bytes) {
Expand All @@ -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<uint32_t>(afh.pending_block_ids.size()));
auto block_id = MakeBlockId(static_cast<uint32_t>(afh.committed_block_count) + afh.staged_block_count);
auto body_stream = Azure::Core::IO::MemoryBodyStream(static_cast<uint8_t *>(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);
Expand Down
21 changes: 10 additions & 11 deletions src/azure_dfs_filesystem.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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) {
Expand Down Expand Up @@ -139,7 +139,7 @@ unique_ptr<AzureFileHandle> AzureDfsStorageFileSystem::CreateHandle(const OpenFi
auto storage_context = GetOrCreateStorageContext(opener, info.path, parsed_url);
auto file_system_client = storage_context->As<AzureDfsContextState>().GetDfsFileSystemClient(parsed_url.container);

auto handle = make_uniq<AzureDfsStorageFileHandle>(*this, info, flags, storage_context->read_options,
auto handle = make_uniq<AzureDfsStorageFileHandle>(*this, info, flags, storage_context->options,
file_system_client.GetFileClient(parsed_url.path));
if (!handle->PostConstruct()) {
return nullptr;
Expand Down Expand Up @@ -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) {
Expand All @@ -349,10 +349,9 @@ void AzureDfsStorageFileSystem::ReadRange(AzureFileHandle &handle, idx_t file_of
shared_ptr<AzureContextState> AzureDfsStorageFileSystem::CreateStorageContext(optional_ptr<FileOpener> opener,
const string &path,
const AzureParsedUrl &parsed_url) {
auto azure_read_options = ParseAzureReadOptions(opener);
auto azure_options = ParseAzureOptions(opener);

return make_shared_ptr<AzureDfsContextState>(ConnectToDfsStorageAccount(opener, path, parsed_url),
azure_read_options);
return make_shared_ptr<AzureDfsContextState>(ConnectToDfsStorageAccount(opener, path, parsed_url), azure_options);
}

int64_t AzureDfsStorageFileSystem::Write(FileHandle &handle, void *buffer, int64_t nr_bytes) {
Expand Down
14 changes: 10 additions & 4 deletions src/azure_extension.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
26 changes: 15 additions & 11 deletions src/azure_filesystem.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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<data_t[]>(new data_t[read_options.buffer_size]);
read_buffer = duckdb::unique_ptr<data_t[]>(new data_t[options.read_buffer_size]);
}

// Set metadata of file when available, it avoids to invoke to the storage to get them.
Expand Down Expand Up @@ -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<idx_t>(hfh.read_options.buffer_size, hfh.length - hfh.file_offset);
auto new_buffer_available = MinValue<idx_t>(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) {
Expand Down Expand Up @@ -256,22 +255,27 @@ shared_ptr<AzureContextState> AzureStorageFileSystem::GetOrCreateStorageContext(
return CreateStorageContext(opener, path, parsed_url);
}

AzureReadOptions AzureStorageFileSystem::ParseAzureReadOptions(optional_ptr<FileOpener> opener) {
AzureReadOptions options;
AzureOptions AzureStorageFileSystem::ParseAzureOptions(optional_ptr<FileOpener> opener) {
AzureOptions options;

Value concurrency_val;
if (FileOpener::TryGetCurrentSetting(opener, "azure_read_transfer_concurrency", concurrency_val)) {
options.transfer_concurrency = concurrency_val.GetValue<int32_t>();
options.read_transfer_concurrency = concurrency_val.GetValue<int32_t>();
}

Value chunk_size_val;
if (FileOpener::TryGetCurrentSetting(opener, "azure_read_transfer_chunk_size", chunk_size_val)) {
options.transfer_chunk_size = chunk_size_val.GetValue<int64_t>();
options.read_transfer_chunk_size = chunk_size_val.GetValue<int64_t>();
}

Value buffer_size_val;
if (FileOpener::TryGetCurrentSetting(opener, "azure_read_buffer_size", buffer_size_val)) {
options.buffer_size = buffer_size_val.GetValue<idx_t>();
options.read_buffer_size = buffer_size_val.GetValue<idx_t>();
}

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<idx_t>();
}

return options;
Expand Down
7 changes: 3 additions & 4 deletions src/include/azure_blob_filesystem.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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();
Expand All @@ -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<std::string> pending_block_ids;
uint32_t staged_block_count = 0;
};

class AzureBlobStorageFileSystem : public AzureStorageFileSystem {
Expand Down
6 changes: 2 additions & 4 deletions src/include/azure_dfs_filesystem.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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);
Expand Down
19 changes: 10 additions & 9 deletions src/include/azure_filesystem.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Was the chunk size change deliberate? If so it probably should go into the PR description

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah yes, I sliced apart the original PR into pre- and post- v1.5.3 parts but missed that one. I'll add comment since it's a good adjustment in any case.

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;
Expand All @@ -40,7 +41,7 @@ class AzureContextState : public ClientContextState {
}

protected:
AzureContextState(const AzureReadOptions &read_options);
explicit AzureContextState(const AzureOptions &options);

protected:
bool is_valid;
Expand All @@ -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;
Expand All @@ -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 {
Expand Down Expand Up @@ -132,7 +133,7 @@ class AzureStorageFileSystem : public FileSystem {
const AzureParsedUrl &parsed_url) = 0;

virtual void LoadRemoteFileInfo(AzureFileHandle &handle) = 0;
static AzureReadOptions ParseAzureReadOptions(optional_ptr<FileOpener> opener);
static AzureOptions ParseAzureOptions(optional_ptr<FileOpener> opener);

public:
static timestamp_t ToTimestamp(const Azure::DateTime &dt);
Expand Down
Loading