Merge branch 'main' into main

Author: eramitmittal

CMakeLists.txt

@@ -3,8 +3,10 @@
 #include "duckdb.hpp"
 #include "geo_parquet.hpp"
 #include "parquet_dbp_encoder.hpp"
+#include "parquet_dlba_encoder.hpp"
 #include "parquet_rle_bp_decoder.hpp"
 #include "parquet_rle_bp_encoder.hpp"
+#include "parquet_bss_encoder.hpp"
 #include "parquet_statistics.hpp"
 #include "parquet_writer.hpp"
 #ifndef DUCKDB_AMALGAMATION
@@ -1059,34 +1061,42 @@ class StandardColumnWriterState : public BasicColumnWriterState {
 	}
 	~StandardColumnWriterState() override = default;
 
-	// analysis state for integer values for DELTA_BINARY_PACKED
+	// analysis state for integer values for DELTA_BINARY_PACKED/DELTA_LENGTH_BYTE_ARRAY
 	idx_t total_value_count = 0;
+	idx_t total_string_size = 0;
 
 	unordered_map<T, uint32_t> dictionary;
 	duckdb_parquet::Encoding::type encoding;
 };
 
-template <class T>
+template <class SRC, class TGT>
 class StandardWriterPageState : public ColumnWriterPageState {
 public:
-	explicit StandardWriterPageState(const idx_t total_value_count, Encoding::type encoding_p,
-	                                 const unordered_map<T, uint32_t> &dictionary_p)
-	    : encoding(encoding_p), dbp_initialized(false), dbp_encoder(total_value_count), dictionary(dictionary_p),
-	      dict_written_value(false), dict_bit_width(RleBpDecoder::ComputeBitWidth(dictionary.size())),
-	      dict_encoder(dict_bit_width) {
+	explicit StandardWriterPageState(const idx_t total_value_count, const idx_t total_string_size,
+	                                 Encoding::type encoding_p, const unordered_map<SRC, uint32_t> &dictionary_p)
+	    : encoding(encoding_p), dbp_initialized(false), dbp_encoder(total_value_count), dlba_initialized(false),
+	      dlba_encoder(total_value_count, total_string_size), bss_encoder(total_value_count, sizeof(TGT)),
+	      dictionary(dictionary_p), dict_written_value(false),
+	      dict_bit_width(RleBpDecoder::ComputeBitWidth(dictionary.size())), dict_encoder(dict_bit_width) {
 	}
 	duckdb_parquet::Encoding::type encoding;
 
 	bool dbp_initialized;
 	DbpEncoder dbp_encoder;
 
-	const unordered_map<T, uint32_t> &dictionary;
+	bool dlba_initialized;
+	DlbaEncoder dlba_encoder;
+
+	BssEncoder bss_encoder;
+
+	const unordered_map<SRC, uint32_t> &dictionary;
 	bool dict_written_value;
 	uint32_t dict_bit_width;
 	RleBpEncoder dict_encoder;
 };
 
 namespace dbp_encoder {
+
 template <class T>
 void BeginWrite(DbpEncoder &encoder, WriteStream &writer, const T &first_value) {
 	throw InternalException("Can't write type to DELTA_BINARY_PACKED column");
@@ -1139,6 +1149,60 @@ void WriteValue(DbpEncoder &encoder, WriteStream &writer, const uint32_t &value)
 
 } // namespace dbp_encoder
 
+namespace dlba_encoder {
+
+template <class T>
+void BeginWrite(DlbaEncoder &encoder, WriteStream &writer, const T &first_value) {
+	throw InternalException("Can't write type to DELTA_LENGTH_BYTE_ARRAY column");
+}
+
+template <>
+void BeginWrite(DlbaEncoder &encoder, WriteStream &writer, const string_t &first_value) {
+	encoder.BeginWrite(writer, first_value);
+}
+
+template <class T>
+void WriteValue(DlbaEncoder &encoder, WriteStream &writer, const T &value) {
+	throw InternalException("Can't write type to DELTA_LENGTH_BYTE_ARRAY column");
+}
+
+template <>
+void WriteValue(DlbaEncoder &encoder, WriteStream &writer, const string_t &value) {
+	encoder.WriteValue(writer, value);
+}
+
+// helpers to get size from strings
+template <class SRC>
+static constexpr idx_t GetDlbaStringSize(const SRC &src_value) {
+	return 0;
+}
+
+template <>
+idx_t GetDlbaStringSize(const string_t &src_value) {
+	return src_value.GetSize();
+}
+
+} // namespace dlba_encoder
+
+namespace bss_encoder {
+
+template <class T>
+void WriteValue(BssEncoder &encoder, const T &value) {
+	throw InternalException("Can't write type to BYTE_STREAM_SPLIT column");
+}
+
+template <>
+void WriteValue(BssEncoder &encoder, const float &value) {
+	encoder.WriteValue(value);
+}
+
+template <>
+void WriteValue(BssEncoder &encoder, const double &value) {
+	encoder.WriteValue(value);
+}
+
+} // namespace bss_encoder
+
 template <class SRC, class TGT, class OP = ParquetCastOperator>
 class StandardColumnWriter : public BasicColumnWriter {
 public:
@@ -1159,13 +1223,13 @@ class StandardColumnWriter : public BasicColumnWriter {
 	unique_ptr<ColumnWriterPageState> InitializePageState(BasicColumnWriterState &state_p) override {
 		auto &state = state_p.Cast<StandardColumnWriterState<SRC>>();
 
-		auto result =
-		    make_uniq<StandardWriterPageState<SRC>>(state.total_value_count, state.encoding, state.dictionary);
+		auto result = make_uniq<StandardWriterPageState<SRC, TGT>>(state.total_value_count, state.total_string_size,
+		                                                           state.encoding, state.dictionary);
 		return std::move(result);
 	}
 
 	void FlushPageState(WriteStream &temp_writer, ColumnWriterPageState *state_p) override {
-		auto &page_state = state_p->Cast<StandardWriterPageState<SRC>>();
+		auto &page_state = state_p->Cast<StandardWriterPageState<SRC, TGT>>();
 		switch (page_state.encoding) {
 		case Encoding::DELTA_BINARY_PACKED:
 			if (!page_state.dbp_initialized) {
@@ -1182,7 +1246,15 @@ class StandardColumnWriter : public BasicColumnWriter {
 				return;
 			}
 			page_state.dict_encoder.FinishWrite(temp_writer);
-
+			break;
+		case Encoding::DELTA_LENGTH_BYTE_ARRAY:
+			if (!page_state.dlba_initialized) {
+				dlba_encoder::BeginWrite<string_t>(page_state.dlba_encoder, temp_writer, string_t(""));
+			}
+			page_state.dlba_encoder.FinishWrite(temp_writer);
+			break;
+		case Encoding::BYTE_STREAM_SPLIT:
+			page_state.bss_encoder.FinishWrite(temp_writer);
 			break;
 		case Encoding::PLAIN:
 			break;
@@ -1220,14 +1292,15 @@ class StandardColumnWriter : public BasicColumnWriter {
 				continue;
 			}
 			if (validity.RowIsValid(vector_index)) {
+				const auto &src_value = data_ptr[vector_index];
 				if (state.dictionary.size() <= writer.DictionarySizeLimit()) {
-					const auto &src_value = data_ptr[vector_index];
 					if (state.dictionary.find(src_value) == state.dictionary.end()) {
 						state.dictionary[src_value] = new_value_index;
 						new_value_index++;
 					}
 				}
 				state.total_value_count++;
+				state.total_string_size += dlba_encoder::GetDlbaStringSize(src_value);
 			}
 			vector_index++;
 		}
@@ -1238,9 +1311,22 @@ class StandardColumnWriter : public BasicColumnWriter {
 
 		auto &state = state_p.Cast<StandardColumnWriterState<SRC>>();
 		if (state.dictionary.size() == 0 || state.dictionary.size() > writer.DictionarySizeLimit()) {
-			// special handling for int column: dpb, otherwise plain
-			state.encoding = (type == Type::type::INT32 || type == Type::type::INT64) ? Encoding::DELTA_BINARY_PACKED
-			                                                                          : Encoding::PLAIN;
+			// If we aren't doing dictionary encoding, the following encodings are virtually always better than PLAIN
+			switch (type) {
+			case Type::type::INT32:
+			case Type::type::INT64:
+				state.encoding = Encoding::DELTA_BINARY_PACKED;
+				break;
+			case Type::type::BYTE_ARRAY:
+				state.encoding = Encoding::DELTA_LENGTH_BYTE_ARRAY;
+				break;
+			case Type::type::FLOAT:
+			case Type::type::DOUBLE:
+				state.encoding = Encoding::BYTE_STREAM_SPLIT;
+				break;
+			default:
+				state.encoding = Encoding::PLAIN;
+			}
 			state.dictionary.clear();
 		}
 	}
@@ -1261,7 +1347,7 @@ class StandardColumnWriter : public BasicColumnWriter {
 
 	void WriteVector(WriteStream &temp_writer, ColumnWriterStatistics *stats, ColumnWriterPageState *page_state_p,
 	                 Vector &input_column, idx_t chunk_start, idx_t chunk_end) override {
-		auto &page_state = page_state_p->Cast<StandardWriterPageState<SRC>>();
+		auto &page_state = page_state_p->Cast<StandardWriterPageState<SRC, TGT>>();
 
 		const auto &mask = FlatVector::Validity(input_column);
 		const auto *data_ptr = FlatVector::GetData<SRC>(input_column);
@@ -1287,7 +1373,6 @@ class StandardColumnWriter : public BasicColumnWriter {
 			}
 			break;
 		}
-
 		case Encoding::DELTA_BINARY_PACKED: {
 			idx_t r = chunk_start;
 			if (!page_state.dbp_initialized) {
@@ -1315,6 +1400,44 @@ class StandardColumnWriter : public BasicColumnWriter {
 			}
 			break;
 		}
+		case Encoding::DELTA_LENGTH_BYTE_ARRAY: {
+			idx_t r = chunk_start;
+			if (!page_state.dlba_initialized) {
+				// find first non-null value
+				for (; r < chunk_end; r++) {
+					if (!mask.RowIsValid(r)) {
+						continue;
+					}
+					const TGT target_value = OP::template Operation<SRC, TGT>(data_ptr[r]);
+					OP::template HandleStats<SRC, TGT>(stats, target_value);
+					dlba_encoder::BeginWrite(page_state.dlba_encoder, temp_writer, target_value);
+					page_state.dlba_initialized = true;
+					r++; // skip over
+					break;
+				}
+			}
+
+			for (; r < chunk_end; r++) {
+				if (!mask.RowIsValid(r)) {
+					continue;
+				}
+				const TGT target_value = OP::template Operation<SRC, TGT>(data_ptr[r]);
+				OP::template HandleStats<SRC, TGT>(stats, target_value);
+				dlba_encoder::WriteValue(page_state.dlba_encoder, temp_writer, target_value);
+			}
+			break;
+		}
+		case Encoding::BYTE_STREAM_SPLIT: {
+			for (idx_t r = chunk_start; r < chunk_end; r++) {
+				if (!mask.RowIsValid(r)) {
+					continue;
+				}
+				const TGT target_value = OP::template Operation<SRC, TGT>(data_ptr[r]);
+				OP::template HandleStats<SRC, TGT>(stats, target_value);
+				bss_encoder::WriteValue(page_state.bss_encoder, target_value);
+			}
+			break;
+		}
 		case Encoding::PLAIN: {
 			D_ASSERT(page_state.encoding == Encoding::PLAIN);
 			TemplatedWritePlain<SRC, TGT, OP>(input_column, stats, chunk_start, chunk_end, mask, temp_writer);

src/duckdb/extension/parquet/column_writer.cpp

@@ -0,0 +1,45 @@
+//===----------------------------------------------------------------------===//
+//                         DuckDB
+//
+// parquet_bss_encoder.hpp
+//
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "decode_utils.hpp"
+
+namespace duckdb {
+
+class BssEncoder {
+public:
+	explicit BssEncoder(const idx_t total_value_count_p, const idx_t bit_width_p)
+	    : total_value_count(total_value_count_p), bit_width(bit_width_p), count(0),
+	      buffer(Allocator::DefaultAllocator().Allocate(total_value_count * bit_width + 1)) {
+	}
+
+public:
+	template <class T>
+	void WriteValue(const T &value) {
+		D_ASSERT(sizeof(T) == bit_width);
+		for (idx_t i = 0; i < sizeof(T); i++) {
+			buffer.get()[i * total_value_count + count] = reinterpret_cast<const_data_ptr_t>(&value)[i];
+		}
+		count++;
+	}
+
+	void FinishWrite(WriteStream &writer) {
+		D_ASSERT(count == total_value_count);
+		writer.WriteData(buffer.get(), total_value_count * bit_width);
+	}
+
+private:
+	const idx_t total_value_count;
+	const idx_t bit_width;
+
+	idx_t count;
+	AllocatedData buffer;
+};
+
+} // namespace duckdb

src/duckdb/extension/parquet/include/parquet_bss_encoder.hpp

@@ -0,0 +1,48 @@
+//===----------------------------------------------------------------------===//
+//                         DuckDB
+//
+// parquet_dlba_encoder.hpp
+//
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "parquet_dbp_encoder.hpp"
+#include "duckdb/common/serializer/memory_stream.hpp"
+
+namespace duckdb {
+
+class DlbaEncoder {
+public:
+	DlbaEncoder(const idx_t total_value_count_p, const idx_t total_string_size_p)
+	    : dbp_encoder(total_value_count_p), total_string_size(total_string_size_p),
+	      buffer(Allocator::DefaultAllocator().Allocate(total_string_size + 1)),
+	      stream(make_unsafe_uniq<MemoryStream>(buffer.get(), buffer.GetSize())) {
+	}
+
+public:
+	void BeginWrite(WriteStream &writer, const string_t &first_value) {
+		dbp_encoder.BeginWrite(writer, UnsafeNumericCast<int64_t>(first_value.GetSize()));
+		stream->WriteData(const_data_ptr_cast(first_value.GetData()), first_value.GetSize());
+	}
+
+	void WriteValue(WriteStream &writer, const string_t &value) {
+		dbp_encoder.WriteValue(writer, UnsafeNumericCast<int64_t>(value.GetSize()));
+		stream->WriteData(const_data_ptr_cast(value.GetData()), value.GetSize());
+	}
+
+	void FinishWrite(WriteStream &writer) {
+		D_ASSERT(stream->GetPosition() == total_string_size);
+		dbp_encoder.FinishWrite(writer);
+		writer.WriteData(buffer.get(), total_string_size);
+	}
+
+private:
+	DbpEncoder dbp_encoder;
+	const idx_t total_string_size;
+	AllocatedData buffer;
+	unsafe_unique_ptr<MemoryStream> stream;
+};
+
+} // namespace duckdb

src/duckdb/extension/parquet/include/parquet_dlba_encoder.hpp

@@ -26,8 +26,10 @@
 #include "duckdb/parser/parsed_data/create_type_info.hpp"
 #include "duckdb/parser/parsed_data/create_view_info.hpp"
 #include "duckdb/parser/parsed_data/drop_info.hpp"
+#include "duckdb/parser/statement/create_statement.hpp"
 #include "duckdb/planner/parsed_data/bound_create_table_info.hpp"
 #include "duckdb/planner/binder.hpp"
+#include "duckdb/planner/expression_binder/index_binder.hpp"
 #include "duckdb/catalog/default/default_types.hpp"
 #include "duckdb/main/extension_entries.hpp"
 #include "duckdb/main/extension/generated_extension_loader.hpp"
@@ -304,6 +306,14 @@ optional_ptr<CatalogEntry> Catalog::CreateIndex(ClientContext &context, CreateIn
 	return CreateIndex(GetCatalogTransaction(context), info);
 }
 
+unique_ptr<LogicalOperator> Catalog::BindCreateIndex(Binder &binder, CreateStatement &stmt, TableCatalogEntry &table,
+                                                     unique_ptr<LogicalOperator> plan) {
+	D_ASSERT(plan->type == LogicalOperatorType::LOGICAL_GET);
+	auto create_index_info = unique_ptr_cast<CreateInfo, CreateIndexInfo>(std::move(stmt.info));
+	IndexBinder index_binder(binder, binder.context);
+	return index_binder.BindCreateIndex(binder.context, std::move(create_index_info), table, std::move(plan), nullptr);
+}
+
 unique_ptr<LogicalOperator> Catalog::BindAlterAddIndex(Binder &binder, TableCatalogEntry &table_entry,
                                                        unique_ptr<LogicalOperator> plan,
                                                        unique_ptr<CreateIndexInfo> create_info,
@@ -953,9 +963,13 @@ optional_ptr<SchemaCatalogEntry> Catalog::GetSchema(CatalogEntryRetriever &retri
                                                     QueryErrorContext error_context) {
 	auto entries = GetCatalogEntries(retriever, catalog_name, schema_name);
 	for (idx_t i = 0; i < entries.size(); i++) {
+		auto catalog = Catalog::GetCatalogEntry(retriever, entries[i].catalog);
+		if (!catalog) {
+			// skip if it is not an attached database
+			continue;
+		}
 		auto on_not_found = i + 1 == entries.size() ? if_not_found : OnEntryNotFound::RETURN_NULL;
-		auto &catalog = Catalog::GetCatalog(retriever, entries[i].catalog);
-		auto result = catalog.GetSchema(retriever.GetContext(), schema_name, on_not_found, error_context);
+		auto result = catalog->GetSchema(retriever.GetContext(), schema_name, on_not_found, error_context);
 		if (result) {
 			return result;
 		}