duckdb 0.3.5-dev987.0 → 0.4.1-dev102.0

package/src/parquet-amalgamation.hpp

@@ -4625,7 +4625,7 @@ std::ostream& operator<<(std::ostream& out, const FileCryptoMetaData& obj);
 // LICENSE_CHANGE_END
 
 
-
+#include <list>
 
 
 // LICENSE_CHANGE_BEGIN
@@ -6916,31 +6916,174 @@ protected:
 
 namespace duckdb {
 
+// A ReadHead for prefetching data in a specific range
+struct ReadHead {
+	ReadHead(idx_t location, uint64_t size) : location(location), size(size) {};
+	// Hint info
+	idx_t location;
+	uint64_t size;
+
+	// Current info
+	unique_ptr<AllocatedData> data;
+	bool data_isset = false;
+
+	idx_t GetEnd() const {
+		return size + location;
+	}
+
+	void Allocate(Allocator &allocator) {
+		data = allocator.Allocate(size);
+	}
+};
+
+// Comparator for ReadHeads that are either overlapping, adjacent, or within ALLOW_GAP bytes from each other
+struct ReadHeadComparator {
+	static constexpr uint64_t ALLOW_GAP = 1 << 14; // 16 KiB
+	bool operator()(const ReadHead *a, const ReadHead *b) const {
+		auto a_start = a->location;
+		auto a_end = a->location + a->size;
+		auto b_start = b->location;
+
+		if (a_end <= NumericLimits<idx_t>::Maximum() - ALLOW_GAP) {
+			a_end += ALLOW_GAP;
+		}
+
+		return a_start < b_start && a_end < b_start;
+	}
+};
+
+// Two-step read ahead buffer
+// 1: register all ranges that will be read, merging ranges that are consecutive
+// 2: prefetch all registered ranges
+struct ReadAheadBuffer {
+	ReadAheadBuffer(Allocator &allocator, FileHandle &handle, FileOpener &opener)
+	    : allocator(allocator), handle(handle), file_opener(opener) {
+	}
+
+	// The list of read heads
+	std::list<ReadHead> read_heads;
+	// Set for merging consecutive ranges
+	std::set<ReadHead *, ReadHeadComparator> merge_set;
+
+	Allocator &allocator;
+	FileHandle &handle;
+	FileOpener &file_opener;
+
+	idx_t total_size = 0;
+
+	// Add a read head to the prefetching list
+	void AddReadHead(idx_t pos, uint64_t len, bool merge_buffers = true) {
+		// Attempt to merge with existing
+		if (merge_buffers) {
+			ReadHead new_read_head {pos, len};
+			auto lookup_set = merge_set.find(&new_read_head);
+			if (lookup_set != merge_set.end()) {
+				auto existing_head = *lookup_set;
+				auto new_start = MinValue<idx_t>(existing_head->location, new_read_head.location);
+				auto new_length = MaxValue<idx_t>(existing_head->GetEnd(), new_read_head.GetEnd()) - new_start;
+				existing_head->location = new_start;
+				existing_head->size = new_length;
+				return;
+			}
+		}
+
+		read_heads.emplace_front(ReadHead(pos, len));
+		total_size += len;
+		auto &read_head = read_heads.front();
+
+		if (merge_buffers) {
+			merge_set.insert(&read_head);
+		}
+
+		if (read_head.GetEnd() > handle.GetFileSize()) {
+			throw std::runtime_error("Prefetch registered for bytes outside file");
+		}
+	}
+
+	// Returns the relevant read head
+	ReadHead *GetReadHead(idx_t pos) {
+		for (auto &read_head : read_heads) {
+			if (pos >= read_head.location && pos < read_head.GetEnd()) {
+				return &read_head;
+			}
+		}
+		return nullptr;
+	}
+
+	// Prefetch all read heads
+	void Prefetch() {
+		for (auto &read_head : read_heads) {
+			read_head.Allocate(allocator);
+
+			if (read_head.GetEnd() > handle.GetFileSize()) {
+				throw std::runtime_error("Prefetch registered requested for bytes outside file");
+			}
+
+			handle.Read(read_head.data->get(), read_head.size, read_head.location);
+			read_head.data_isset = true;
+		}
+	}
+};
+
 class ThriftFileTransport : public duckdb_apache::thrift::transport::TVirtualTransport<ThriftFileTransport> {
 public:
-	ThriftFileTransport(Allocator &allocator, FileHandle &handle_p)
-	    : allocator(allocator), handle(handle_p), location(0) {
+	static constexpr uint64_t PREFETCH_FALLBACK_BUFFERSIZE = 1000000;
+
+	ThriftFileTransport(Allocator &allocator, FileHandle &handle_p, FileOpener &opener, bool prefetch_mode_p)
+	    : handle(handle_p), location(0), allocator(allocator), ra_buffer(ReadAheadBuffer(allocator, handle_p, opener)),
+	      prefetch_mode(prefetch_mode_p) {
 	}
 
 	uint32_t read(uint8_t *buf, uint32_t len) {
-		if (prefetched_data && location >= prefetch_location &&
-		    location + len < prefetch_location + prefetched_data->GetSize()) {
-			memcpy(buf, prefetched_data->get() + location - prefetch_location, len);
+		auto prefetch_buffer = ra_buffer.GetReadHead(location);
+		if (prefetch_buffer != nullptr && location - prefetch_buffer->location + len <= prefetch_buffer->size) {
+			D_ASSERT(location - prefetch_buffer->location + len <= prefetch_buffer->size);
+
+			if (!prefetch_buffer->data_isset) {
+				prefetch_buffer->Allocate(allocator);
+				handle.Read(prefetch_buffer->data->get(), prefetch_buffer->size, prefetch_buffer->location);
+				prefetch_buffer->data_isset = true;
+			}
+			memcpy(buf, prefetch_buffer->data->get() + location - prefetch_buffer->location, len);
 		} else {
-			handle.Read(buf, len, location);
+			if (prefetch_mode && len < PREFETCH_FALLBACK_BUFFERSIZE && len > 0) {
+				Prefetch(location, MinValue<uint64_t>(PREFETCH_FALLBACK_BUFFERSIZE, handle.GetFileSize() - location));
+				auto prefetch_buffer_fallback = ra_buffer.GetReadHead(location);
+				D_ASSERT(location - prefetch_buffer_fallback->location + len <= prefetch_buffer_fallback->size);
+				memcpy(buf, prefetch_buffer_fallback->data->get() + location - prefetch_buffer_fallback->location, len);
+			} else {
+				handle.Read(buf, len, location);
+			}
 		}
 		location += len;
 		return len;
 	}
 
-	void Prefetch(idx_t pos, idx_t len) {
-		prefetch_location = pos;
-		prefetched_data = allocator.Allocate(len);
-		handle.Read(prefetched_data->get(), len, prefetch_location);
+	// Prefetch a single buffer
+	void Prefetch(idx_t pos, uint64_t len) {
+		RegisterPrefetch(pos, len, false);
+		FinalizeRegistration();
+		PrefetchRegistered();
+	}
+
+	// Register a buffer for prefixing
+	void RegisterPrefetch(idx_t pos, uint64_t len, bool can_merge = true) {
+		ra_buffer.AddReadHead(pos, len, can_merge);
+	}
+
+	// Prevents any further merges, should be called before PrefetchRegistered
+	void FinalizeRegistration() {
+		ra_buffer.merge_set.clear();
+	}
+
+	// Prefetch all previously registered ranges
+	void PrefetchRegistered() {
+		ra_buffer.Prefetch();
 	}
 
 	void ClearPrefetch() {
-		prefetched_data.reset();
+		ra_buffer.read_heads.clear();
+		ra_buffer.merge_set.clear();
 	}
 
 	void SetLocation(idx_t location_p) {
@@ -6955,12 +7098,17 @@ public:
 	}
 
 private:
-	Allocator &allocator;
 	FileHandle &handle;
 	idx_t location;
 
-	unique_ptr<AllocatedData> prefetched_data;
-	idx_t prefetch_location;
+	Allocator &allocator;
+
+	// Multi-buffer prefetch
+	ReadAheadBuffer ra_buffer;
+
+	// Whether the prefetch mode is enabled. In this mode the DirectIO flag of the handle will be set and the parquet
+	// reader will manage the read buffering.
+	bool prefetch_mode;
 };
 
 } // namespace duckdb
@@ -7417,8 +7565,13 @@ public:
 	idx_t MaxDefine() const;
 	idx_t MaxRepeat() const;
 
+	virtual idx_t FileOffset() const;
+	virtual uint64_t TotalCompressedSize();
 	virtual idx_t GroupRowsAvailable();
 
+	// register the range this reader will touch for prefetching
+	virtual void RegisterPrefetch(ThriftFileTransport &transport, bool allow_merge);
+
 	virtual unique_ptr<BaseStatistics> Stats(const std::vector<ColumnChunk> &columns);
 
 protected:
@@ -7433,6 +7586,9 @@ protected:
 	virtual void DictReference(Vector &result);
 	virtual void PlainReference(shared_ptr<ByteBuffer>, Vector &result);
 
+	// applies any skips that were registered using Skip()
+	virtual void ApplyPendingSkips(idx_t num_values);
+
 	bool HasDefines() {
 		return max_define > 0;
 	}
@@ -7451,13 +7607,15 @@ protected:
 	ParquetReader &reader;
 	LogicalType type;
 
+	idx_t pending_skips = 0;
+
 private:
 	void PrepareRead(parquet_filter_t &filter);
 	void PreparePage(idx_t compressed_page_size, idx_t uncompressed_page_size);
 	void PrepareDataPage(PageHeader &page_hdr);
 	void PreparePageV2(PageHeader &page_hdr);
 
-	const duckdb_parquet::format::ColumnChunk *chunk;
+	const duckdb_parquet::format::ColumnChunk *chunk = nullptr;
 
 	duckdb_apache::thrift::protocol::TProtocol *protocol;
 	idx_t page_rows_available;
@@ -7545,6 +7703,11 @@ class ChunkCollection;
 class BaseStatistics;
 class TableFilterSet;
 
+struct ParquetReaderPrefetchConfig {
+	// Percentage of data in a row group span that should be scanned for enabling whole group prefetch
+	static constexpr double WHOLE_GROUP_PREFETCH_MINIMUM_SCAN = 0.95;
+};
+
 struct ParquetReaderScanState {
 	vector<idx_t> group_idx_list;
 	int64_t current_group;
@@ -7560,6 +7723,9 @@ struct ParquetReaderScanState {
 
 	ResizeableBuffer define_buf;
 	ResizeableBuffer repeat_buf;
+
+	bool prefetch_mode = false;
+	bool current_group_prefetched = false;
 };
 
 struct ParquetOptions {
@@ -7624,6 +7790,10 @@ private:
 	                                               idx_t depth, idx_t max_define, idx_t max_repeat,
 	                                               idx_t &next_schema_idx, idx_t &next_file_idx);
 	const duckdb_parquet::format::RowGroup &GetGroup(ParquetReaderScanState &state);
+	uint64_t GetGroupCompressedSize(ParquetReaderScanState &state);
+	idx_t GetGroupOffset(ParquetReaderScanState &state);
+	// Group span is the distance between the min page offset and the max page offset plus the max page compressed size
+	uint64_t GetGroupSpan(ParquetReaderScanState &state);
 	void PrepareRowGroupBuffer(ParquetReaderScanState &state, idx_t out_col_idx);
 	LogicalType DeriveLogicalType(const SchemaElement &s_ele);
 

package/test/pathnames.test.js

@@ -0,0 +1,82 @@
+var sqlite3 = require("..");
+var assert = require("assert");
+
+describe("pathname search support", function () {
+  let db;
+  describe("without search paths", () => {
+    before((done) => {
+      db = new sqlite3.Database(":memory:", done);
+    });
+
+    it("supports a full path", function (done) {
+      db.prepare('select * from "test/support/prepare.csv"').all(
+        (err, result) => {
+          assert(err === null);
+          assert(result.length === 5000);
+          done();
+        }
+      );
+    });
+
+    it("don't not support a partial path", function (done) {
+      db.prepare('select * from "prepare.csv"').all((err, result) => {
+        assert(err.code === "DUCKDB_NODEJS_ERROR");
+        assert(err.errno === -1);
+        assert(result == null);
+        done();
+      });
+    });
+  });
+
+  describe("with search paths", () => {
+    before((done) => {
+      db = new sqlite3.Database(":memory:", () => {
+        db.prepare("SET FILE_SEARCH_PATH='test/support'").run(done);
+      });
+    });
+
+    it("supports a full path", function (done) {
+      db.prepare('select * from "test/support/prepare.csv"').all(
+        (err, result) => {
+          assert(err === null);
+          assert(result.length === 5000);
+          done();
+        }
+      );
+    });
+
+    it("supports a partial path", function (done) {
+      db.prepare('select * from "prepare.csv"').all((err, result) => {
+        assert(err === null);
+        assert(result.length === 5000);
+        done();
+      });
+    });
+  });
+
+  describe("with multiple search paths", () => {
+    before((done) => {
+      db = new sqlite3.Database(":memory:", () => {
+        db.prepare("SET FILE_SEARCH_PATH='test/support'").run(done);
+      });
+    });
+
+    it("supports a full path", function (done) {
+      db.prepare('select * from "test/support/prepare.csv"').all(
+        (err, result) => {
+          assert(err === null);
+          assert(result.length === 5000);
+          done();
+        }
+      );
+    });
+
+    it("supports a partial path", function (done) {
+      db.prepare('select * from "prepare.csv"').all((err, result) => {
+        assert(err === null);
+        assert(result.length === 5000);
+        done();
+      });
+    });
+  });
+});