2022-04-26 12:57:02 +00:00
|
|
|
#include "WithFileSize.h"
|
|
|
|
#include <IO/ReadBufferFromFile.h>
|
|
|
|
#include <IO/CompressedReadBufferWrapper.h>
|
|
|
|
#include <IO/ParallelReadBuffer.h>
|
2024-02-13 11:02:46 +00:00
|
|
|
#include <IO/ReadBufferFromFileDecorator.h>
|
2023-06-16 15:51:18 +00:00
|
|
|
#include <IO/PeekableReadBuffer.h>
|
2022-04-26 12:57:02 +00:00
|
|
|
|
|
|
|
namespace DB
|
|
|
|
{
|
|
|
|
|
2022-05-25 14:49:40 +00:00
|
|
|
namespace ErrorCodes
|
|
|
|
{
|
|
|
|
extern const int UNKNOWN_FILE_SIZE;
|
|
|
|
}
|
|
|
|
|
2022-04-26 12:57:02 +00:00
|
|
|
template <typename T>
|
2022-05-25 14:49:40 +00:00
|
|
|
static size_t getFileSize(T & in)
|
2022-04-26 12:57:02 +00:00
|
|
|
{
|
|
|
|
if (auto * with_file_size = dynamic_cast<WithFileSize *>(&in))
|
2024-02-13 11:02:46 +00:00
|
|
|
{
|
2022-04-26 12:57:02 +00:00
|
|
|
return with_file_size->getFileSize();
|
2024-02-13 11:02:46 +00:00
|
|
|
}
|
2022-04-26 12:57:02 +00:00
|
|
|
|
2022-05-25 14:49:40 +00:00
|
|
|
throw Exception(ErrorCodes::UNKNOWN_FILE_SIZE, "Cannot find out file size");
|
2022-04-26 12:57:02 +00:00
|
|
|
}
|
|
|
|
|
2022-05-25 14:49:40 +00:00
|
|
|
size_t getFileSizeFromReadBuffer(ReadBuffer & in)
|
2022-04-26 12:57:02 +00:00
|
|
|
{
|
2024-02-13 11:02:46 +00:00
|
|
|
if (auto * delegate = dynamic_cast<ReadBufferFromFileDecorator *>(&in))
|
|
|
|
{
|
|
|
|
return getFileSize(delegate->getWrappedReadBuffer());
|
|
|
|
}
|
|
|
|
else if (auto * compressed = dynamic_cast<CompressedReadBufferWrapper *>(&in))
|
|
|
|
{
|
2022-04-26 12:57:02 +00:00
|
|
|
return getFileSize(compressed->getWrappedReadBuffer());
|
2024-02-13 11:02:46 +00:00
|
|
|
}
|
2022-04-26 12:57:02 +00:00
|
|
|
|
|
|
|
return getFileSize(in);
|
|
|
|
}
|
|
|
|
|
2023-06-16 15:51:18 +00:00
|
|
|
std::optional<size_t> tryGetFileSizeFromReadBuffer(ReadBuffer & in)
|
|
|
|
{
|
|
|
|
try
|
|
|
|
{
|
|
|
|
return getFileSizeFromReadBuffer(in);
|
|
|
|
}
|
|
|
|
catch (...)
|
|
|
|
{
|
|
|
|
return std::nullopt;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2022-04-26 12:57:02 +00:00
|
|
|
bool isBufferWithFileSize(const ReadBuffer & in)
|
|
|
|
{
|
2024-02-13 11:02:46 +00:00
|
|
|
if (const auto * delegate = dynamic_cast<const ReadBufferFromFileDecorator *>(&in))
|
|
|
|
{
|
|
|
|
return delegate->isWithFileSize();
|
|
|
|
}
|
|
|
|
else if (const auto * compressed = dynamic_cast<const CompressedReadBufferWrapper *>(&in))
|
2022-04-26 12:57:02 +00:00
|
|
|
{
|
|
|
|
return isBufferWithFileSize(compressed->getWrappedReadBuffer());
|
|
|
|
}
|
|
|
|
|
|
|
|
return dynamic_cast<const WithFileSize *>(&in) != nullptr;
|
|
|
|
}
|
|
|
|
|
2023-06-16 15:51:18 +00:00
|
|
|
size_t getDataOffsetMaybeCompressed(const ReadBuffer & in)
|
|
|
|
{
|
2024-02-13 11:02:46 +00:00
|
|
|
if (const auto * delegate = dynamic_cast<const ReadBufferFromFileDecorator *>(&in))
|
|
|
|
{
|
|
|
|
return getDataOffsetMaybeCompressed(delegate->getWrappedReadBuffer());
|
|
|
|
}
|
|
|
|
else if (const auto * compressed = dynamic_cast<const CompressedReadBufferWrapper *>(&in))
|
2023-06-16 15:51:18 +00:00
|
|
|
{
|
|
|
|
return getDataOffsetMaybeCompressed(compressed->getWrappedReadBuffer());
|
|
|
|
}
|
|
|
|
else if (const auto * peekable = dynamic_cast<const PeekableReadBuffer *>(&in))
|
|
|
|
{
|
|
|
|
return getDataOffsetMaybeCompressed(peekable->getSubBuffer());
|
|
|
|
}
|
|
|
|
|
|
|
|
return in.count();
|
|
|
|
}
|
|
|
|
|
|
|
|
|
2022-04-26 12:57:02 +00:00
|
|
|
}
|