duckdb 0.8.2-dev37.0 → 0.8.2-dev3949.0

package/src/duckdb/extension/parquet/parquet_writer.cpp

@@ -29,6 +29,30 @@ using duckdb_parquet::format::PageType;
 using ParquetRowGroup = duckdb_parquet::format::RowGroup;
 using duckdb_parquet::format::Type;
 
+ChildFieldIDs::ChildFieldIDs() {
+	ids = make_uniq<case_insensitive_map_t<FieldID>>();
+}
+
+ChildFieldIDs ChildFieldIDs::Copy() const {
+	ChildFieldIDs result;
+	for (const auto &id : *ids) {
+		result.ids->emplace(id.first, id.second.Copy());
+	}
+	return result;
+}
+
+FieldID::FieldID() : set(false) {
+}
+
+FieldID::FieldID(int32_t field_id_p) : set(true), field_id(field_id_p) {
+}
+
+FieldID FieldID::Copy() const {
+	auto result = set ? FieldID(field_id) : FieldID();
+	result.child_field_ids = child_field_ids.Copy();
+	return result;
+}
+
 class MyTransport : public TTransport {
 public:
 	explicit MyTransport(Serializer &serializer) : serializer(serializer) {
@@ -52,27 +76,32 @@ private:
 	Serializer &serializer;
 };
 
-Type::type ParquetWriter::DuckDBTypeToParquetType(const LogicalType &duckdb_type) {
+bool ParquetWriter::DuckDBTypeToParquetTypeInternal(const LogicalType &duckdb_type, Type::type &parquet_type) {
 	switch (duckdb_type.id()) {
 	case LogicalTypeId::BOOLEAN:
-		return Type::BOOLEAN;
+		parquet_type = Type::BOOLEAN;
+		break;
 	case LogicalTypeId::TINYINT:
 	case LogicalTypeId::SMALLINT:
 	case LogicalTypeId::INTEGER:
 	case LogicalTypeId::DATE:
-		return Type::INT32;
+		parquet_type = Type::INT32;
+		break;
 	case LogicalTypeId::BIGINT:
-		return Type::INT64;
+		parquet_type = Type::INT64;
+		break;
 	case LogicalTypeId::FLOAT:
-		return Type::FLOAT;
+		parquet_type = Type::FLOAT;
+		break;
 	case LogicalTypeId::DOUBLE:
 	case LogicalTypeId::HUGEINT:
-		return Type::DOUBLE;
+		parquet_type = Type::DOUBLE;
+		break;
 	case LogicalTypeId::ENUM:
 	case LogicalTypeId::BLOB:
 	case LogicalTypeId::VARCHAR:
-	case LogicalTypeId::BIT:
-		return Type::BYTE_ARRAY;
+		parquet_type = Type::BYTE_ARRAY;
+		break;
 	case LogicalTypeId::TIME:
 	case LogicalTypeId::TIME_TZ:
 	case LogicalTypeId::TIMESTAMP:
@@ -80,31 +109,80 @@ Type::type ParquetWriter::DuckDBTypeToParquetType(const LogicalType &duckdb_type
 	case LogicalTypeId::TIMESTAMP_MS:
 	case LogicalTypeId::TIMESTAMP_NS:
 	case LogicalTypeId::TIMESTAMP_SEC:
-		return Type::INT64;
+		parquet_type = Type::INT64;
+		break;
 	case LogicalTypeId::UTINYINT:
 	case LogicalTypeId::USMALLINT:
 	case LogicalTypeId::UINTEGER:
-		return Type::INT32;
+		parquet_type = Type::INT32;
+		break;
 	case LogicalTypeId::UBIGINT:
-		return Type::INT64;
+		parquet_type = Type::INT64;
+		break;
 	case LogicalTypeId::INTERVAL:
 	case LogicalTypeId::UUID:
-		return Type::FIXED_LEN_BYTE_ARRAY;
+		parquet_type = Type::FIXED_LEN_BYTE_ARRAY;
+		break;
 	case LogicalTypeId::DECIMAL:
 		switch (duckdb_type.InternalType()) {
 		case PhysicalType::INT16:
 		case PhysicalType::INT32:
-			return Type::INT32;
+			parquet_type = Type::INT32;
+			break;
 		case PhysicalType::INT64:
-			return Type::INT64;
+			parquet_type = Type::INT64;
+			break;
 		case PhysicalType::INT128:
-			return Type::FIXED_LEN_BYTE_ARRAY;
+			parquet_type = Type::FIXED_LEN_BYTE_ARRAY;
+			break;
 		default:
 			throw InternalException("Unsupported internal decimal type");
 		}
+		break;
 	default:
+		// Anything that is not supported returns false
+		return false;
+	}
+	return true;
+}
+
+Type::type ParquetWriter::DuckDBTypeToParquetType(const LogicalType &duckdb_type) {
+	Type::type result;
+	if (!DuckDBTypeToParquetTypeInternal(duckdb_type, result)) {
 		throw NotImplementedException("Unimplemented type for Parquet \"%s\"", duckdb_type.ToString());
 	}
+	return result;
+}
+
+bool ParquetWriter::TypeIsSupported(const LogicalType &type) {
+	Type::type unused;
+	auto id = type.id();
+	if (id == LogicalTypeId::LIST) {
+		auto &child_type = ListType::GetChildType(type);
+		return TypeIsSupported(child_type);
+	}
+	if (id == LogicalTypeId::STRUCT) {
+		auto &children = StructType::GetChildTypes(type);
+		for (auto &child : children) {
+			auto &child_type = child.second;
+			if (!TypeIsSupported(child_type)) {
+				return false;
+			}
+		}
+		return true;
+	}
+	if (id == LogicalTypeId::MAP) {
+		auto &key_type = MapType::KeyType(type);
+		auto &value_type = MapType::ValueType(type);
+		if (!TypeIsSupported(key_type)) {
+			return false;
+		}
+		if (!TypeIsSupported(value_type)) {
+			return false;
+		}
+		return true;
+	}
+	return DuckDBTypeToParquetTypeInternal(type, unused);
 }
 
 void ParquetWriter::SetSchemaProperties(const LogicalType &duckdb_type,
@@ -226,8 +304,9 @@ void VerifyUniqueNames(const vector<string> &names) {
 }
 
 ParquetWriter::ParquetWriter(FileSystem &fs, string file_name_p, vector<LogicalType> types_p, vector<string> names_p,
-                             CompressionCodec::type codec)
-    : file_name(std::move(file_name_p)), sql_types(std::move(types_p)), column_names(std::move(names_p)), codec(codec) {
+                             CompressionCodec::type codec, ChildFieldIDs field_ids_p)
+    : file_name(std::move(file_name_p)), sql_types(std::move(types_p)), column_names(std::move(names_p)), codec(codec),
+      field_ids(std::move(field_ids_p)) {
 	// initialize the file writer
 	writer = make_uniq<BufferedFileWriter>(fs, file_name.c_str(),
 	                                       FileFlags::FILE_FLAGS_WRITE | FileFlags::FILE_FLAGS_FILE_CREATE_NEW);
@@ -257,37 +336,73 @@ ParquetWriter::ParquetWriter(FileSystem &fs, string file_name_p, vector<LogicalT
 	vector<string> schema_path;
 	for (idx_t i = 0; i < sql_types.size(); i++) {
 		column_writers.push_back(ColumnWriter::CreateWriterRecursive(file_meta_data.schema, *this, sql_types[i],
-		                                                             unique_names[i], schema_path));
+		                                                             unique_names[i], schema_path, &field_ids));
 	}
 }
 
 void ParquetWriter::PrepareRowGroup(ColumnDataCollection &buffer, PreparedRowGroup &result) {
+	// We write 8 columns at a time so that iterating over ColumnDataCollection is more efficient
+	static constexpr idx_t COLUMNS_PER_PASS = 8;
+
+	// We want these to be in-memory/hybrid so we don't have to copy over strings to the dictionary
+	D_ASSERT(buffer.GetAllocatorType() == ColumnDataAllocatorType::IN_MEMORY_ALLOCATOR ||
+	         buffer.GetAllocatorType() == ColumnDataAllocatorType::HYBRID);
+
 	// set up a new row group for this chunk collection
 	auto &row_group = result.row_group;
 	row_group.num_rows = buffer.Count();
+	row_group.total_byte_size = buffer.SizeInBytes();
 	row_group.__isset.file_offset = true;
 
 	auto &states = result.states;
 	// iterate over each of the columns of the chunk collection and write them
 	D_ASSERT(buffer.ColumnCount() == column_writers.size());
-	for (idx_t col_idx = 0; col_idx < buffer.ColumnCount(); col_idx++) {
-		const auto &col_writer = column_writers[col_idx];
-		auto write_state = col_writer->InitializeWriteState(row_group, buffer.GetAllocator());
-		if (col_writer->HasAnalyze()) {
-			for (auto &chunk : buffer.Chunks()) {
-				col_writer->Analyze(*write_state, nullptr, chunk.data[col_idx], chunk.size());
+	for (idx_t col_idx = 0; col_idx < buffer.ColumnCount(); col_idx += COLUMNS_PER_PASS) {
+		const auto next = MinValue<idx_t>(buffer.ColumnCount() - col_idx, COLUMNS_PER_PASS);
+		vector<column_t> column_ids;
+		vector<reference<ColumnWriter>> col_writers;
+		vector<unique_ptr<ColumnWriterState>> write_states;
+		for (idx_t i = 0; i < next; i++) {
+			column_ids.emplace_back(col_idx + i);
+			col_writers.emplace_back(*column_writers[column_ids.back()]);
+			write_states.emplace_back(col_writers.back().get().InitializeWriteState(row_group));
+		}
+
+		for (auto &chunk : buffer.Chunks({column_ids})) {
+			for (idx_t i = 0; i < next; i++) {
+				if (col_writers[i].get().HasAnalyze()) {
+					col_writers[i].get().Analyze(*write_states[i], nullptr, chunk.data[i], chunk.size());
+				}
 			}
-			col_writer->FinalizeAnalyze(*write_state);
 		}
-		for (auto &chunk : buffer.Chunks()) {
-			col_writer->Prepare(*write_state, nullptr, chunk.data[col_idx], chunk.size());
+
+		for (idx_t i = 0; i < next; i++) {
+			if (col_writers[i].get().HasAnalyze()) {
+				col_writers[i].get().FinalizeAnalyze(*write_states[i]);
+			}
 		}
-		col_writer->BeginWrite(*write_state);
-		for (auto &chunk : buffer.Chunks()) {
-			col_writer->Write(*write_state, chunk.data[col_idx], chunk.size());
+
+		for (auto &chunk : buffer.Chunks({column_ids})) {
+			for (idx_t i = 0; i < next; i++) {
+				col_writers[i].get().Prepare(*write_states[i], nullptr, chunk.data[i], chunk.size());
+			}
+		}
+
+		for (idx_t i = 0; i < next; i++) {
+			col_writers[i].get().BeginWrite(*write_states[i]);
+		}
+
+		for (auto &chunk : buffer.Chunks({column_ids})) {
+			for (idx_t i = 0; i < next; i++) {
+				col_writers[i].get().Write(*write_states[i], chunk.data[i], chunk.size());
+			}
+		}
+
+		for (auto &write_state : write_states) {
+			states.push_back(std::move(write_state));
 		}
-		states.push_back(std::move(write_state));
 	}
+	result.heaps = buffer.GetHeapReferences();
 }
 
 void ParquetWriter::FlushRowGroup(PreparedRowGroup &prepared) {
@@ -307,6 +422,8 @@ void ParquetWriter::FlushRowGroup(PreparedRowGroup &prepared) {
 	// append the row group to the file meta data
 	file_meta_data.row_groups.push_back(row_group);
 	file_meta_data.num_rows += row_group.num_rows;
+
+	prepared.heaps.clear();
 }
 
 void ParquetWriter::Flush(ColumnDataCollection &buffer) {
@@ -316,6 +433,7 @@ void ParquetWriter::Flush(ColumnDataCollection &buffer) {
 
 	PreparedRowGroup prepared_row_group;
 	PrepareRowGroup(buffer, prepared_row_group);
+	buffer.Reset();
 
 	FlushRowGroup(prepared_row_group);
 }

package/src/duckdb/extension/parquet/serialize_parquet.cpp

@@ -0,0 +1,26 @@
+//===----------------------------------------------------------------------===//
+// This file is automatically generated by scripts/generate_serialization.py
+// Do not edit this file manually, your changes will be overwritten
+//===----------------------------------------------------------------------===//
+
+#include "duckdb/common/serializer/format_serializer.hpp"
+#include "duckdb/common/serializer/format_deserializer.hpp"
+#include "parquet_reader.hpp"
+
+namespace duckdb {
+
+void ParquetOptions::FormatSerialize(FormatSerializer &serializer) const {
+	serializer.WriteProperty(100, "binary_as_string", binary_as_string);
+	serializer.WriteProperty(101, "file_row_number", file_row_number);
+	serializer.WriteProperty(102, "file_options", file_options);
+}
+
+ParquetOptions ParquetOptions::FormatDeserialize(FormatDeserializer &deserializer) {
+	ParquetOptions result;
+	deserializer.ReadProperty(100, "binary_as_string", result.binary_as_string);
+	deserializer.ReadProperty(101, "file_row_number", result.file_row_number);
+	deserializer.ReadProperty(102, "file_options", result.file_options);
+	return result;
+}
+
+} // namespace duckdb

package/src/duckdb/extension/parquet/zstd_file_system.cpp

@@ -155,7 +155,7 @@ void ZstdStreamWrapper::Close() {
 
 class ZStdFile : public CompressedFile {
 public:
-	ZStdFile(duckdb::unique_ptr<FileHandle> child_handle_p, const string &path, bool write)
+	ZStdFile(unique_ptr<FileHandle> child_handle_p, const string &path, bool write)
 	    : CompressedFile(zstd_fs, std::move(child_handle_p), path) {
 		Initialize(write);
 	}
@@ -163,7 +163,7 @@ public:
 	ZStdFileSystem zstd_fs;
 };
 
-unique_ptr<FileHandle> ZStdFileSystem::OpenCompressedFile(duckdb::unique_ptr<FileHandle> handle, bool write) {
+unique_ptr<FileHandle> ZStdFileSystem::OpenCompressedFile(unique_ptr<FileHandle> handle, bool write) {
 	auto path = handle->path;
 	return make_uniq<ZStdFile>(std::move(handle), path, write);
 }