duckdb 0.3.2-dev300.0 → 0.3.2-dev324.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/package.json CHANGED
@@ -1,7 +1,7 @@
1
1
  {
2
2
  "name": "duckdb",
3
3
  "main": "./lib/duckdb.js",
4
- "version": "0.3.2-dev300.0",
4
+ "version": "0.3.2-dev324.0",
5
5
  "description": "DuckDB node.js API",
6
6
  "gypfile": true,
7
7
  "dependencies": {
package/src/duckdb.cpp CHANGED
@@ -10187,6 +10187,217 @@ uint64_t Checksum(uint8_t *buffer, size_t size) {
10187
10187
  }
10188
10188
 
10189
10189
  } // namespace duckdb
10190
+ //===----------------------------------------------------------------------===//
10191
+ // DuckDB
10192
+ //
10193
+ // duckdb/common/compressed_file_system.hpp
10194
+ //
10195
+ //
10196
+ //===----------------------------------------------------------------------===//
10197
+
10198
+
10199
+
10200
+
10201
+
10202
+
10203
+ namespace duckdb {
10204
+ class CompressedFile;
10205
+
10206
+ struct StreamData {
10207
+ // various buffers & pointers
10208
+ bool write = false;
10209
+ unique_ptr<data_t[]> in_buff;
10210
+ unique_ptr<data_t[]> out_buff;
10211
+ data_ptr_t out_buff_start = nullptr;
10212
+ data_ptr_t out_buff_end = nullptr;
10213
+ data_ptr_t in_buff_start = nullptr;
10214
+ data_ptr_t in_buff_end = nullptr;
10215
+
10216
+ idx_t in_buf_size = 0;
10217
+ idx_t out_buf_size = 0;
10218
+ };
10219
+
10220
+ struct StreamWrapper {
10221
+ DUCKDB_API virtual ~StreamWrapper();
10222
+
10223
+ DUCKDB_API virtual void Initialize(CompressedFile &file, bool write) = 0;
10224
+ DUCKDB_API virtual bool Read(StreamData &stream_data) = 0;
10225
+ DUCKDB_API virtual void Write(CompressedFile &file, StreamData &stream_data, data_ptr_t buffer,
10226
+ int64_t nr_bytes) = 0;
10227
+ DUCKDB_API virtual void Close() = 0;
10228
+ };
10229
+
10230
+ class CompressedFileSystem : public FileSystem {
10231
+ public:
10232
+ DUCKDB_API int64_t Read(FileHandle &handle, void *buffer, int64_t nr_bytes) override;
10233
+ DUCKDB_API int64_t Write(FileHandle &handle, void *buffer, int64_t nr_bytes) override;
10234
+
10235
+ DUCKDB_API void Reset(FileHandle &handle) override;
10236
+
10237
+ DUCKDB_API int64_t GetFileSize(FileHandle &handle) override;
10238
+
10239
+ DUCKDB_API bool OnDiskFile(FileHandle &handle) override;
10240
+ DUCKDB_API bool CanSeek() override;
10241
+
10242
+ DUCKDB_API virtual unique_ptr<StreamWrapper> CreateStream() = 0;
10243
+ DUCKDB_API virtual idx_t InBufferSize() = 0;
10244
+ DUCKDB_API virtual idx_t OutBufferSize() = 0;
10245
+ };
10246
+
10247
+ class CompressedFile : public FileHandle {
10248
+ public:
10249
+ DUCKDB_API CompressedFile(CompressedFileSystem &fs, unique_ptr<FileHandle> child_handle_p, const string &path);
10250
+ DUCKDB_API virtual ~CompressedFile() override;
10251
+
10252
+ CompressedFileSystem &compressed_fs;
10253
+ unique_ptr<FileHandle> child_handle;
10254
+ //! Whether the file is opened for reading or for writing
10255
+ bool write = false;
10256
+ StreamData stream_data;
10257
+
10258
+ public:
10259
+ DUCKDB_API void Initialize(bool write);
10260
+ DUCKDB_API int64_t ReadData(void *buffer, int64_t nr_bytes);
10261
+ DUCKDB_API int64_t WriteData(data_ptr_t buffer, int64_t nr_bytes);
10262
+ DUCKDB_API void Close() override;
10263
+
10264
+ private:
10265
+ unique_ptr<StreamWrapper> stream_wrapper;
10266
+ };
10267
+
10268
+ } // namespace duckdb
10269
+
10270
+
10271
+ namespace duckdb {
10272
+
10273
+ StreamWrapper::~StreamWrapper() {
10274
+ }
10275
+
10276
+ CompressedFile::CompressedFile(CompressedFileSystem &fs, unique_ptr<FileHandle> child_handle_p, const string &path)
10277
+ : FileHandle(fs, path), compressed_fs(fs), child_handle(move(child_handle_p)) {
10278
+ }
10279
+
10280
+ CompressedFile::~CompressedFile() {
10281
+ Close();
10282
+ }
10283
+
10284
+ void CompressedFile::Initialize(bool write) {
10285
+ Close();
10286
+
10287
+ this->write = write;
10288
+ stream_data.in_buf_size = compressed_fs.InBufferSize();
10289
+ stream_data.out_buf_size = compressed_fs.OutBufferSize();
10290
+ stream_data.in_buff = unique_ptr<data_t[]>(new data_t[stream_data.in_buf_size]);
10291
+ stream_data.in_buff_start = stream_data.in_buff.get();
10292
+ stream_data.in_buff_end = stream_data.in_buff.get();
10293
+ stream_data.out_buff = unique_ptr<data_t[]>(new data_t[stream_data.out_buf_size]);
10294
+ stream_data.out_buff_start = stream_data.out_buff.get();
10295
+ stream_data.out_buff_end = stream_data.out_buff.get();
10296
+
10297
+ stream_wrapper = compressed_fs.CreateStream();
10298
+ stream_wrapper->Initialize(*this, write);
10299
+ }
10300
+
10301
+ int64_t CompressedFile::ReadData(void *buffer, int64_t remaining) {
10302
+ idx_t total_read = 0;
10303
+ while (true) {
10304
+ // first check if there are input bytes available in the output buffers
10305
+ if (stream_data.out_buff_start != stream_data.out_buff_end) {
10306
+ // there is! copy it into the output buffer
10307
+ idx_t available = MinValue<idx_t>(remaining, stream_data.out_buff_end - stream_data.out_buff_start);
10308
+ memcpy(data_ptr_t(buffer) + total_read, stream_data.out_buff_start, available);
10309
+
10310
+ // increment the total read variables as required
10311
+ stream_data.out_buff_start += available;
10312
+ total_read += available;
10313
+ remaining -= available;
10314
+ if (remaining == 0) {
10315
+ // done! read enough
10316
+ return total_read;
10317
+ }
10318
+ }
10319
+ if (!stream_wrapper) {
10320
+ return total_read;
10321
+ }
10322
+
10323
+ // ran out of buffer: read more data from the child stream
10324
+ stream_data.out_buff_start = stream_data.out_buff.get();
10325
+ stream_data.out_buff_end = stream_data.out_buff.get();
10326
+ D_ASSERT(stream_data.in_buff_start <= stream_data.in_buff_end);
10327
+ D_ASSERT(stream_data.in_buff_end <= stream_data.in_buff_start + stream_data.in_buf_size);
10328
+
10329
+ // read more input if none available
10330
+ if (stream_data.in_buff_start == stream_data.in_buff_end) {
10331
+ // empty input buffer: refill from the start
10332
+ stream_data.in_buff_start = stream_data.in_buff.get();
10333
+ stream_data.in_buff_end = stream_data.in_buff_start;
10334
+ auto sz = child_handle->Read(stream_data.in_buff.get(), stream_data.in_buf_size);
10335
+ if (sz <= 0) {
10336
+ stream_wrapper.reset();
10337
+ break;
10338
+ }
10339
+ stream_data.in_buff_end = stream_data.in_buff_start + sz;
10340
+ }
10341
+
10342
+ auto finished = stream_wrapper->Read(stream_data);
10343
+ if (finished) {
10344
+ stream_wrapper.reset();
10345
+ }
10346
+ }
10347
+ return total_read;
10348
+ }
10349
+
10350
+ int64_t CompressedFile::WriteData(data_ptr_t buffer, int64_t nr_bytes) {
10351
+ stream_wrapper->Write(*this, stream_data, buffer, nr_bytes);
10352
+ return nr_bytes;
10353
+ }
10354
+
10355
+ void CompressedFile::Close() {
10356
+ if (stream_wrapper) {
10357
+ stream_wrapper->Close();
10358
+ stream_wrapper.reset();
10359
+ }
10360
+ stream_data.in_buff.reset();
10361
+ stream_data.out_buff.reset();
10362
+ stream_data.out_buff_start = nullptr;
10363
+ stream_data.out_buff_end = nullptr;
10364
+ stream_data.in_buff_start = nullptr;
10365
+ stream_data.in_buff_end = nullptr;
10366
+ stream_data.in_buf_size = 0;
10367
+ stream_data.out_buf_size = 0;
10368
+ }
10369
+
10370
+ int64_t CompressedFileSystem::Read(FileHandle &handle, void *buffer, int64_t nr_bytes) {
10371
+ auto &compressed_file = (CompressedFile &)handle;
10372
+ return compressed_file.ReadData(buffer, nr_bytes);
10373
+ }
10374
+
10375
+ int64_t CompressedFileSystem::Write(FileHandle &handle, void *buffer, int64_t nr_bytes) {
10376
+ auto &compressed_file = (CompressedFile &)handle;
10377
+ return compressed_file.WriteData((data_ptr_t)buffer, nr_bytes);
10378
+ }
10379
+
10380
+ void CompressedFileSystem::Reset(FileHandle &handle) {
10381
+ auto &compressed_file = (CompressedFile &)handle;
10382
+ compressed_file.child_handle->Reset();
10383
+ compressed_file.Initialize(compressed_file.write);
10384
+ }
10385
+
10386
+ int64_t CompressedFileSystem::GetFileSize(FileHandle &handle) {
10387
+ auto &compressed_file = (CompressedFile &)handle;
10388
+ return compressed_file.child_handle->GetFileSize();
10389
+ }
10390
+
10391
+ bool CompressedFileSystem::OnDiskFile(FileHandle &handle) {
10392
+ auto &compressed_file = (CompressedFile &)handle;
10393
+ return compressed_file.child_handle->OnDiskFile();
10394
+ }
10395
+
10396
+ bool CompressedFileSystem::CanSeek() {
10397
+ return false;
10398
+ }
10399
+
10400
+ } // namespace duckdb
10190
10401
 
10191
10402
 
10192
10403
 
@@ -10931,6 +11142,27 @@ ExpressionType FlipComparisionExpression(ExpressionType type) {
10931
11142
  } // namespace duckdb
10932
11143
 
10933
11144
 
11145
+
11146
+ namespace duckdb {
11147
+
11148
+ FileCompressionType FileCompressionTypeFromString(const string &input) {
11149
+ auto parameter = StringUtil::Lower(input);
11150
+ if (parameter == "infer" || parameter == "auto") {
11151
+ return FileCompressionType::AUTO_DETECT;
11152
+ } else if (parameter == "gzip") {
11153
+ return FileCompressionType::GZIP;
11154
+ } else if (parameter == "zstd") {
11155
+ return FileCompressionType::ZSTD;
11156
+ } else if (parameter == "uncompressed" || parameter == "none" || parameter.empty()) {
11157
+ return FileCompressionType::UNCOMPRESSED;
11158
+ } else {
11159
+ throw ParserException("Unrecognized file compression type \"%s\"", input);
11160
+ }
11161
+ }
11162
+
11163
+ } // namespace duckdb
11164
+
11165
+
10934
11166
  namespace duckdb {
10935
11167
 
10936
11168
  string JoinTypeToString(JoinType type) {
@@ -18692,6 +18924,10 @@ void FileSystem::RegisterSubSystem(unique_ptr<FileSystem> sub_fs) {
18692
18924
  throw NotImplementedException("%s: Can't register a sub system on a non-virtual file system", GetName());
18693
18925
  }
18694
18926
 
18927
+ void FileSystem::RegisterSubSystem(FileCompressionType compression_type, unique_ptr<FileSystem> sub_fs) {
18928
+ throw NotImplementedException("%s: Can't register a sub system on a non-virtual file system", GetName());
18929
+ }
18930
+
18695
18931
  bool FileSystem::CanHandleFile(const string &fpath) {
18696
18932
  throw NotImplementedException("%s: CanHandleFile is not implemented!", GetName());
18697
18933
  }
@@ -18712,6 +18948,10 @@ bool FileSystem::CanSeek() {
18712
18948
  throw NotImplementedException("%s: CanSeek is not implemented!", GetName());
18713
18949
  }
18714
18950
 
18951
+ unique_ptr<FileHandle> FileSystem::OpenCompressedFile(unique_ptr<FileHandle> handle, bool write) {
18952
+ throw NotImplementedException("%s: OpenCompressedFile is not implemented!", GetName());
18953
+ }
18954
+
18715
18955
  bool FileSystem::OnDiskFile(FileHandle &handle) {
18716
18956
  throw NotImplementedException("%s: OnDiskFile is not implemented!", GetName());
18717
18957
  }
@@ -18798,19 +19038,15 @@ FileType FileHandle::GetType() {
18798
19038
 
18799
19039
  namespace duckdb {
18800
19040
 
18801
- class GZipFileSystem : public FileSystem {
18802
- public:
18803
- static unique_ptr<FileHandle> OpenCompressedFile(unique_ptr<FileHandle> handle);
18804
-
18805
- int64_t Read(FileHandle &handle, void *buffer, int64_t nr_bytes) override;
18806
-
18807
- void Reset(FileHandle &handle) override;
19041
+ class GZipFileSystem : public CompressedFileSystem {
19042
+ // 32 KB
19043
+ static constexpr const idx_t BUFFER_SIZE = 1 << 15;
18808
19044
 
18809
- int64_t GetFileSize(FileHandle &handle) override;
19045
+ public:
19046
+ unique_ptr<FileHandle> OpenCompressedFile(unique_ptr<FileHandle> handle, bool write) override;
18810
19047
 
18811
- bool OnDiskFile(FileHandle &handle) override;
18812
- bool CanSeek() override {
18813
- return false;
19048
+ std::string GetName() const override {
19049
+ return "GZipFileSystem";
18814
19050
  }
18815
19051
 
18816
19052
  //! Verifies that a buffer contains a valid GZIP header
@@ -18818,9 +19054,9 @@ public:
18818
19054
  //! Consumes a byte stream as a gzip string, returning the decompressed string
18819
19055
  static string UncompressGZIPString(const string &in);
18820
19056
 
18821
- std::string GetName() const override {
18822
- return "GZipFileSystem";
18823
- }
19057
+ unique_ptr<StreamWrapper> CreateStream() override;
19058
+ idx_t InBufferSize() override;
19059
+ idx_t OutBufferSize() override;
18824
19060
  };
18825
19061
 
18826
19062
  static constexpr const uint8_t GZIP_COMPRESSION_DEFLATE = 0x08;
@@ -20142,6 +20378,155 @@ void *mz_zip_extract_archive_file_to_heap_v2(const char *pZip_filename, const ch
20142
20378
 
20143
20379
 
20144
20380
 
20381
+ // LICENSE_CHANGE_BEGIN
20382
+ // The following code up to LICENSE_CHANGE_END is subject to THIRD PARTY LICENSE #3
20383
+ // See the end of this file for a list
20384
+
20385
+ //===----------------------------------------------------------------------===//
20386
+ // DuckDB
20387
+ //
20388
+ // miniz_wrapper.hpp
20389
+ //
20390
+ //
20391
+ //===----------------------------------------------------------------------===//
20392
+
20393
+
20394
+
20395
+
20396
+ #include <string>
20397
+ #include <stdexcept>
20398
+
20399
+ namespace duckdb {
20400
+
20401
+ enum class MiniZStreamType {
20402
+ MINIZ_TYPE_NONE,
20403
+ MINIZ_TYPE_INFLATE,
20404
+ MINIZ_TYPE_DEFLATE
20405
+ };
20406
+
20407
+ struct MiniZStream {
20408
+ static constexpr uint8_t GZIP_HEADER_MINSIZE = 10;
20409
+ static constexpr uint8_t GZIP_FOOTER_SIZE = 8;
20410
+ static constexpr uint8_t GZIP_COMPRESSION_DEFLATE = 0x08;
20411
+ static constexpr unsigned char GZIP_FLAG_UNSUPPORTED = 0x1 | 0x2 | 0x4 | 0x10 | 0x20;
20412
+
20413
+ public:
20414
+ MiniZStream() : type(MiniZStreamType::MINIZ_TYPE_NONE) {
20415
+ memset(&stream, 0, sizeof(duckdb_miniz::mz_stream));
20416
+ }
20417
+ ~MiniZStream() {
20418
+ switch(type) {
20419
+ case MiniZStreamType::MINIZ_TYPE_INFLATE:
20420
+ duckdb_miniz::mz_inflateEnd(&stream);
20421
+ break;
20422
+ case MiniZStreamType::MINIZ_TYPE_DEFLATE:
20423
+ duckdb_miniz::mz_deflateEnd(&stream);
20424
+ break;
20425
+ default:
20426
+ break;
20427
+ }
20428
+ }
20429
+ void FormatException(std::string error_msg) {
20430
+ throw std::runtime_error(error_msg);
20431
+ }
20432
+ void FormatException(const char *error_msg, int mz_ret) {
20433
+ auto err = duckdb_miniz::mz_error(mz_ret);
20434
+ FormatException(error_msg + std::string(": ") + (err ? err : "Unknown error code"));
20435
+ }
20436
+ void Decompress(const char *compressed_data, size_t compressed_size, char *out_data, size_t out_size) {
20437
+ auto mz_ret = mz_inflateInit2(&stream, -MZ_DEFAULT_WINDOW_BITS);
20438
+ if (mz_ret != duckdb_miniz::MZ_OK) {
20439
+ FormatException("Failed to initialize miniz", mz_ret);
20440
+ }
20441
+ type = MiniZStreamType::MINIZ_TYPE_INFLATE;
20442
+
20443
+ if (compressed_size < GZIP_HEADER_MINSIZE) {
20444
+ FormatException("Failed to decompress GZIP block: compressed size is less than gzip header size");
20445
+ }
20446
+ auto gzip_hdr = (const unsigned char *)compressed_data;
20447
+ if (gzip_hdr[0] != 0x1F || gzip_hdr[1] != 0x8B || gzip_hdr[2] != GZIP_COMPRESSION_DEFLATE ||
20448
+ gzip_hdr[3] & GZIP_FLAG_UNSUPPORTED) {
20449
+ FormatException("Input is invalid/unsupported GZIP stream");
20450
+ }
20451
+
20452
+ stream.next_in = (const unsigned char *)compressed_data + GZIP_HEADER_MINSIZE;
20453
+ stream.avail_in = compressed_size - GZIP_HEADER_MINSIZE;
20454
+ stream.next_out = (unsigned char *)out_data;
20455
+ stream.avail_out = out_size;
20456
+
20457
+ mz_ret = mz_inflate(&stream, duckdb_miniz::MZ_FINISH);
20458
+ if (mz_ret != duckdb_miniz::MZ_OK && mz_ret != duckdb_miniz::MZ_STREAM_END) {
20459
+ FormatException("Failed to decompress GZIP block", mz_ret);
20460
+ }
20461
+ }
20462
+ size_t MaxCompressedLength(size_t input_size) {
20463
+ return duckdb_miniz::mz_compressBound(input_size) + GZIP_HEADER_MINSIZE + GZIP_FOOTER_SIZE;
20464
+ }
20465
+ static void InitializeGZIPHeader(unsigned char *gzip_header) {
20466
+ memset(gzip_header, 0, GZIP_HEADER_MINSIZE);
20467
+ gzip_header[0] = 0x1F;
20468
+ gzip_header[1] = 0x8B;
20469
+ gzip_header[2] = GZIP_COMPRESSION_DEFLATE;
20470
+ gzip_header[3] = 0;
20471
+ gzip_header[4] = 0;
20472
+ gzip_header[5] = 0;
20473
+ gzip_header[6] = 0;
20474
+ gzip_header[7] = 0;
20475
+ gzip_header[8] = 0;
20476
+ gzip_header[9] = 0xFF;
20477
+ }
20478
+
20479
+ static void InitializeGZIPFooter(unsigned char *gzip_footer, duckdb_miniz::mz_ulong crc, idx_t uncompressed_size) {
20480
+ gzip_footer[0] = crc & 0xFF;
20481
+ gzip_footer[1] = (crc >> 8) & 0xFF;
20482
+ gzip_footer[2] = (crc >> 16) & 0xFF;
20483
+ gzip_footer[3] = (crc >> 24) & 0xFF;
20484
+ gzip_footer[4] = uncompressed_size & 0xFF;
20485
+ gzip_footer[5] = (uncompressed_size >> 8) & 0xFF;
20486
+ gzip_footer[6] = (uncompressed_size >> 16) & 0xFF;
20487
+ gzip_footer[7] = (uncompressed_size >> 24) & 0xFF;
20488
+ }
20489
+
20490
+ void Compress(const char *uncompressed_data, size_t uncompressed_size, char *out_data, size_t *out_size) {
20491
+ auto mz_ret = mz_deflateInit2(&stream, duckdb_miniz::MZ_DEFAULT_LEVEL, MZ_DEFLATED, -MZ_DEFAULT_WINDOW_BITS, 1, 0);
20492
+ if (mz_ret != duckdb_miniz::MZ_OK) {
20493
+ FormatException("Failed to initialize miniz", mz_ret);
20494
+ }
20495
+ type = MiniZStreamType::MINIZ_TYPE_DEFLATE;
20496
+
20497
+ auto gzip_header = (unsigned char*) out_data;
20498
+ InitializeGZIPHeader(gzip_header);
20499
+
20500
+ auto gzip_body = gzip_header + GZIP_HEADER_MINSIZE;
20501
+
20502
+ stream.next_in = (const unsigned char*) uncompressed_data;
20503
+ stream.avail_in = uncompressed_size;
20504
+ stream.next_out = gzip_body;
20505
+ stream.avail_out = *out_size - GZIP_HEADER_MINSIZE;
20506
+
20507
+ mz_ret = mz_deflate(&stream, duckdb_miniz::MZ_FINISH);
20508
+ if (mz_ret != duckdb_miniz::MZ_OK && mz_ret != duckdb_miniz::MZ_STREAM_END) {
20509
+ FormatException("Failed to compress GZIP block", mz_ret);
20510
+ }
20511
+ auto gzip_footer = gzip_body + stream.total_out;
20512
+ auto crc = duckdb_miniz::mz_crc32(MZ_CRC32_INIT, (const unsigned char*) uncompressed_data, uncompressed_size);
20513
+ InitializeGZIPFooter(gzip_footer, crc, uncompressed_size);
20514
+
20515
+ *out_size = stream.total_out + GZIP_HEADER_MINSIZE + GZIP_FOOTER_SIZE;
20516
+ }
20517
+
20518
+ private:
20519
+ duckdb_miniz::mz_stream stream;
20520
+ MiniZStreamType type;
20521
+ };
20522
+
20523
+ }
20524
+
20525
+
20526
+ // LICENSE_CHANGE_END
20527
+
20528
+
20529
+
20145
20530
 
20146
20531
  namespace duckdb {
20147
20532
 
@@ -20200,206 +20585,202 @@ static idx_t GZipConsumeString(FileHandle &input) {
20200
20585
  return size;
20201
20586
  }
20202
20587
 
20203
- struct MiniZStreamWrapper {
20204
- ~MiniZStreamWrapper() {
20205
- Close();
20206
- }
20588
+ struct MiniZStreamWrapper : public StreamWrapper {
20589
+ ~MiniZStreamWrapper() override;
20207
20590
 
20591
+ CompressedFile *file = nullptr;
20208
20592
  duckdb_miniz::mz_stream *mz_stream_ptr = nullptr;
20593
+ bool writing = false;
20594
+ duckdb_miniz::mz_ulong crc;
20595
+ idx_t total_size;
20209
20596
 
20210
20597
  public:
20211
- void Initialize() {
20212
- Close();
20213
- mz_stream_ptr = new duckdb_miniz::mz_stream();
20214
- memset(mz_stream_ptr, 0, sizeof(duckdb_miniz::mz_stream));
20215
- }
20598
+ void Initialize(CompressedFile &file, bool write) override;
20216
20599
 
20217
- void Close() {
20218
- if (!mz_stream_ptr) {
20219
- return;
20220
- }
20221
- duckdb_miniz::mz_inflateEnd(mz_stream_ptr);
20222
- delete mz_stream_ptr;
20223
- mz_stream_ptr = nullptr;
20224
- }
20225
- };
20226
-
20227
- class GZipFile : public FileHandle {
20228
- static constexpr const idx_t BUFFER_SIZE = 1024;
20229
-
20230
- public:
20231
- GZipFile(unique_ptr<FileHandle> child_handle_p, const string &path)
20232
- : FileHandle(gzip_fs, path), child_handle(move(child_handle_p)) {
20233
- Initialize();
20234
- }
20235
- ~GZipFile() override {
20236
- Close();
20237
- }
20600
+ bool Read(StreamData &stream_data) override;
20601
+ void Write(CompressedFile &file, StreamData &stream_data, data_ptr_t buffer, int64_t nr_bytes) override;
20238
20602
 
20239
- void Initialize();
20240
- int64_t ReadData(void *buffer, int64_t nr_bytes);
20603
+ void Close() override;
20241
20604
 
20242
- GZipFileSystem gzip_fs;
20243
- unique_ptr<FileHandle> child_handle;
20244
-
20245
- protected:
20246
- void Close() override {
20247
- miniz_stream.reset();
20248
- in_buff.reset();
20249
- out_buff.reset();
20250
- }
20251
-
20252
- private:
20253
- idx_t data_start = 0;
20254
- unique_ptr<MiniZStreamWrapper> miniz_stream;
20255
- // various buffers & pointers
20256
- unique_ptr<data_t[]> in_buff;
20257
- unique_ptr<data_t[]> out_buff;
20258
- data_ptr_t out_buff_start = nullptr;
20259
- data_ptr_t out_buff_end = nullptr;
20260
- data_ptr_t in_buff_start = nullptr;
20261
- data_ptr_t in_buff_end = nullptr;
20605
+ void FlushStream();
20262
20606
  };
20263
20607
 
20264
- void GZipFileSystem::VerifyGZIPHeader(uint8_t gzip_hdr[], idx_t read_count) {
20265
- // check for incorrectly formatted files
20266
- if (read_count != GZIP_HEADER_MINSIZE) {
20267
- throw IOException("Input is not a GZIP stream");
20268
- }
20269
- if (gzip_hdr[0] != 0x1F || gzip_hdr[1] != 0x8B) { // magic header
20270
- throw IOException("Input is not a GZIP stream");
20271
- }
20272
- if (gzip_hdr[2] != GZIP_COMPRESSION_DEFLATE) { // compression method
20273
- throw IOException("Unsupported GZIP compression method");
20608
+ MiniZStreamWrapper::~MiniZStreamWrapper() {
20609
+ // avoid closing if destroyed during stack unwinding
20610
+ if (std::uncaught_exception()) {
20611
+ return;
20274
20612
  }
20275
- if (gzip_hdr[3] & GZIP_FLAG_UNSUPPORTED) {
20276
- throw IOException("Unsupported GZIP archive");
20613
+ try {
20614
+ Close();
20615
+ } catch (...) {
20277
20616
  }
20278
20617
  }
20279
20618
 
20280
- void GZipFile::Initialize() {
20619
+ void MiniZStreamWrapper::Initialize(CompressedFile &file, bool write) {
20281
20620
  Close();
20621
+ this->file = &file;
20622
+ mz_stream_ptr = new duckdb_miniz::mz_stream();
20623
+ memset(mz_stream_ptr, 0, sizeof(duckdb_miniz::mz_stream));
20624
+ this->writing = write;
20282
20625
 
20283
- D_ASSERT(BUFFER_SIZE >= 3); // found to work fine with 3
20626
+ // TODO use custom alloc/free methods in miniz to throw exceptions on OOM
20284
20627
  uint8_t gzip_hdr[GZIP_HEADER_MINSIZE];
20285
- data_start = GZIP_HEADER_MINSIZE;
20286
-
20287
- in_buff = unique_ptr<data_t[]>(new data_t[BUFFER_SIZE]);
20288
- in_buff_start = in_buff.get();
20289
- in_buff_end = in_buff.get();
20290
- out_buff = unique_ptr<data_t[]>(new data_t[BUFFER_SIZE]);
20291
- out_buff_start = out_buff.get();
20292
- out_buff_end = out_buff.get();
20293
-
20294
- miniz_stream = make_unique<MiniZStreamWrapper>();
20295
- miniz_stream->Initialize();
20628
+ if (write) {
20629
+ crc = MZ_CRC32_INIT;
20630
+ total_size = 0;
20296
20631
 
20297
- auto &mz_stream_ptr = miniz_stream->mz_stream_ptr;
20632
+ MiniZStream::InitializeGZIPHeader(gzip_hdr);
20633
+ file.child_handle->Write(gzip_hdr, GZIP_HEADER_MINSIZE);
20298
20634
 
20299
- // TODO use custom alloc/free methods in miniz to throw exceptions on OOM
20300
- auto read_count = child_handle->Read(gzip_hdr, GZIP_HEADER_MINSIZE);
20301
- GZipFileSystem::VerifyGZIPHeader(gzip_hdr, read_count);
20302
-
20303
- if (gzip_hdr[3] & GZIP_FLAG_NAME) {
20304
- child_handle->Seek(data_start);
20305
- data_start += GZipConsumeString(*child_handle);
20306
- }
20307
- child_handle->Seek(data_start);
20308
- // stream is now set to beginning of payload data
20309
- auto ret = duckdb_miniz::mz_inflateInit2((duckdb_miniz::mz_streamp)mz_stream_ptr, -MZ_DEFAULT_WINDOW_BITS);
20310
- if (ret != duckdb_miniz::MZ_OK) {
20311
- throw InternalException("Failed to initialize miniz");
20635
+ auto ret = mz_deflateInit2((duckdb_miniz::mz_streamp)mz_stream_ptr, duckdb_miniz::MZ_DEFAULT_LEVEL, MZ_DEFLATED,
20636
+ -MZ_DEFAULT_WINDOW_BITS, 1, 0);
20637
+ if (ret != duckdb_miniz::MZ_OK) {
20638
+ throw InternalException("Failed to initialize miniz");
20639
+ }
20640
+ } else {
20641
+ idx_t data_start = GZIP_HEADER_MINSIZE;
20642
+ auto read_count = file.child_handle->Read(gzip_hdr, GZIP_HEADER_MINSIZE);
20643
+ GZipFileSystem::VerifyGZIPHeader(gzip_hdr, read_count);
20644
+
20645
+ if (gzip_hdr[3] & GZIP_FLAG_NAME) {
20646
+ file.child_handle->Seek(data_start);
20647
+ data_start += GZipConsumeString(*file.child_handle);
20648
+ }
20649
+ file.child_handle->Seek(data_start);
20650
+ // stream is now set to beginning of payload data
20651
+ auto ret = duckdb_miniz::mz_inflateInit2((duckdb_miniz::mz_streamp)mz_stream_ptr, -MZ_DEFAULT_WINDOW_BITS);
20652
+ if (ret != duckdb_miniz::MZ_OK) {
20653
+ throw InternalException("Failed to initialize miniz");
20654
+ }
20655
+ }
20656
+ }
20657
+
20658
+ bool MiniZStreamWrapper::Read(StreamData &sd) {
20659
+ // actually decompress
20660
+ mz_stream_ptr->next_in = (data_ptr_t)sd.in_buff_start;
20661
+ D_ASSERT(sd.in_buff_end - sd.in_buff_start < NumericLimits<int32_t>::Maximum());
20662
+ mz_stream_ptr->avail_in = (uint32_t)(sd.in_buff_end - sd.in_buff_start);
20663
+ mz_stream_ptr->next_out = (data_ptr_t)sd.out_buff_end;
20664
+ mz_stream_ptr->avail_out = (uint32_t)((sd.out_buff.get() + sd.out_buf_size) - sd.out_buff_end);
20665
+ auto ret = duckdb_miniz::mz_inflate(mz_stream_ptr, duckdb_miniz::MZ_NO_FLUSH);
20666
+ if (ret != duckdb_miniz::MZ_OK && ret != duckdb_miniz::MZ_STREAM_END) {
20667
+ throw IOException("Failed to decode gzip stream: %s", duckdb_miniz::mz_error(ret));
20668
+ }
20669
+ // update pointers following inflate()
20670
+ sd.in_buff_start = (data_ptr_t)mz_stream_ptr->next_in;
20671
+ sd.in_buff_end = sd.in_buff_start + mz_stream_ptr->avail_in;
20672
+ sd.out_buff_end = (data_ptr_t)mz_stream_ptr->next_out;
20673
+ D_ASSERT(sd.out_buff_end + mz_stream_ptr->avail_out == sd.out_buff.get() + sd.out_buf_size);
20674
+ // if stream ended, deallocate inflator
20675
+ if (ret == duckdb_miniz::MZ_STREAM_END) {
20676
+ Close();
20677
+ return true;
20312
20678
  }
20679
+ return false;
20313
20680
  }
20314
20681
 
20315
- unique_ptr<FileHandle> GZipFileSystem::OpenCompressedFile(unique_ptr<FileHandle> handle) {
20316
- auto path = handle->path;
20317
- return make_unique<GZipFile>(move(handle), path);
20318
- }
20682
+ void MiniZStreamWrapper::Write(CompressedFile &file, StreamData &sd, data_ptr_t uncompressed_data,
20683
+ int64_t uncompressed_size) {
20684
+ // update the src and the total size
20685
+ crc = duckdb_miniz::mz_crc32(crc, (const unsigned char *)uncompressed_data, uncompressed_size);
20686
+ total_size += uncompressed_size;
20319
20687
 
20320
- int64_t GZipFile::ReadData(void *buffer, int64_t remaining) {
20321
- auto &mz_stream_ptr = miniz_stream->mz_stream_ptr;
20322
- idx_t total_read = 0;
20323
- while (true) {
20324
- // first check if there are input bytes available in the output buffers
20325
- if (out_buff_start != out_buff_end) {
20326
- // there is! copy it into the output buffer
20327
- idx_t available = MinValue<idx_t>(remaining, out_buff_end - out_buff_start);
20328
- memcpy(data_ptr_t(buffer) + total_read, out_buff_start, available);
20688
+ auto remaining = uncompressed_size;
20689
+ while (remaining > 0) {
20690
+ idx_t output_remaining = (sd.out_buff.get() + sd.out_buf_size) - sd.out_buff_start;
20329
20691
 
20330
- // increment the total read variables as required
20331
- out_buff_start += available;
20332
- total_read += available;
20333
- remaining -= available;
20334
- if (remaining == 0) {
20335
- // done! read enough
20336
- return total_read;
20337
- }
20692
+ mz_stream_ptr->next_in = (const unsigned char *)uncompressed_data;
20693
+ mz_stream_ptr->avail_in = remaining;
20694
+ mz_stream_ptr->next_out = sd.out_buff_start;
20695
+ mz_stream_ptr->avail_out = output_remaining;
20696
+
20697
+ auto res = mz_deflate(mz_stream_ptr, duckdb_miniz::MZ_NO_FLUSH);
20698
+ if (res != duckdb_miniz::MZ_OK) {
20699
+ D_ASSERT(res != duckdb_miniz::MZ_STREAM_END);
20700
+ throw InternalException("Failed to compress GZIP block");
20338
20701
  }
20339
- if (!mz_stream_ptr) {
20340
- return total_read;
20702
+ sd.out_buff_start += output_remaining - mz_stream_ptr->avail_out;
20703
+ if (mz_stream_ptr->avail_out == 0) {
20704
+ // no more output buffer available: flush
20705
+ file.child_handle->Write(sd.out_buff.get(), sd.out_buff_start - sd.out_buff.get());
20706
+ sd.out_buff_start = sd.out_buff.get();
20341
20707
  }
20708
+ idx_t written = remaining - mz_stream_ptr->avail_in;
20709
+ uncompressed_data += written;
20710
+ remaining = mz_stream_ptr->avail_in;
20711
+ }
20712
+ }
20342
20713
 
20343
- // ran out of buffer: read more data from the child stream
20344
- out_buff_start = out_buff.get();
20345
- out_buff_end = out_buff.get();
20346
- D_ASSERT(in_buff_start <= in_buff_end);
20347
- D_ASSERT(in_buff_end <= in_buff_start + GZipFile::BUFFER_SIZE);
20714
+ void MiniZStreamWrapper::FlushStream() {
20715
+ auto &sd = file->stream_data;
20716
+ mz_stream_ptr->next_in = nullptr;
20717
+ mz_stream_ptr->avail_in = 0;
20718
+ while (true) {
20719
+ auto output_remaining = (sd.out_buff.get() + sd.out_buf_size) - sd.out_buff_start;
20720
+ mz_stream_ptr->next_out = sd.out_buff_start;
20721
+ mz_stream_ptr->avail_out = output_remaining;
20348
20722
 
20349
- // read more input if none available
20350
- if (in_buff_start == in_buff_end) {
20351
- // empty input buffer: refill from the start
20352
- in_buff_start = in_buff.get();
20353
- auto sz = child_handle->Read(in_buff.get(), BUFFER_SIZE);
20354
- if (sz <= 0) {
20355
- break;
20356
- }
20357
- in_buff_end = in_buff_start + sz;
20723
+ auto res = mz_deflate(mz_stream_ptr, duckdb_miniz::MZ_FINISH);
20724
+ sd.out_buff_start += (output_remaining - mz_stream_ptr->avail_out);
20725
+ if (sd.out_buff_start > sd.out_buff.get()) {
20726
+ file->child_handle->Write(sd.out_buff.get(), sd.out_buff_start - sd.out_buff.get());
20727
+ sd.out_buff_start = sd.out_buff.get();
20358
20728
  }
20359
-
20360
- // actually decompress
20361
- mz_stream_ptr->next_in = (data_ptr_t)in_buff_start;
20362
- D_ASSERT(in_buff_end - in_buff_start < NumericLimits<int32_t>::Maximum());
20363
- mz_stream_ptr->avail_in = (uint32_t)(in_buff_end - in_buff_start);
20364
- mz_stream_ptr->next_out = (data_ptr_t)out_buff_end;
20365
- D_ASSERT((out_buff.get() + BUFFER_SIZE) - out_buff_end < NumericLimits<int32_t>::Maximum());
20366
- mz_stream_ptr->avail_out = (uint32_t)((out_buff.get() + BUFFER_SIZE) - out_buff_end);
20367
- auto ret = duckdb_miniz::mz_inflate(mz_stream_ptr, duckdb_miniz::MZ_NO_FLUSH);
20368
- if (ret != duckdb_miniz::MZ_OK && ret != duckdb_miniz::MZ_STREAM_END) {
20369
- throw Exception(duckdb_miniz::mz_error(ret));
20729
+ if (res == duckdb_miniz::MZ_STREAM_END) {
20730
+ break;
20370
20731
  }
20371
- // update pointers following inflate()
20372
- in_buff_start = (data_ptr_t)mz_stream_ptr->next_in;
20373
- in_buff_end = in_buff_start + mz_stream_ptr->avail_in;
20374
- out_buff_end = (data_ptr_t)mz_stream_ptr->next_out;
20375
- D_ASSERT(out_buff_end + mz_stream_ptr->avail_out == out_buff.get() + BUFFER_SIZE);
20376
- // if stream ended, deallocate inflator
20377
- if (ret == duckdb_miniz::MZ_STREAM_END) {
20378
- miniz_stream->Close();
20732
+ if (res != duckdb_miniz::MZ_OK) {
20733
+ throw InternalException("Failed to compress GZIP block");
20379
20734
  }
20380
20735
  }
20381
- return total_read;
20382
20736
  }
20383
20737
 
20384
- int64_t GZipFileSystem::Read(FileHandle &handle, void *buffer, int64_t nr_bytes) {
20385
- auto &gzip_file = (GZipFile &)handle;
20386
- return gzip_file.ReadData(buffer, nr_bytes);
20387
- }
20738
+ void MiniZStreamWrapper::Close() {
20739
+ if (!mz_stream_ptr) {
20740
+ return;
20741
+ }
20742
+ if (writing) {
20743
+ // flush anything remaining in the stream
20744
+ FlushStream();
20388
20745
 
20389
- void GZipFileSystem::Reset(FileHandle &handle) {
20390
- auto &gzip_file = (GZipFile &)handle;
20391
- gzip_file.child_handle->Reset();
20392
- gzip_file.Initialize();
20393
- }
20746
+ // write the footer
20747
+ unsigned char gzip_footer[MiniZStream::GZIP_FOOTER_SIZE];
20748
+ MiniZStream::InitializeGZIPFooter(gzip_footer, crc, total_size);
20749
+ file->child_handle->Write(gzip_footer, MiniZStream::GZIP_FOOTER_SIZE);
20394
20750
 
20395
- int64_t GZipFileSystem::GetFileSize(FileHandle &handle) {
20396
- auto &gzip_file = (GZipFile &)handle;
20397
- return gzip_file.child_handle->GetFileSize();
20751
+ duckdb_miniz::mz_deflateEnd(mz_stream_ptr);
20752
+ } else {
20753
+ duckdb_miniz::mz_inflateEnd(mz_stream_ptr);
20754
+ }
20755
+ delete mz_stream_ptr;
20756
+ mz_stream_ptr = nullptr;
20757
+ file = nullptr;
20398
20758
  }
20399
20759
 
20400
- bool GZipFileSystem::OnDiskFile(FileHandle &handle) {
20401
- auto &gzip_file = (GZipFile &)handle;
20402
- return gzip_file.child_handle->OnDiskFile();
20760
+ class GZipFile : public CompressedFile {
20761
+ public:
20762
+ GZipFile(unique_ptr<FileHandle> child_handle_p, const string &path, bool write)
20763
+ : CompressedFile(gzip_fs, move(child_handle_p), path) {
20764
+ Initialize(write);
20765
+ }
20766
+
20767
+ GZipFileSystem gzip_fs;
20768
+ };
20769
+
20770
+ void GZipFileSystem::VerifyGZIPHeader(uint8_t gzip_hdr[], idx_t read_count) {
20771
+ // check for incorrectly formatted files
20772
+ if (read_count != GZIP_HEADER_MINSIZE) {
20773
+ throw IOException("Input is not a GZIP stream");
20774
+ }
20775
+ if (gzip_hdr[0] != 0x1F || gzip_hdr[1] != 0x8B) { // magic header
20776
+ throw IOException("Input is not a GZIP stream");
20777
+ }
20778
+ if (gzip_hdr[2] != GZIP_COMPRESSION_DEFLATE) { // compression method
20779
+ throw IOException("Unsupported GZIP compression method");
20780
+ }
20781
+ if (gzip_hdr[3] & GZIP_FLAG_UNSUPPORTED) {
20782
+ throw IOException("Unsupported GZIP archive");
20783
+ }
20403
20784
  }
20404
20785
 
20405
20786
  string GZipFileSystem::UncompressGZIPString(const string &in) {
@@ -20458,6 +20839,23 @@ string GZipFileSystem::UncompressGZIPString(const string &in) {
20458
20839
  return decompressed;
20459
20840
  }
20460
20841
 
20842
+ unique_ptr<FileHandle> GZipFileSystem::OpenCompressedFile(unique_ptr<FileHandle> handle, bool write) {
20843
+ auto path = handle->path;
20844
+ return make_unique<GZipFile>(move(handle), path, write);
20845
+ }
20846
+
20847
+ unique_ptr<StreamWrapper> GZipFileSystem::CreateStream() {
20848
+ return make_unique<MiniZStreamWrapper>();
20849
+ }
20850
+
20851
+ idx_t GZipFileSystem::InBufferSize() {
20852
+ return BUFFER_SIZE;
20853
+ }
20854
+
20855
+ idx_t GZipFileSystem::OutBufferSize() {
20856
+ return BUFFER_SIZE;
20857
+ }
20858
+
20461
20859
  } // namespace duckdb
20462
20860
 
20463
20861
 
@@ -20793,15 +21191,14 @@ public:
20793
21191
  Close();
20794
21192
  }
20795
21193
 
20796
- protected:
21194
+ int fd;
21195
+
21196
+ public:
20797
21197
  void Close() override {
20798
21198
  if (fd != -1) {
20799
21199
  close(fd);
20800
21200
  }
20801
21201
  };
20802
-
20803
- public:
20804
- int fd;
20805
21202
  };
20806
21203
 
20807
21204
  static FileType GetFileTypeInternal(int fd) { // LCOV_EXCL_START
@@ -21169,14 +21566,13 @@ public:
21169
21566
  Close();
21170
21567
  }
21171
21568
 
21172
- protected:
21569
+ idx_t position;
21570
+ HANDLE fd;
21571
+
21572
+ public:
21173
21573
  void Close() override {
21174
21574
  CloseHandle(fd);
21175
21575
  };
21176
-
21177
- public:
21178
- idx_t position;
21179
- HANDLE fd;
21180
21576
  };
21181
21577
 
21182
21578
  unique_ptr<FileHandle> LocalFileSystem::OpenFile(const string &path, uint8_t flags, FileLockType lock_type,
@@ -26351,13 +26747,13 @@ public:
26351
26747
  : FileHandle(pipe_fs, path), child_handle(move(child_handle_p)) {
26352
26748
  }
26353
26749
 
26354
- int64_t ReadChunk(void *buffer, int64_t nr_bytes);
26355
- int64_t WriteChunk(void *buffer, int64_t nr_bytes);
26356
-
26357
26750
  PipeFileSystem pipe_fs;
26358
26751
  unique_ptr<FileHandle> child_handle;
26359
26752
 
26360
- protected:
26753
+ public:
26754
+ int64_t ReadChunk(void *buffer, int64_t nr_bytes);
26755
+ int64_t WriteChunk(void *buffer, int64_t nr_bytes);
26756
+
26361
26757
  void Close() override {
26362
26758
  }
26363
26759
  };
@@ -48084,6 +48480,7 @@ void VectorOperations::ReadFromStorage(data_ptr_t source, idx_t count, Vector &r
48084
48480
 
48085
48481
 
48086
48482
 
48483
+
48087
48484
  namespace duckdb {
48088
48485
 
48089
48486
  // bunch of wrappers to allow registering protocol handlers
@@ -48165,6 +48562,10 @@ public:
48165
48562
  sub_systems.push_back(move(fs));
48166
48563
  }
48167
48564
 
48565
+ void RegisterSubSystem(FileCompressionType compression_type, unique_ptr<FileSystem> fs) override {
48566
+ compressed_fs[compression_type] = move(fs);
48567
+ }
48568
+
48168
48569
  std::string GetName() const override {
48169
48570
  return "VirtualFileSystem";
48170
48571
  }
@@ -48181,6 +48582,7 @@ private:
48181
48582
 
48182
48583
  private:
48183
48584
  vector<unique_ptr<FileSystem>> sub_systems;
48585
+ map<FileCompressionType, unique_ptr<FileSystem>> compressed_fs;
48184
48586
  const unique_ptr<FileSystem> default_fs;
48185
48587
  };
48186
48588
 
@@ -48194,6 +48596,7 @@ private:
48194
48596
  namespace duckdb {
48195
48597
 
48196
48598
  VirtualFileSystem::VirtualFileSystem() : default_fs(FileSystem::CreateLocal()) {
48599
+ RegisterSubSystem(FileCompressionType::GZIP, make_unique<GZipFileSystem>());
48197
48600
  }
48198
48601
 
48199
48602
  unique_ptr<FileHandle> VirtualFileSystem::OpenFile(const string &path, uint8_t flags, FileLockType lock,
@@ -48203,6 +48606,8 @@ unique_ptr<FileHandle> VirtualFileSystem::OpenFile(const string &path, uint8_t f
48203
48606
  auto lower_path = StringUtil::Lower(path);
48204
48607
  if (StringUtil::EndsWith(lower_path, ".gz")) {
48205
48608
  compression = FileCompressionType::GZIP;
48609
+ } else if (StringUtil::EndsWith(lower_path, ".zst")) {
48610
+ compression = FileCompressionType::ZSTD;
48206
48611
  } else {
48207
48612
  compression = FileCompressionType::UNCOMPRESSED;
48208
48613
  }
@@ -48212,13 +48617,12 @@ unique_ptr<FileHandle> VirtualFileSystem::OpenFile(const string &path, uint8_t f
48212
48617
  if (file_handle->GetType() == FileType::FILE_TYPE_FIFO) {
48213
48618
  file_handle = PipeFileSystem::OpenPipe(move(file_handle));
48214
48619
  } else if (compression != FileCompressionType::UNCOMPRESSED) {
48215
- switch (compression) {
48216
- case FileCompressionType::GZIP:
48217
- file_handle = GZipFileSystem::OpenCompressedFile(move(file_handle));
48218
- break;
48219
- default:
48220
- throw NotImplementedException("Unimplemented compression type");
48620
+ auto entry = compressed_fs.find(compression);
48621
+ if (entry == compressed_fs.end()) {
48622
+ throw NotImplementedException(
48623
+ "Attempting to open a compressed file, but the compression type is not supported");
48221
48624
  }
48625
+ file_handle = entry->second->OpenCompressedFile(move(file_handle), flags & FileFlags::FILE_FLAGS_WRITE);
48222
48626
  }
48223
48627
  return file_handle;
48224
48628
  }
@@ -62229,8 +62633,8 @@ struct BufferedCSVReaderOptions {
62229
62633
  //! The file path of the CSV file to read
62230
62634
  string file_path;
62231
62635
  //! Whether file is compressed or not, and if so which compression type
62232
- //! ("infer" (default; infer from file extention), "gzip", "none")
62233
- string compression = "infer";
62636
+ //! AUTO_DETECT (default; infer from file extension)
62637
+ FileCompressionType compression = FileCompressionType::AUTO_DETECT;
62234
62638
  //! Whether or not to automatically detect dialect and datatypes
62235
62639
  bool auto_detect = false;
62236
62640
  //! Whether or not a delimiter was defined by the user
@@ -62302,7 +62706,6 @@ public:
62302
62706
  unique_ptr<FileHandle> file_handle;
62303
62707
  bool plain_file_source = false;
62304
62708
  idx_t file_size = 0;
62305
- FileCompressionType compression = FileCompressionType::UNCOMPRESSED;
62306
62709
 
62307
62710
  unique_ptr<char[]> buffer;
62308
62711
  idx_t buffer_size;
@@ -62584,15 +62987,8 @@ void BufferedCSVReader::PrepareComplexParser() {
62584
62987
  }
62585
62988
 
62586
62989
  unique_ptr<FileHandle> BufferedCSVReader::OpenCSV(const BufferedCSVReaderOptions &options) {
62587
- this->compression = FileCompressionType::UNCOMPRESSED;
62588
- if (options.compression == "infer" || options.compression == "auto") {
62589
- this->compression = FileCompressionType::AUTO_DETECT;
62590
- } else if (options.compression == "gzip") {
62591
- this->compression = FileCompressionType::GZIP;
62592
- }
62593
-
62594
62990
  auto result = fs.OpenFile(options.file_path.c_str(), FileFlags::FILE_FLAGS_READ, FileLockType::NO_LOCK,
62595
- this->compression, this->opener);
62991
+ options.compression, this->opener);
62596
62992
  plain_file_source = result->OnDiskFile() && result->CanSeek();
62597
62993
  file_size = result->GetFileSize();
62598
62994
  return result;
@@ -63235,9 +63631,8 @@ void BufferedCSVReader::DetectHeader(const vector<vector<LogicalType>> &best_sql
63235
63631
 
63236
63632
  } else {
63237
63633
  options.header = false;
63238
- idx_t total_columns = parse_chunk.ColumnCount();
63239
- for (idx_t col = 0; col < total_columns; col++) {
63240
- string column_name = GenerateColumnName(total_columns, col);
63634
+ for (idx_t col = 0; col < options.num_cols; col++) {
63635
+ string column_name = GenerateColumnName(options.num_cols, col);
63241
63636
  col_names.push_back(column_name);
63242
63637
  }
63243
63638
  }
@@ -63389,6 +63784,7 @@ vector<LogicalType> BufferedCSVReader::SniffCSV(const vector<LogicalType> &reque
63389
63784
  // #######
63390
63785
  // ### header detection
63391
63786
  // #######
63787
+ options.num_cols = best_num_cols;
63392
63788
  DetectHeader(best_sql_types_candidates, best_header_row);
63393
63789
 
63394
63790
  // #######
@@ -95536,11 +95932,7 @@ static bool ParseBaseOption(BufferedCSVReaderOptions &options, string &loption,
95536
95932
  throw BinderException("Copy is only supported for UTF-8 encoded files, ENCODING 'UTF-8'");
95537
95933
  }
95538
95934
  } else if (loption == "compression") {
95539
- options.compression = ParseString(set);
95540
- if (!(options.compression == "infer" || options.compression == "gzip" || options.compression == "none" ||
95541
- options.compression.empty())) {
95542
- throw BinderException("read_csv currently only supports 'gzip' compression.");
95543
- }
95935
+ options.compression = FileCompressionTypeFromString(ParseString(set));
95544
95936
  } else if (loption == "skip") {
95545
95937
  options.skip_rows = ParseInteger(set);
95546
95938
  } else {
@@ -95856,14 +96248,15 @@ struct LocalReadCSVData : public LocalFunctionData {
95856
96248
  };
95857
96249
 
95858
96250
  struct GlobalWriteCSVData : public GlobalFunctionData {
95859
- GlobalWriteCSVData(FileSystem &fs, const string &file_path, FileOpener *opener) : fs(fs) {
96251
+ GlobalWriteCSVData(FileSystem &fs, const string &file_path, FileOpener *opener, FileCompressionType compression)
96252
+ : fs(fs) {
95860
96253
  handle = fs.OpenFile(file_path, FileFlags::FILE_FLAGS_WRITE | FileFlags::FILE_FLAGS_FILE_CREATE_NEW,
95861
- FileLockType::WRITE_LOCK, FileSystem::DEFAULT_COMPRESSION, opener);
96254
+ FileLockType::WRITE_LOCK, compression, opener);
95862
96255
  }
95863
96256
 
95864
96257
  void WriteData(const_data_ptr_t data, idx_t size) {
95865
96258
  lock_guard<mutex> flock(lock);
95866
- fs.Write(*handle, (void *)data, size);
96259
+ handle->Write((void *)data, size);
95867
96260
  }
95868
96261
 
95869
96262
  FileSystem &fs;
@@ -95889,7 +96282,7 @@ static unique_ptr<GlobalFunctionData> WriteCSVInitializeGlobal(ClientContext &co
95889
96282
  auto &csv_data = (WriteCSVData &)bind_data;
95890
96283
  auto &options = csv_data.options;
95891
96284
  auto global_data = make_unique<GlobalWriteCSVData>(FileSystem::GetFileSystem(context), csv_data.files[0],
95892
- FileSystem::GetFileOpener(context));
96285
+ FileSystem::GetFileOpener(context), options.compression);
95893
96286
 
95894
96287
  if (options.header) {
95895
96288
  BufferedSerializer serializer;
@@ -95977,6 +96370,15 @@ static void WriteCSVCombine(ClientContext &context, FunctionData &bind_data, Glo
95977
96370
  }
95978
96371
  }
95979
96372
 
96373
+ //===--------------------------------------------------------------------===//
96374
+ // Finalize
96375
+ //===--------------------------------------------------------------------===//
96376
+ void WriteCSVFinalize(ClientContext &context, FunctionData &bind_data, GlobalFunctionData &gstate) {
96377
+ auto &global_state = (GlobalWriteCSVData &)gstate;
96378
+ global_state.handle->Close();
96379
+ global_state.handle.reset();
96380
+ }
96381
+
95980
96382
  void CSVCopyFunction::RegisterFunction(BuiltinFunctions &set) {
95981
96383
  CopyFunction info("csv");
95982
96384
  info.copy_to_bind = WriteCSVBind;
@@ -95984,6 +96386,7 @@ void CSVCopyFunction::RegisterFunction(BuiltinFunctions &set) {
95984
96386
  info.copy_to_initialize_global = WriteCSVInitializeGlobal;
95985
96387
  info.copy_to_sink = WriteCSVSink;
95986
96388
  info.copy_to_combine = WriteCSVCombine;
96389
+ info.copy_to_finalize = WriteCSVFinalize;
95987
96390
 
95988
96391
  info.copy_from_bind = ReadCSVBind;
95989
96392
  info.copy_from_function = ReadCSVTableFunction::GetFunction();
@@ -96887,7 +97290,7 @@ static unique_ptr<FunctionData> ReadCSVBind(ClientContext &context, vector<Value
96887
97290
  throw BinderException("read_csv requires at least a single column as input!");
96888
97291
  }
96889
97292
  } else if (kv.first == "compression") {
96890
- options.compression = kv.second.str_value;
97293
+ options.compression = FileCompressionTypeFromString(kv.second.str_value);
96891
97294
  } else if (kv.first == "filename") {
96892
97295
  result->include_file_name = kv.second.value_.boolean;
96893
97296
  } else if (kv.first == "skip") {
@@ -96898,10 +97301,6 @@ static unique_ptr<FunctionData> ReadCSVBind(ClientContext &context, vector<Value
96898
97301
  throw BinderException("read_csv requires columns to be specified. Use read_csv_auto or set read_csv(..., "
96899
97302
  "AUTO_DETECT=TRUE) to automatically guess columns.");
96900
97303
  }
96901
- if (!(options.compression == "infer" || options.compression == "gzip" || options.compression == "none" ||
96902
- options.compression.empty())) {
96903
- throw BinderException("read_csv currently only supports 'gzip' compression.");
96904
- }
96905
97304
  if (options.auto_detect) {
96906
97305
  options.file_path = result->files[0];
96907
97306
  auto initial_reader = make_unique<BufferedCSVReader>(context, options);
@@ -97029,8 +97428,13 @@ void ReadCSVTableFunction::RegisterFunction(BuiltinFunctions &set) {
97029
97428
 
97030
97429
  unique_ptr<TableFunctionRef> ReadCSVReplacement(const string &table_name, void *data) {
97031
97430
  auto lower_name = StringUtil::Lower(table_name);
97032
- if (!StringUtil::EndsWith(lower_name, ".csv") && !StringUtil::EndsWith(lower_name, ".tsv") &&
97033
- !StringUtil::EndsWith(lower_name, ".csv.gz")) {
97431
+ // remove any compression
97432
+ if (StringUtil::EndsWith(lower_name, ".gz")) {
97433
+ lower_name = lower_name.substr(0, lower_name.size() - 3);
97434
+ } else if (StringUtil::EndsWith(lower_name, ".zst")) {
97435
+ lower_name = lower_name.substr(0, lower_name.size() - 4);
97436
+ }
97437
+ if (!StringUtil::EndsWith(lower_name, ".csv") && !StringUtil::EndsWith(lower_name, ".tsv")) {
97034
97438
  return nullptr;
97035
97439
  }
97036
97440
  auto table_function = make_unique<TableFunctionRef>();
package/src/duckdb.hpp CHANGED
@@ -10,8 +10,8 @@ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLI
10
10
 
11
11
  #pragma once
12
12
  #define DUCKDB_AMALGAMATION 1
13
- #define DUCKDB_SOURCE_ID "c9a9053a2"
14
- #define DUCKDB_VERSION "0.3.2-dev300"
13
+ #define DUCKDB_SOURCE_ID "5d577b5d9"
14
+ #define DUCKDB_VERSION "0.3.2-dev324"
15
15
  //===----------------------------------------------------------------------===//
16
16
  // DuckDB
17
17
  //
@@ -1719,7 +1719,9 @@ using std::unordered_map;
1719
1719
 
1720
1720
  namespace duckdb {
1721
1721
 
1722
- enum class FileCompressionType : uint8_t { AUTO_DETECT = 0, UNCOMPRESSED = 1, GZIP = 2 };
1722
+ enum class FileCompressionType : uint8_t { AUTO_DETECT = 0, UNCOMPRESSED = 1, GZIP = 2, ZSTD = 3 };
1723
+
1724
+ FileCompressionType FileCompressionTypeFromString(const string &input);
1723
1725
 
1724
1726
  } // namespace duckdb
1725
1727
 
@@ -1779,7 +1781,7 @@ public:
1779
1781
  idx_t GetFileSize();
1780
1782
  FileType GetType();
1781
1783
 
1782
- protected:
1784
+ //! Closes the file handle.
1783
1785
  virtual void Close() = 0;
1784
1786
 
1785
1787
  public:
@@ -1883,6 +1885,7 @@ public:
1883
1885
 
1884
1886
  //! registers a sub-file system to handle certain file name prefixes, e.g. http:// etc.
1885
1887
  virtual void RegisterSubSystem(unique_ptr<FileSystem> sub_fs);
1888
+ virtual void RegisterSubSystem(FileCompressionType compression_type, unique_ptr<FileSystem> fs);
1886
1889
 
1887
1890
  //! Whether or not a sub-system can handle a specific file path
1888
1891
  virtual bool CanHandleFile(const string &fpath);
@@ -1899,6 +1902,8 @@ public:
1899
1902
  //! in a file on-disk are much cheaper than e.g. random reads in a file over the network
1900
1903
  virtual bool OnDiskFile(FileHandle &handle);
1901
1904
 
1905
+ virtual unique_ptr<FileHandle> OpenCompressedFile(unique_ptr<FileHandle> handle, bool write);
1906
+
1902
1907
  //! Create a LocalFileSystem.
1903
1908
  static unique_ptr<FileSystem> CreateLocal();
1904
1909