duckdb 0.8.2-dev33.0 → 0.8.2-dev3334.0

package/src/duckdb/src/common/arrow/arrow_converter.cpp

@@ -15,12 +15,21 @@
 
 namespace duckdb {
 
-void ArrowConverter::ToArrowArray(DataChunk &input, ArrowArray *out_array, ArrowOptions options) {
+void ArrowConverter::ToArrowArray(DataChunk &input, ArrowArray *out_array, ClientProperties options) {
 	ArrowAppender appender(input.GetTypes(), input.size(), std::move(options));
 	appender.Append(input, 0, input.size(), input.size());
 	*out_array = appender.Finalize();
 }
 
+unsafe_unique_array<char> AddName(const string &name) {
+	auto name_ptr = make_unsafe_uniq_array<char>(name.size() + 1);
+	for (size_t i = 0; i < name.size(); i++) {
+		name_ptr[i] = name[i];
+	}
+	name_ptr[name.size()] = '\0';
+	return name_ptr;
+}
+
 //===--------------------------------------------------------------------===//
 // Arrow Schema
 //===--------------------------------------------------------------------===//
@@ -34,6 +43,7 @@ struct DuckDBArrowSchemaHolder {
 	std::list<vector<ArrowSchema *>> nested_children_ptr;
 	//! This holds strings created to represent decimal types
 	vector<unsafe_unique_array<char>> owned_type_names;
+	vector<unsafe_unique_array<char>> owned_column_names;
 };
 
 static void ReleaseDuckDBArrowSchema(ArrowSchema *schema) {
@@ -45,24 +55,27 @@ static void ReleaseDuckDBArrowSchema(ArrowSchema *schema) {
 	delete holder;
 }
 
-void InitializeChild(ArrowSchema &child, const string &name = "") {
+void InitializeChild(ArrowSchema &child, DuckDBArrowSchemaHolder &root_holder, const string &name = "") {
 	//! Child is cleaned up by parent
 	child.private_data = nullptr;
 	child.release = ReleaseDuckDBArrowSchema;
 
-	//! Store the child schema
+	// Store the child schema
 	child.flags = ARROW_FLAG_NULLABLE;
-	child.name = name.c_str();
+	root_holder.owned_type_names.push_back(AddName(name));
+
+	child.name = root_holder.owned_type_names.back().get();
 	child.n_children = 0;
 	child.children = nullptr;
 	child.metadata = nullptr;
 	child.dictionary = nullptr;
 }
+
 void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, const LogicalType &type,
-                    const ArrowOptions &options);
+                    const ClientProperties &options);
 
 void SetArrowMapFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, const LogicalType &type,
-                       const ArrowOptions &options) {
+                       const ClientProperties &options) {
 	child.format = "+m";
 	//! Map has one child which is a struct
 	child.n_children = 1;
@@ -70,14 +83,14 @@ void SetArrowMapFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child,
 	root_holder.nested_children.back().resize(1);
 	root_holder.nested_children_ptr.emplace_back();
 	root_holder.nested_children_ptr.back().push_back(&root_holder.nested_children.back()[0]);
-	InitializeChild(root_holder.nested_children.back()[0]);
+	InitializeChild(root_holder.nested_children.back()[0], root_holder);
 	child.children = &root_holder.nested_children_ptr.back()[0];
 	child.children[0]->name = "entries";
 	SetArrowFormat(root_holder, **child.children, ListType::GetChildType(type), options);
 }
 
 void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, const LogicalType &type,
-                    const ArrowOptions &options) {
+                    const ClientProperties &options) {
 	switch (type.id()) {
 	case LogicalTypeId::BOOLEAN:
 		child.format = "b";
@@ -117,7 +130,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		break;
 	case LogicalTypeId::UUID:
 	case LogicalTypeId::VARCHAR:
-		if (options.offset_size == ArrowOffsetSize::LARGE) {
+		if (options.arrow_offset_size == ArrowOffsetSize::LARGE) {
 			child.format = "U";
 		} else {
 			child.format = "u";
@@ -127,7 +140,6 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		child.format = "tdD";
 		break;
 	case LogicalTypeId::TIME:
-	case LogicalTypeId::TIME_TZ:
 		child.format = "ttu";
 		break;
 	case LogicalTypeId::TIMESTAMP:
@@ -135,12 +147,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		break;
 	case LogicalTypeId::TIMESTAMP_TZ: {
 		string format = "tsu:" + options.time_zone;
-		auto format_ptr = make_unsafe_uniq_array<char>(format.size() + 1);
-		for (size_t i = 0; i < format.size(); i++) {
-			format_ptr[i] = format[i];
-		}
-		format_ptr[format.size()] = '\0';
-		root_holder.owned_type_names.push_back(std::move(format_ptr));
+		root_holder.owned_type_names.push_back(AddName(format));
 		child.format = root_holder.owned_type_names.back().get();
 		break;
 	}
@@ -160,12 +167,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		uint8_t width, scale;
 		type.GetDecimalProperties(width, scale);
 		string format = "d:" + to_string(width) + "," + to_string(scale);
-		auto format_ptr = make_unsafe_uniq_array<char>(format.size() + 1);
-		for (size_t i = 0; i < format.size(); i++) {
-			format_ptr[i] = format[i];
-		}
-		format_ptr[format.size()] = '\0';
-		root_holder.owned_type_names.push_back(std::move(format_ptr));
+		root_holder.owned_type_names.push_back(AddName(format));
 		child.format = root_holder.owned_type_names.back().get();
 		break;
 	}
@@ -175,7 +177,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 	}
 	case LogicalTypeId::BLOB:
 	case LogicalTypeId::BIT: {
-		if (options.offset_size == ArrowOffsetSize::LARGE) {
+		if (options.arrow_offset_size == ArrowOffsetSize::LARGE) {
 			child.format = "Z";
 		} else {
 			child.format = "z";
@@ -189,7 +191,7 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		root_holder.nested_children.back().resize(1);
 		root_holder.nested_children_ptr.emplace_back();
 		root_holder.nested_children_ptr.back().push_back(&root_holder.nested_children.back()[0]);
-		InitializeChild(root_holder.nested_children.back()[0]);
+		InitializeChild(root_holder.nested_children.back()[0], root_holder);
 		child.children = &root_holder.nested_children_ptr.back()[0];
 		child.children[0]->name = "l";
 		SetArrowFormat(root_holder, **child.children, ListType::GetChildType(type), options);
@@ -209,15 +211,9 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		child.children = &root_holder.nested_children_ptr.back()[0];
 		for (size_t type_idx = 0; type_idx < child_types.size(); type_idx++) {
 
-			InitializeChild(*child.children[type_idx]);
+			InitializeChild(*child.children[type_idx], root_holder);
 
-			auto &struct_col_name = child_types[type_idx].first;
-			auto name_ptr = make_unsafe_uniq_array<char>(struct_col_name.size() + 1);
-			for (size_t i = 0; i < struct_col_name.size(); i++) {
-				name_ptr[i] = struct_col_name[i];
-			}
-			name_ptr[struct_col_name.size()] = '\0';
-			root_holder.owned_type_names.push_back(std::move(name_ptr));
+			root_holder.owned_type_names.push_back(AddName(child_types[type_idx].first));
 
 			child.children[type_idx]->name = root_holder.owned_type_names.back().get();
 			SetArrowFormat(root_holder, *child.children[type_idx], child_types[type_idx].second, options);
@@ -228,6 +224,38 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		SetArrowMapFormat(root_holder, child, type, options);
 		break;
 	}
+	case LogicalTypeId::UNION: {
+		std::string format = "+us:";
+
+		auto &child_types = UnionType::CopyMemberTypes(type);
+		child.n_children = child_types.size();
+		root_holder.nested_children.emplace_back();
+		root_holder.nested_children.back().resize(child_types.size());
+		root_holder.nested_children_ptr.emplace_back();
+		root_holder.nested_children_ptr.back().resize(child_types.size());
+		for (idx_t type_idx = 0; type_idx < child_types.size(); type_idx++) {
+			root_holder.nested_children_ptr.back()[type_idx] = &root_holder.nested_children.back()[type_idx];
+		}
+		child.children = &root_holder.nested_children_ptr.back()[0];
+		for (size_t type_idx = 0; type_idx < child_types.size(); type_idx++) {
+
+			InitializeChild(*child.children[type_idx], root_holder);
+
+			root_holder.owned_type_names.push_back(AddName(child_types[type_idx].first));
+
+			child.children[type_idx]->name = root_holder.owned_type_names.back().get();
+			SetArrowFormat(root_holder, *child.children[type_idx], child_types[type_idx].second, options);
+
+			format += to_string(type_idx) + ",";
+		}
+
+		format.pop_back();
+
+		root_holder.owned_type_names.push_back(AddName(format));
+		child.format = root_holder.owned_type_names.back().get();
+
+		break;
+	}
 	case LogicalTypeId::ENUM: {
 		// TODO what do we do with pointer enums here?
 		switch (EnumType::GetPhysicalType(type)) {
@@ -247,18 +275,18 @@ void SetArrowFormat(DuckDBArrowSchemaHolder &root_holder, ArrowSchema &child, co
 		root_holder.nested_children.back().resize(1);
 		root_holder.nested_children_ptr.emplace_back();
 		root_holder.nested_children_ptr.back().push_back(&root_holder.nested_children.back()[0]);
-		InitializeChild(root_holder.nested_children.back()[0]);
+		InitializeChild(root_holder.nested_children.back()[0], root_holder);
 		child.dictionary = root_holder.nested_children_ptr.back()[0];
 		child.dictionary->format = "u";
 		break;
 	}
 	default:
-		throw InternalException("Unsupported Arrow type " + type.ToString());
+		throw NotImplementedException("Unsupported Arrow type " + type.ToString());
 	}
 }
 
 void ArrowConverter::ToArrowSchema(ArrowSchema *out_schema, const vector<LogicalType> &types,
-                                   const vector<string> &names, const ArrowOptions &options) {
+                                   const vector<string> &names, const ClientProperties &options) {
 	D_ASSERT(out_schema);
 	D_ASSERT(types.size() == names.size());
 	idx_t column_count = types.size();
@@ -283,9 +311,9 @@ void ArrowConverter::ToArrowSchema(ArrowSchema *out_schema, const vector<Logical
 
 	// Configure all child schemas
 	for (idx_t col_idx = 0; col_idx < column_count; col_idx++) {
-
+		root_holder->owned_column_names.push_back(AddName(names[col_idx]));
 		auto &child = root_holder->children[col_idx];
-		InitializeChild(child, names[col_idx]);
+		InitializeChild(child, *root_holder, names[col_idx]);
 		SetArrowFormat(*root_holder, child, types[col_idx], options);
 	}
 

package/src/duckdb/src/common/arrow/arrow_wrapper.cpp

@@ -9,6 +9,7 @@
 #include "duckdb/common/arrow/result_arrow_wrapper.hpp"
 #include "duckdb/common/arrow/arrow_appender.hpp"
 #include "duckdb/main/query_result.hpp"
+#include "duckdb/main/chunk_scan_state/query_result.hpp"
 
 namespace duckdb {
 
@@ -68,7 +69,7 @@ int ResultArrowArrayStreamWrapper::MyStreamGetSchema(struct ArrowArrayStream *st
 	auto my_stream = reinterpret_cast<ResultArrowArrayStreamWrapper *>(stream->private_data);
 	if (!my_stream->column_types.empty()) {
 		ArrowConverter::ToArrowSchema(out, my_stream->column_types, my_stream->column_names,
-		                              QueryResult::GetArrowOptions(*my_stream->result));
+		                              my_stream->result->client_properties);
 		return 0;
 	}
 
@@ -89,7 +90,7 @@ int ResultArrowArrayStreamWrapper::MyStreamGetSchema(struct ArrowArrayStream *st
 		my_stream->column_names = result.names;
 	}
 	ArrowConverter::ToArrowSchema(out, my_stream->column_types, my_stream->column_names,
-	                              QueryResult::GetArrowOptions(*my_stream->result));
+	                              my_stream->result->client_properties);
 	return 0;
 }
 
@@ -99,6 +100,7 @@ int ResultArrowArrayStreamWrapper::MyStreamGetNext(struct ArrowArrayStream *stre
 	}
 	auto my_stream = reinterpret_cast<ResultArrowArrayStreamWrapper *>(stream->private_data);
 	auto &result = *my_stream->result;
+	auto &scan_state = *my_stream->scan_state;
 	if (result.HasError()) {
 		my_stream->last_error = result.GetErrorObject();
 		return -1;
@@ -117,7 +119,8 @@ int ResultArrowArrayStreamWrapper::MyStreamGetNext(struct ArrowArrayStream *stre
 	}
 	idx_t result_count;
 	PreservedError error;
-	if (!ArrowUtil::TryFetchChunk(&result, my_stream->batch_size, out, result_count, error)) {
+	if (!ArrowUtil::TryFetchChunk(scan_state, result.client_properties, my_stream->batch_size, out, result_count,
+	                              error)) {
 		D_ASSERT(error);
 		my_stream->last_error = error;
 		return -1;
@@ -147,7 +150,7 @@ const char *ResultArrowArrayStreamWrapper::MyStreamGetLastError(struct ArrowArra
 }
 
 ResultArrowArrayStreamWrapper::ResultArrowArrayStreamWrapper(unique_ptr<QueryResult> result_p, idx_t batch_size_p)
-    : result(std::move(result_p)) {
+    : result(std::move(result_p)), scan_state(make_uniq<QueryResultChunkScanState>(*result)) {
 	//! We first initialize the private data of the stream
 	stream.private_data = this;
 	//! Ceil Approx_Batch_Size/STANDARD_VECTOR_SIZE
@@ -162,51 +165,43 @@ ResultArrowArrayStreamWrapper::ResultArrowArrayStreamWrapper(unique_ptr<QueryRes
 	stream.get_last_error = ResultArrowArrayStreamWrapper::MyStreamGetLastError;
 }
 
-bool ArrowUtil::TryFetchNext(QueryResult &result, unique_ptr<DataChunk> &chunk, PreservedError &error) {
-	if (result.type == QueryResultType::STREAM_RESULT) {
-		auto &stream_result = result.Cast<StreamQueryResult>();
-		if (!stream_result.IsOpen()) {
-			return true;
-		}
-	}
-	return result.TryFetch(chunk, error);
-}
-
-bool ArrowUtil::TryFetchChunk(QueryResult *result, idx_t chunk_size, ArrowArray *out, idx_t &count,
-                              PreservedError &error) {
+bool ArrowUtil::TryFetchChunk(ChunkScanState &scan_state, ClientProperties options, idx_t batch_size, ArrowArray *out,
+                              idx_t &count, PreservedError &error) {
 	count = 0;
-	ArrowAppender appender(result->types, chunk_size, QueryResult::GetArrowOptions(*result));
-	auto &current_chunk = result->current_chunk;
-	if (current_chunk.Valid()) {
+	ArrowAppender appender(scan_state.Types(), batch_size, std::move(options));
+	auto remaining_tuples_in_chunk = scan_state.RemainingInChunk();
+	if (remaining_tuples_in_chunk) {
 		// We start by scanning the non-finished current chunk
-		idx_t cur_consumption = current_chunk.RemainingSize() > chunk_size ? chunk_size : current_chunk.RemainingSize();
+		idx_t cur_consumption = MinValue(remaining_tuples_in_chunk, batch_size);
 		count += cur_consumption;
-		appender.Append(*current_chunk.data_chunk, current_chunk.position, current_chunk.position + cur_consumption,
-		                current_chunk.data_chunk->size());
-		current_chunk.position += cur_consumption;
-	}
-	while (count < chunk_size) {
-		unique_ptr<DataChunk> data_chunk;
-		if (!TryFetchNext(*result, data_chunk, error)) {
-			if (result->HasError()) {
-				error = result->GetErrorObject();
+		auto &current_chunk = scan_state.CurrentChunk();
+		appender.Append(current_chunk, scan_state.CurrentOffset(), scan_state.CurrentOffset() + cur_consumption,
+		                current_chunk.size());
+		scan_state.IncreaseOffset(cur_consumption);
+	}
+	while (count < batch_size) {
+		if (!scan_state.LoadNextChunk(error)) {
+			if (scan_state.HasError()) {
+				error = scan_state.GetError();
 			}
 			return false;
 		}
-		if (!data_chunk || data_chunk->size() == 0) {
+		if (scan_state.ChunkIsEmpty()) {
+			// The scan was successful, but an empty chunk was returned
 			break;
 		}
-		if (count + data_chunk->size() > chunk_size) {
-			// We have to split the chunk between this and the next batch
-			idx_t available_space = chunk_size - count;
-			appender.Append(*data_chunk, 0, available_space, data_chunk->size());
-			count += available_space;
-			current_chunk.data_chunk = std::move(data_chunk);
-			current_chunk.position = available_space;
-		} else {
-			count += data_chunk->size();
-			appender.Append(*data_chunk, 0, data_chunk->size(), data_chunk->size());
+		auto &current_chunk = scan_state.CurrentChunk();
+		if (scan_state.Finished() || current_chunk.size() == 0) {
+			break;
 		}
+		// The amount we still need to append into this chunk
+		auto remaining = batch_size - count;
+
+		// The amount remaining, capped by the amount left in the current chunk
+		auto to_append_to_batch = MinValue(remaining, scan_state.RemainingInChunk());
+		appender.Append(current_chunk, 0, to_append_to_batch, current_chunk.size());
+		count += to_append_to_batch;
+		scan_state.IncreaseOffset(to_append_to_batch);
 	}
 	if (count > 0) {
 		*out = appender.Finalize();
@@ -214,10 +209,10 @@ bool ArrowUtil::TryFetchChunk(QueryResult *result, idx_t chunk_size, ArrowArray
 	return true;
 }
 
-idx_t ArrowUtil::FetchChunk(QueryResult *result, idx_t chunk_size, ArrowArray *out) {
+idx_t ArrowUtil::FetchChunk(ChunkScanState &scan_state, ClientProperties options, idx_t chunk_size, ArrowArray *out) {
 	PreservedError error;
 	idx_t result_count;
-	if (!TryFetchChunk(result, chunk_size, out, result_count, error)) {
+	if (!TryFetchChunk(scan_state, std::move(options), chunk_size, out, result_count, error)) {
 		error.Throw();
 	}
 	return result_count;

package/src/duckdb/src/common/assert.cpp

@@ -4,6 +4,9 @@
 namespace duckdb {
 
 void DuckDBAssertInternal(bool condition, const char *condition_name, const char *file, int linenr) {
+#ifdef DISABLE_ASSERTIONS
+	return;
+#endif
 	if (condition) {
 		return;
 	}

package/src/duckdb/src/common/constants.cpp

@@ -7,7 +7,8 @@
 namespace duckdb {
 
 constexpr const idx_t DConstants::INVALID_INDEX;
-const row_t MAX_ROW_ID = 4611686018427388000ULL; // 2^62
+const row_t MAX_ROW_ID = 36028797018960000ULL;       // 2^55
+const row_t MAX_ROW_ID_LOCAL = 72057594037920000ULL; // 2^56
 const column_t COLUMN_IDENTIFIER_ROW_ID = (column_t)-1;
 const sel_t ZERO_VECTOR[STANDARD_VECTOR_SIZE] = {0};
 const double PI = 3.141592653589793;