duckdb 0.3.5-dev113.0 → 0.3.5-dev1181.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/src/statement.cpp

@@ -121,6 +121,89 @@ static duckdb::Value bind_parameter(const Napi::Value source) {
 	return duckdb::Value();
 }
 
+static Napi::Value convert_col_val(Napi::Env &env, duckdb::Value dval, duckdb::LogicalTypeId id) {
+	Napi::Value value;
+
+	// TODO templateroo here
+	switch (id) {
+	case duckdb::LogicalTypeId::BOOLEAN: {
+		value = Napi::Boolean::New(env, duckdb::BooleanValue::Get(dval));
+	} break;
+	case duckdb::LogicalTypeId::INTEGER: {
+		value = Napi::Number::New(env, duckdb::IntegerValue::Get(dval));
+	} break;
+	case duckdb::LogicalTypeId::FLOAT: {
+		value = Napi::Number::New(env, duckdb::FloatValue::Get(dval));
+	} break;
+	case duckdb::LogicalTypeId::DOUBLE: {
+		value = Napi::Number::New(env, duckdb::DoubleValue::Get(dval));
+	} break;
+	case duckdb::LogicalTypeId::BIGINT: {
+		value = Napi::Number::New(env, duckdb::BigIntValue::Get(dval));
+	} break;
+	case duckdb::LogicalTypeId::HUGEINT: {
+		value = Napi::Number::New(env, dval.GetValue<double>());
+	} break;
+	case duckdb::LogicalTypeId::DECIMAL: {
+		value = Napi::Number::New(env, dval.GetValue<double>());
+	} break;
+	case duckdb::LogicalTypeId::INTERVAL: {
+		auto interval = duckdb::IntervalValue::Get(dval);
+		auto object_value = Napi::Object::New(env);
+		object_value.Set("months", interval.months);
+		object_value.Set("days", interval.days);
+		object_value.Set("micros", interval.micros);
+		value = object_value;
+	} break;
+#if (NAPI_VERSION > 4)
+	case duckdb::LogicalTypeId::DATE: {
+		const auto scale = duckdb::Interval::SECS_PER_DAY * duckdb::Interval::MSECS_PER_SEC;
+		value = Napi::Date::New(env, double(dval.GetValue<int32_t>() * scale));
+	} break;
+	case duckdb::LogicalTypeId::TIMESTAMP:
+	case duckdb::LogicalTypeId::TIMESTAMP_TZ: {
+		value = Napi::Date::New(env, double(dval.GetValue<int64_t>() / duckdb::Interval::MICROS_PER_MSEC));
+	} break;
+#endif
+	case duckdb::LogicalTypeId::VARCHAR: {
+		value = Napi::String::New(env, duckdb::StringValue::Get(dval));
+	} break;
+	case duckdb::LogicalTypeId::BLOB: {
+		auto &blob = duckdb::StringValue::Get(dval);
+		value = Napi::Buffer<char>::Copy(env, blob.c_str(), blob.length());
+	} break;
+	case duckdb::LogicalTypeId::SQLNULL: {
+		value = env.Null();
+	} break;
+	case duckdb::LogicalTypeId::LIST: {
+		auto child_type = duckdb::ListType::GetChildType(dval.type());
+		auto &child_values = duckdb::ListValue::GetChildren(dval);
+		auto object_value = Napi::Array::New(env);
+		for (duckdb::idx_t child_idx = 0; child_idx < child_values.size(); child_idx++) {
+			auto child_value = child_values.at(child_idx);
+			object_value.Set(child_idx, convert_col_val(env, child_value, child_type.id()));
+		}
+		value = object_value;
+	} break;
+	case duckdb::LogicalTypeId::STRUCT: {
+		auto &child_types = duckdb::StructType::GetChildTypes(dval.type());
+		auto &child_values = duckdb::StructValue::GetChildren(dval);
+		auto object_value = Napi::Object::New(env);
+		for (duckdb::idx_t child_idx = 0; child_idx < child_values.size(); child_idx++) {
+			auto child_value = child_values.at(child_idx);
+			auto child_type = child_types.at(child_idx);
+			object_value.Set(child_type.first, convert_col_val(env, child_value, child_type.second.id()));
+		}
+		value = object_value;
+	} break;
+	default:
+		Napi::Error::New(env, "Data type is not supported " + dval.type().ToString()).ThrowAsJavaScriptException();
+		return env.Null();
+	}
+
+	return value;
+}
+
 static Napi::Value convert_chunk(Napi::Env &env, std::vector<std::string> names, duckdb::DataChunk &chunk) {
 	Napi::EscapableHandleScope scope(env);
 	std::vector<Napi::String> node_names;
@@ -134,75 +217,13 @@ static Napi::Value convert_chunk(Napi::Env &env, std::vector<std::string> names,
 		Napi::Object row_result = Napi::Object::New(env);
 
 		for (duckdb::idx_t col_idx = 0; col_idx < chunk.ColumnCount(); col_idx++) {
-			Napi::Value value;
-			// set up a new Napi::Object for some data types, e.g. INTERVAL
-			Napi::Object object_value;
-
-			bool is_object_value {false};
-
-			auto dval = chunk.GetValue(col_idx, row_idx);
+			duckdb::Value dval = chunk.GetValue(col_idx, row_idx);
 			if (dval.IsNull()) {
 				row_result.Set(node_names[col_idx], env.Null());
 				continue;
 			}
 
-			// TODO templateroo here
-			switch (chunk.data[col_idx].GetType().id()) {
-			case duckdb::LogicalTypeId::BOOLEAN: {
-				value = Napi::Boolean::New(env, duckdb::BooleanValue::Get(dval));
-			} break;
-			case duckdb::LogicalTypeId::INTEGER: {
-				value = Napi::Number::New(env, duckdb::IntegerValue::Get(dval));
-			} break;
-			case duckdb::LogicalTypeId::FLOAT: {
-				value = Napi::Number::New(env, duckdb::FloatValue::Get(dval));
-			} break;
-			case duckdb::LogicalTypeId::DOUBLE: {
-				value = Napi::Number::New(env, duckdb::DoubleValue::Get(dval));
-			} break;
-			case duckdb::LogicalTypeId::BIGINT: {
-				value = Napi::Number::New(env, duckdb::BigIntValue::Get(dval));
-			} break;
-			case duckdb::LogicalTypeId::HUGEINT: {
-				value = Napi::Number::New(env, dval.GetValue<double>());
-			} break;
-			case duckdb::LogicalTypeId::INTERVAL: {
-				auto interval = duckdb::IntervalValue::Get(dval);
-				is_object_value = true;
-				object_value = Napi::Object::New(env);
-				object_value.Set("months", interval.months);
-				object_value.Set("days", interval.days);
-				object_value.Set("micros", interval.micros);
-			} break;
-#if (NAPI_VERSION > 4)
-			case duckdb::LogicalTypeId::DATE: {
-				const auto scale = duckdb::Interval::SECS_PER_DAY * duckdb::Interval::MSECS_PER_SEC;
-				value = Napi::Date::New(env, double(dval.GetValue<int32_t>() * scale));
-			} break;
-			case duckdb::LogicalTypeId::TIMESTAMP: {
-				value = Napi::Date::New(env, double(dval.GetValue<int64_t>() / duckdb::Interval::MICROS_PER_MSEC));
-			} break;
-#endif
-			case duckdb::LogicalTypeId::VARCHAR: {
-				value = Napi::String::New(env, duckdb::StringValue::Get(dval));
-			} break;
-			case duckdb::LogicalTypeId::BLOB: {
-				auto &blob = duckdb::StringValue::Get(dval);
-				value = Napi::Buffer<char>::Copy(env, blob.c_str(), blob.length());
-			} break;
-			case duckdb::LogicalTypeId::SQLNULL: {
-				value = env.Null();
-			} break;
-			default:
-				Napi::Error::New(env, "Data type is not supported " + dval.type().ToString())
-				    .ThrowAsJavaScriptException();
-				return env.Null();
-			}
-			if (is_object_value == true) {
-				row_result.Set(node_names[col_idx], object_value);
-			} else {
-				row_result.Set(node_names[col_idx], value);
-			}
+			row_result.Set(node_names[col_idx], convert_col_val(env, dval, chunk.data[col_idx].GetType().id()));
 		}
 		result.Set(row_idx, row_result);
 	}

package/test/data_type_support.test.js

@@ -21,16 +21,73 @@ describe("data type support", function () {
     });
   });
   it("supports INTERVAL values", function (done) {
-    db.prepare(`SELECT 
+    db.prepare(`SELECT
     INTERVAL 1 MINUTE as minutes,
     INTERVAL 5 DAY as days,
     INTERVAL 4 MONTH as months,
-    INTERVAL 4 MONTH + INTERVAL 5 DAY + INTERVAL 1 MINUTE as combined;`).each((err, row) => {
+    INTERVAL 4 MONTH + INTERVAL 5 DAY + INTERVAL 1 MINUTE as combined;`
+    ).each((err, row) => {
       assert(err === null);
-      assert.deepEqual(row.minutes, { months: 0, days: 0, micros: 60 * 1000 * 1000});
-      assert.deepEqual(row.days, { months: 0, days: 5, micros: 0});
-      assert.deepEqual(row.months, {months: 4, days: 0, micros: 0});
-      assert.deepEqual(row.combined, {months: 4, days: 5, micros: 60 * 1000 * 1000});
+      assert.deepEqual(row.minutes, {
+        months: 0,
+        days: 0,
+        micros: 60 * 1000 * 1000,
+      });
+      assert.deepEqual(row.days, { months: 0, days: 5, micros: 0 });
+      assert.deepEqual(row.months, { months: 4, days: 0, micros: 0 });
+      assert.deepEqual(row.combined, {
+        months: 4,
+        days: 5,
+        micros: 60 * 1000 * 1000,
+      });
+      done();
+    });
+  });
+  it("supports STRUCT values", function (done) {
+    db.prepare(`SELECT {'x': 1, 'y': 2, 'z': {'a': 'b'}} as struct`).each(
+      (err, row) => {
+        assert.deepEqual(row.struct, { x: 1, y: 2, z: { a: "b" } });
+        done();
+      }
+    );
+  });
+  it("supports LIST values", function (done) {
+    db.prepare(`SELECT ['duck', 'duck', 'goose'] as list`).each((err, row) => {
+      assert.deepEqual(row.list, ["duck", "duck", "goose"]);
+      done();
+    });
+  });
+  it("supports DATE values", function (done) {
+    db.prepare(`SELECT '2021-01-01'::DATE as dt;`).each((err, row) => {
+      assert(err === null);
+      assert.deepEqual(row.dt, new Date(Date.UTC(2021, 0, 1)));
+      done();
+    });
+  });
+  it("supports TIMESTAMP values", function (done) {
+    db.prepare(`SELECT '2021-01-01T00:00:00'::TIMESTAMP as ts;`).each((err, row) => {
+      assert(err === null);
+      assert.deepEqual(row.ts, new Date(Date.UTC(2021, 0, 1)));
+      done();
+    });
+  });
+  it("supports TIMESTAMP WITH TIME ZONE values", function (done) {
+    db.prepare(`SELECT '2021-01-01T00:00:00Z'::TIMESTAMPTZ as tstz;`).each((err, row) => {
+      assert(err === null);
+      assert.deepEqual(row.tstz, new Date(Date.UTC(2021, 0, 1)));
+      done();
+    });
+  });
+  it("supports DECIMAL values", function (done) {
+    db.run("CREATE TABLE decimal_table (d DECIMAL(24, 6))");
+    const stmt = db.prepare("INSERT INTO decimal_table VALUES (?)");
+    const values = [0, -1, 23534642362547.543463];
+    values.forEach((d) => {
+      stmt.run(d);
+    });
+    db.prepare("SELECT d from decimal_table;").all((err, res) => {
+      assert(err === null);
+      assert(res.every((v, i) => v.d === values[i]));
       done();
     });
   });

package/test/prepare.test.js

@@ -543,7 +543,7 @@ describe('prepare', function() {
             });
             it("should aggregate approx_count_distinct(flt)", function (done) {
                 db.all("SELECT approx_count_distinct(flt) as approx_count_distinct FROM foo", function (err, res) {
-                    assert.ok(res[0].approx_count_distinct >= 1000);
+                    assert.ok(res[0].approx_count_distinct >= 950);
                     done(err);
                 });
             });
@@ -587,7 +587,7 @@ describe('prepare', function() {
             });
             it("should aggregate approx_count_distinct(num)", function (done) {
                 db.all("SELECT approx_count_distinct(num) as approx_count_distinct FROM foo", function (err, res) {
-                    assert.ok(res[0].approx_count_distinct >= 1000);
+                    assert.ok(res[0].approx_count_distinct >= 950);
                     done(err);
                 });
             });