duckdb 1.1.3 → 1.1.4-dev11.0

package/src/duckdb/src/storage/compression/roaring/compress.cpp

@@ -0,0 +1,481 @@
+#include "duckdb/storage/compression/roaring/roaring.hpp"
+#include "duckdb/storage/compression/roaring/appender.hpp"
+
+#include "duckdb/common/limits.hpp"
+#include "duckdb/common/likely.hpp"
+#include "duckdb/common/numeric_utils.hpp"
+#include "duckdb/function/compression/compression.hpp"
+#include "duckdb/function/compression_function.hpp"
+#include "duckdb/main/config.hpp"
+#include "duckdb/storage/buffer_manager.hpp"
+#include "duckdb/storage/table/column_data_checkpointer.hpp"
+#include "duckdb/storage/table/column_segment.hpp"
+#include "duckdb/storage/table/scan_state.hpp"
+#include "duckdb/storage/segment/uncompressed.hpp"
+#include "duckdb/common/fast_mem.hpp"
+#include "duckdb/common/bitpacking.hpp"
+
+namespace duckdb {
+
+namespace roaring {
+
+ContainerCompressionState::ContainerCompressionState() {
+	Reset();
+}
+
+inline void AppendBitset(ContainerCompressionState &state, bool null, uint16_t amount) {
+	D_ASSERT(state.uncompressed);
+	if (null) {
+		ValidityMask mask(state.uncompressed, ROARING_CONTAINER_SIZE);
+		SetInvalidRange(mask, state.appended_count, state.appended_count + amount);
+	}
+}
+
+inline void AppendRun(ContainerCompressionState &state, bool null, uint16_t amount) {
+	// Adjust the run
+	auto run_idx = state.run_idx;
+	auto appended_count = state.appended_count;
+	if (!null && run_idx < MAX_RUN_IDX && appended_count && (null != state.last_is_null)) {
+		if (run_idx < COMPRESSED_RUN_THRESHOLD) {
+			auto &last_run = state.runs[run_idx];
+			// End the last run
+			last_run.length = (appended_count - last_run.start) - 1;
+		}
+		state.compressed_runs[(run_idx * 2) + 1] = static_cast<uint8_t>(appended_count & (COMPRESSED_SEGMENT_SIZE - 1));
+		state.run_counts[appended_count >> COMPRESSED_SEGMENT_SHIFT_AMOUNT]++;
+		state.run_idx++;
+	} else if (null && run_idx < MAX_RUN_IDX && (!appended_count || null != state.last_is_null)) {
+		if (run_idx < COMPRESSED_RUN_THRESHOLD) {
+			auto &current_run = state.runs[run_idx];
+			// Initialize a new run
+			current_run.start = appended_count;
+		}
+		state.compressed_runs[(run_idx * 2) + 0] = static_cast<uint8_t>(appended_count & (COMPRESSED_SEGMENT_SIZE - 1));
+		state.run_counts[appended_count >> COMPRESSED_SEGMENT_SHIFT_AMOUNT]++;
+	}
+}
+
+template <bool INVERTED>
+inline void AppendToArray(ContainerCompressionState &state, bool null, uint16_t amount) {
+	if (DUCKDB_LIKELY(INVERTED != null)) {
+		return;
+	}
+
+	auto current_array_idx = state.array_idx[null];
+	if (current_array_idx + amount > MAX_ARRAY_IDX) {
+		return;
+	}
+	auto appended_count = state.appended_count;
+	auto array_count = state.array_counts[null];
+	auto compressed_array = state.compressed_arrays[null];
+	uint16_t appended = 0;
+	while (appended < amount) {
+		uint16_t remaining = amount - appended;
+		uint8_t segment_offset = appended ? 0 : (appended_count + appended) & (COMPRESSED_SEGMENT_SIZE - 1);
+		uint8_t to_append =
+		    static_cast<uint8_t>(MinValue<uint16_t>(remaining, COMPRESSED_SEGMENT_SIZE - segment_offset));
+		for (uint8_t i = 0; i < to_append; i++) {
+			auto index = current_array_idx + appended + i;
+			compressed_array[index] = segment_offset + i;
+		}
+
+		idx_t segment_index = (appended_count + appended) / COMPRESSED_SEGMENT_SIZE;
+		array_count[segment_index] += to_append;
+		appended += to_append;
+	}
+
+	if (current_array_idx + amount < COMPRESSED_ARRAY_THRESHOLD) {
+		auto &array = state.arrays[null];
+		for (uint16_t i = 0; i < amount; i++) {
+			array[current_array_idx + i] = appended_count + i;
+		}
+	}
+	state.array_idx[null] += amount;
+}
+
+void ContainerCompressionState::Append(bool null, uint16_t amount) {
+	append_function(*this, null, amount);
+	last_is_null = null;
+	null_count += null * amount;
+	appended_count += amount;
+}
+
+void ContainerCompressionState::OverrideArray(data_ptr_t &destination, bool nulls, idx_t count) {
+	if (nulls) {
+		append_function = AppendToArray<true>;
+	} else {
+		append_function = AppendToArray<false>;
+	}
+
+	if (count >= COMPRESSED_ARRAY_THRESHOLD) {
+		memset(destination, 0, sizeof(uint8_t) * COMPRESSED_SEGMENT_COUNT);
+		array_counts[nulls] = reinterpret_cast<uint8_t *>(destination);
+		auto data_start = destination + sizeof(uint8_t) * COMPRESSED_SEGMENT_COUNT;
+		compressed_arrays[nulls] = reinterpret_cast<uint8_t *>(data_start);
+	} else {
+		destination = AlignValue<sizeof(uint16_t)>(destination);
+		arrays[nulls] = reinterpret_cast<uint16_t *>(destination);
+	}
+}
+
+void ContainerCompressionState::OverrideRun(data_ptr_t &destination, idx_t count) {
+	append_function = AppendRun;
+
+	if (count >= COMPRESSED_RUN_THRESHOLD) {
+		memset(destination, 0, sizeof(uint8_t) * COMPRESSED_SEGMENT_COUNT);
+		run_counts = reinterpret_cast<uint8_t *>(destination);
+		auto data_start = destination + sizeof(uint8_t) * COMPRESSED_SEGMENT_COUNT;
+		compressed_runs = reinterpret_cast<uint8_t *>(data_start);
+	} else {
+		destination = AlignValue<sizeof(RunContainerRLEPair)>(destination);
+		runs = reinterpret_cast<RunContainerRLEPair *>(destination);
+	}
+}
+
+void ContainerCompressionState::OverrideUncompressed(data_ptr_t &destination) {
+	append_function = AppendBitset;
+	destination = AlignValue<sizeof(idx_t)>(destination);
+	uncompressed = reinterpret_cast<validity_t *>(destination);
+}
+
+void ContainerCompressionState::Finalize() {
+	D_ASSERT(!finalized);
+	if (appended_count && last_is_null && run_idx < MAX_RUN_IDX) {
+		if (run_idx < COMPRESSED_RUN_THRESHOLD) {
+			auto &last_run = runs[run_idx];
+			// End the last run
+			last_run.length = (appended_count - last_run.start);
+		}
+		compressed_runs[(run_idx * 2) + 1] = static_cast<uint8_t>(appended_count % COMPRESSED_SEGMENT_SIZE);
+		if (appended_count != ROARING_CONTAINER_SIZE) {
+			run_counts[appended_count >> COMPRESSED_SEGMENT_SHIFT_AMOUNT]++;
+		}
+		run_idx++;
+	}
+	finalized = true;
+}
+
+ContainerMetadata ContainerCompressionState::GetResult() {
+	if (uncompressed) {
+		return ContainerMetadata::BitsetContainer(appended_count);
+	}
+	D_ASSERT(finalized);
+	return ContainerMetadata::CreateMetadata(appended_count, array_idx[NULLS], array_idx[NON_NULLS], run_idx);
+}
+
+void ContainerCompressionState::Reset() {
+	length = 0;
+
+	appended_count = 0;
+	null_count = 0;
+	run_idx = 0;
+	array_idx[NON_NULLS] = 0;
+	array_idx[NULLS] = 0;
+	finalized = false;
+	last_is_null = false;
+
+	// Reset the arrays + runs
+	arrays[NULLS] = base_arrays[NULLS];
+	arrays[NON_NULLS] = base_arrays[NON_NULLS];
+	runs = base_runs;
+
+	compressed_arrays[NULLS] = base_compressed_arrays[NULLS];
+	compressed_arrays[NON_NULLS] = base_compressed_arrays[NON_NULLS];
+	compressed_runs = base_compressed_runs;
+
+	array_counts[NULLS] = base_array_counts[NULLS];
+	array_counts[NON_NULLS] = base_array_counts[NON_NULLS];
+	run_counts = base_run_counts;
+
+	memset(array_counts[NULLS], 0, sizeof(uint8_t) * COMPRESSED_SEGMENT_COUNT);
+	memset(array_counts[NON_NULLS], 0, sizeof(uint8_t) * COMPRESSED_SEGMENT_COUNT);
+	memset(run_counts, 0, sizeof(uint8_t) * COMPRESSED_SEGMENT_COUNT);
+	uncompressed = nullptr;
+}
+
+//===--------------------------------------------------------------------===//
+// Compress
+//===--------------------------------------------------------------------===//
+RoaringCompressState::RoaringCompressState(ColumnDataCheckpointData &checkpoint_data,
+                                           unique_ptr<AnalyzeState> analyze_state_p)
+    : CompressionState(analyze_state_p->info), owned_analyze_state(std::move(analyze_state_p)),
+      analyze_state(owned_analyze_state->Cast<RoaringAnalyzeState>()), container_state(),
+      container_metadata(analyze_state.container_metadata), checkpoint_data(checkpoint_data),
+      function(checkpoint_data.GetCompressionFunction(CompressionType::COMPRESSION_ROARING)) {
+	CreateEmptySegment(checkpoint_data.GetRowGroup().start);
+	total_count = 0;
+	InitializeContainer();
+}
+
+idx_t RoaringCompressState::GetContainerIndex() {
+	idx_t index = total_count / ROARING_CONTAINER_SIZE;
+	return index;
+}
+
+idx_t RoaringCompressState::GetRemainingSpace() {
+	return static_cast<idx_t>(metadata_ptr - data_ptr);
+}
+
+bool RoaringCompressState::CanStore(idx_t container_size, const ContainerMetadata &metadata) {
+	idx_t required_space = 0;
+	if (metadata.IsUncompressed()) {
+		// Account for the alignment we might need for this container
+		required_space += (AlignValue<idx_t>(reinterpret_cast<idx_t>(data_ptr))) - reinterpret_cast<idx_t>(data_ptr);
+	}
+	required_space += metadata.GetDataSizeInBytes(container_size);
+
+	idx_t runs_count = metadata_collection.GetRunContainerCount();
+	idx_t arrays_count = metadata_collection.GetArrayAndBitsetContainerCount();
+#ifdef DEBUG
+	idx_t current_size = metadata_collection.GetMetadataSize(runs_count + arrays_count, runs_count, arrays_count);
+	(void)current_size;
+	D_ASSERT(required_space + current_size <= GetRemainingSpace());
+#endif
+	if (metadata.IsRun()) {
+		runs_count++;
+	} else {
+		arrays_count++;
+	}
+	idx_t metadata_size = metadata_collection.GetMetadataSize(runs_count + arrays_count, runs_count, arrays_count);
+	required_space += metadata_size;
+
+	if (required_space > GetRemainingSpace()) {
+		return false;
+	}
+	return true;
+}
+
+void RoaringCompressState::InitializeContainer() {
+	if (total_count == analyze_state.total_count) {
+		// No more containers left
+		return;
+	}
+	auto container_index = GetContainerIndex();
+	D_ASSERT(container_index < container_metadata.size());
+	auto metadata = container_metadata[container_index];
+
+	idx_t container_size = AlignValue<idx_t, ValidityMask::BITS_PER_VALUE>(
+	    MinValue<idx_t>(analyze_state.total_count - container_state.appended_count, ROARING_CONTAINER_SIZE));
+	if (!CanStore(container_size, metadata)) {
+		idx_t row_start = current_segment->start + current_segment->count;
+		FlushSegment();
+		CreateEmptySegment(row_start);
+	}
+
+	// Override the pointer to write directly into the block
+	if (metadata.IsUncompressed()) {
+		data_ptr = reinterpret_cast<data_ptr_t>(AlignValue<idx_t>(reinterpret_cast<idx_t>(data_ptr)));
+		FastMemset(data_ptr, ~0, sizeof(validity_t) * (container_size / ValidityMask::BITS_PER_VALUE));
+		container_state.OverrideUncompressed(data_ptr);
+	} else if (metadata.IsRun()) {
+		auto number_of_runs = metadata.NumberOfRuns();
+		container_state.OverrideRun(data_ptr, number_of_runs);
+	} else {
+		auto cardinality = metadata.Cardinality();
+		container_state.OverrideArray(data_ptr, metadata.IsInverted(), cardinality);
+	}
+	data_ptr += metadata.GetDataSizeInBytes(container_size);
+	metadata_collection.AddMetadata(metadata);
+}
+
+void RoaringCompressState::CreateEmptySegment(idx_t row_start) {
+	auto &db = checkpoint_data.GetDatabase();
+	auto &type = checkpoint_data.GetType();
+
+	auto compressed_segment =
+	    ColumnSegment::CreateTransientSegment(db, function, type, row_start, info.GetBlockSize(), info.GetBlockSize());
+	current_segment = std::move(compressed_segment);
+
+	auto &buffer_manager = BufferManager::GetBufferManager(db);
+	handle = buffer_manager.Pin(current_segment->block);
+	data_ptr = handle.Ptr();
+	data_ptr += sizeof(idx_t);
+	metadata_ptr = handle.Ptr() + info.GetBlockSize();
+}
+
+void RoaringCompressState::FlushSegment() {
+	auto &state = checkpoint_data.GetCheckpointState();
+	auto base_ptr = handle.Ptr();
+	// +======================================+
+	// |x|ddddddddddddddd||mmm|               |
+	// +======================================+
+
+	// x: metadata_offset (to the "right" of it)
+	// d: data of the containers
+	// m: metadata of the containers
+
+	// This is after 'x'
+	base_ptr += sizeof(idx_t);
+
+	// Size of the 'd' part
+	idx_t data_size = NumericCast<idx_t>(data_ptr - base_ptr);
+	data_size = AlignValue(data_size);
+
+	// Size of the 'm' part
+	idx_t metadata_size = metadata_collection.GetMetadataSizeForSegment();
+
+	if (current_segment->count.load() == 0) {
+		D_ASSERT(metadata_size == 0);
+		return;
+	}
+
+	idx_t serialized_metadata_size = metadata_collection.Serialize(data_ptr);
+	metadata_collection.FlushSegment();
+	(void)serialized_metadata_size;
+	D_ASSERT(metadata_size == serialized_metadata_size);
+	idx_t metadata_start = static_cast<idx_t>(data_ptr - base_ptr);
+	Store<idx_t>(metadata_start, handle.Ptr());
+	idx_t total_segment_size = sizeof(idx_t) + data_size + metadata_size;
+	state.FlushSegment(std::move(current_segment), std::move(handle), total_segment_size);
+}
+
+void RoaringCompressState::Finalize() {
+	FlushContainer();
+	FlushSegment();
+	current_segment.reset();
+}
+
+void RoaringCompressState::FlushContainer() {
+	if (container_state.length) {
+		container_state.Append(!container_state.last_bit_set, container_state.length);
+		container_state.length = 0;
+	}
+
+	if (!container_state.appended_count) {
+		return;
+	}
+	container_state.Finalize();
+#ifdef DEBUG
+	auto container_index = GetContainerIndex();
+	auto metadata = container_metadata[container_index];
+
+	idx_t container_size = container_state.appended_count;
+	if (!metadata.IsUncompressed()) {
+		unique_ptr<ContainerScanState> scan_state;
+		if (metadata.IsRun()) {
+			D_ASSERT(metadata.IsInverted());
+			auto number_of_runs = metadata.NumberOfRuns();
+			if (number_of_runs >= COMPRESSED_RUN_THRESHOLD) {
+				auto segments = container_state.run_counts;
+				auto data_ptr = container_state.compressed_runs;
+				scan_state = make_uniq<CompressedRunContainerScanState>(container_index, container_size, number_of_runs,
+				                                                        segments, data_ptr);
+			} else {
+				auto data_ptr = reinterpret_cast<data_ptr_t>(container_state.runs);
+				scan_state =
+				    make_uniq<RunContainerScanState>(container_index, container_size, number_of_runs, data_ptr);
+			}
+		} else {
+			auto cardinality = metadata.Cardinality();
+			if (cardinality >= COMPRESSED_ARRAY_THRESHOLD) {
+				if (metadata.IsInverted()) {
+					auto segments = reinterpret_cast<data_ptr_t>(container_state.array_counts[NULLS]);
+					auto data_ptr = reinterpret_cast<data_ptr_t>(container_state.compressed_arrays[NULLS]);
+					scan_state = make_uniq<CompressedArrayContainerScanState<NULLS>>(container_index, container_size,
+					                                                                 cardinality, segments, data_ptr);
+				} else {
+					auto segments = reinterpret_cast<data_ptr_t>(container_state.array_counts[NON_NULLS]);
+					auto data_ptr = reinterpret_cast<data_ptr_t>(container_state.compressed_arrays[NON_NULLS]);
+					scan_state = make_uniq<CompressedArrayContainerScanState<NON_NULLS>>(
+					    container_index, container_size, cardinality, segments, data_ptr);
+				}
+			} else {
+				if (metadata.IsInverted()) {
+					auto data_ptr = reinterpret_cast<data_ptr_t>(container_state.arrays[NULLS]);
+					scan_state = make_uniq<ArrayContainerScanState<NULLS>>(container_index, container_size, cardinality,
+					                                                       data_ptr);
+				} else {
+					auto data_ptr = reinterpret_cast<data_ptr_t>(container_state.arrays[NON_NULLS]);
+					scan_state = make_uniq<ArrayContainerScanState<NON_NULLS>>(container_index, container_size,
+					                                                           cardinality, data_ptr);
+				}
+			}
+		}
+		scan_state->Verify();
+	}
+
+#endif
+	total_count += container_state.appended_count;
+	bool has_nulls = container_state.null_count != 0;
+	bool has_non_nulls = container_state.null_count != container_state.appended_count;
+	if (has_nulls || container_state.uncompressed) {
+		current_segment->stats.statistics.SetHasNullFast();
+	}
+	if (has_non_nulls || container_state.uncompressed) {
+		current_segment->stats.statistics.SetHasNoNullFast();
+	}
+	current_segment->count += container_state.appended_count;
+	container_state.Reset();
+}
+
+void RoaringCompressState::NextContainer() {
+	FlushContainer();
+	InitializeContainer();
+}
+
+void RoaringCompressState::HandleByte(RoaringCompressState &state, uint8_t array_index) {
+	if (array_index == NumericLimits<uint8_t>::Maximum()) {
+		HandleAllValid(state, 8);
+	} else if (array_index == 0) {
+		HandleNoneValid(state, 8);
+	} else {
+		HandleRaggedByte(state, array_index, 8);
+	}
+}
+
+static inline void HandleBit(RoaringCompressState &state, bool bit_set) {
+	auto &container_state = state.container_state;
+	if (container_state.length && container_state.last_bit_set != bit_set) {
+		container_state.Append(!container_state.last_bit_set, container_state.length);
+		container_state.length = 0;
+	}
+	container_state.length += 1;
+	container_state.last_bit_set = bit_set;
+}
+
+void RoaringCompressState::HandleRaggedByte(RoaringCompressState &state, uint8_t array_index, idx_t relevant_bits) {
+	D_ASSERT(relevant_bits <= 8);
+	for (idx_t i = 0; i < relevant_bits; i++) {
+		const bool bit_set = array_index & (1 << i);
+		HandleBit(state, bit_set);
+	}
+}
+
+void RoaringCompressState::HandleAllValid(RoaringCompressState &state, idx_t amount) {
+	auto &container_state = state.container_state;
+	if (container_state.length && container_state.last_bit_set == false) {
+		container_state.Append(!container_state.last_bit_set, container_state.length);
+		container_state.length = 0;
+	}
+	container_state.length += amount;
+	container_state.last_bit_set = true;
+}
+
+void RoaringCompressState::HandleNoneValid(RoaringCompressState &state, idx_t amount) {
+	auto &container_state = state.container_state;
+	if (container_state.length && container_state.last_bit_set == true) {
+		container_state.Append(!container_state.last_bit_set, container_state.length);
+		container_state.length = 0;
+	}
+	container_state.length += amount;
+	container_state.last_bit_set = false;
+}
+
+idx_t RoaringCompressState::Count(RoaringCompressState &state) {
+	auto &container_state = state.container_state;
+	// How much is appended and waiting to be appended
+	return container_state.appended_count + container_state.length;
+}
+
+void RoaringCompressState::Flush(RoaringCompressState &state) {
+	state.NextContainer();
+}
+
+void RoaringCompressState::Compress(Vector &input, idx_t count) {
+	auto &self = *this;
+	RoaringStateAppender<RoaringCompressState>::AppendVector(self, input, count);
+}
+
+} // namespace roaring
+
+} // namespace duckdb