duckdb 1.1.2-dev6.0 → 1.1.4-dev11.0

package/src/duckdb/src/include/duckdb/storage/compression/roaring/roaring.hpp

@@ -0,0 +1,618 @@
+//===----------------------------------------------------------------------===//
+//                         DuckDB
+//
+// duckdb/storage/compression/roaring/roaring.hpp
+//
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "duckdb/common/typedefs.hpp"
+#include "duckdb/common/helper.hpp"
+#include "duckdb/common/types/validity_mask.hpp"
+#include "duckdb/common/types/vector.hpp"
+#include "duckdb/function/compression_function.hpp"
+#include "duckdb/storage/table/scan_state.hpp"
+
+namespace duckdb {
+namespace roaring {
+
+//! Used for compressed runs/arrays
+static constexpr uint16_t COMPRESSED_SEGMENT_SIZE = 256;
+//! compresed segment size is 256, instead of division we can make use of shifting
+static constexpr uint16_t COMPRESSED_SEGMENT_SHIFT_AMOUNT = 8;
+//! The amount of values that are encoded per container
+static constexpr idx_t ROARING_CONTAINER_SIZE = 2048;
+static constexpr bool NULLS = true;
+static constexpr bool NON_NULLS = false;
+static constexpr uint16_t UNCOMPRESSED_SIZE = (ROARING_CONTAINER_SIZE / sizeof(validity_t));
+static constexpr uint16_t COMPRESSED_SEGMENT_COUNT = (ROARING_CONTAINER_SIZE / COMPRESSED_SEGMENT_SIZE);
+
+static constexpr uint16_t MAX_RUN_IDX = (UNCOMPRESSED_SIZE - COMPRESSED_SEGMENT_COUNT) / (sizeof(uint8_t) * 2);
+static constexpr uint16_t MAX_ARRAY_IDX = (UNCOMPRESSED_SIZE - COMPRESSED_SEGMENT_COUNT) / (sizeof(uint8_t) * 1);
+//! The value used to indicate that a container is a bitset container
+static constexpr uint16_t BITSET_CONTAINER_SENTINEL_VALUE = MAX_ARRAY_IDX + 1;
+static constexpr uint16_t COMPRESSED_ARRAY_THRESHOLD = 8;
+static constexpr uint16_t COMPRESSED_RUN_THRESHOLD = 4;
+
+static constexpr uint16_t CONTAINER_TYPE_BITWIDTH = 2;
+static constexpr uint16_t RUN_CONTAINER_SIZE_BITWIDTH = 7;
+static constexpr uint16_t ARRAY_CONTAINER_SIZE_BITWIDTH = 8;
+
+//! This can be increased, but requires additional changes beyond just changing the constant
+static_assert(MAX_RUN_IDX < NumericLimits<uint8_t>::Maximum(), "The run size is encoded in a maximum of 8 bits");
+
+//! This can be increased, but requires additional changes beyond just changing the constant
+static_assert(BITSET_CONTAINER_SENTINEL_VALUE < NumericLimits<uint8_t>::Maximum(),
+              "The array/bitset size is encoded in a maximum of 8 bits");
+
+static_assert(ROARING_CONTAINER_SIZE % COMPRESSED_SEGMENT_SIZE == 0,
+              "The (maximum) container size has to be cleanly divisable by the segment size");
+
+static_assert((1 << RUN_CONTAINER_SIZE_BITWIDTH) - 1 >= MAX_RUN_IDX,
+              "The bitwidth used to store the size of a run container has to be big enough to store the maximum size");
+static_assert(
+    (1 << ARRAY_CONTAINER_SIZE_BITWIDTH) - 1 >= MAX_ARRAY_IDX + 1,
+    "The bitwidth used to store the size of an array/bitset container has to be big enough to store the maximum size");
+
+void SetInvalidRange(ValidityMask &result, idx_t start, idx_t end);
+
+struct RunContainerRLEPair {
+	uint16_t start;
+	uint16_t length;
+};
+
+enum class ContainerType : uint8_t { RUN_CONTAINER, ARRAY_CONTAINER, BITSET_CONTAINER };
+
+struct ContainerMetadata {
+public:
+	bool operator==(const ContainerMetadata &other) const {
+		if (container_type != other.container_type) {
+			return false;
+		}
+		if (count != other.count) {
+			return false;
+		}
+		if (nulls != other.nulls) {
+			return false;
+		}
+		return true;
+	}
+
+	static ContainerMetadata CreateMetadata(uint16_t count, uint16_t array_null, uint16_t array_non_null,
+	                                        uint16_t runs);
+
+	static ContainerMetadata RunContainer(uint16_t runs) {
+		auto res = ContainerMetadata();
+		res.container_type = ContainerType::RUN_CONTAINER;
+		res.nulls = true;
+		res.count = runs;
+		return res;
+	}
+
+	static ContainerMetadata ArrayContainer(uint16_t array_size, bool nulls) {
+		auto res = ContainerMetadata();
+		res.container_type = ContainerType::ARRAY_CONTAINER;
+		res.nulls = nulls;
+		res.count = array_size;
+		return res;
+	}
+
+	static ContainerMetadata BitsetContainer(uint16_t container_size) {
+		auto res = ContainerMetadata();
+		res.container_type = ContainerType::BITSET_CONTAINER;
+		res.nulls = true;
+		res.count = container_size;
+		return res;
+	}
+
+public:
+	idx_t GetDataSizeInBytes(idx_t container_size) const;
+	bool IsUncompressed() const {
+		return container_type == ContainerType::BITSET_CONTAINER;
+	}
+	bool IsRun() const {
+		return container_type == ContainerType::RUN_CONTAINER;
+	}
+	bool IsInverted() const {
+		return nulls;
+	}
+	bool IsArray() const {
+		return container_type == ContainerType::ARRAY_CONTAINER;
+	}
+	idx_t NumberOfRuns() const {
+		D_ASSERT(IsRun());
+		return count;
+	}
+	idx_t Cardinality() const {
+		D_ASSERT(IsArray());
+		return count;
+	}
+
+public:
+	ContainerType container_type;
+	//! Whether nulls are being encoded or non-nulls
+	bool nulls;
+	//! The amount (meaning depends on container_type)
+	uint16_t count;
+
+private:
+	ContainerMetadata() {
+	}
+};
+
+struct ContainerMetadataCollection {
+	static constexpr uint8_t IS_RUN_FLAG = 1 << 1;
+	static constexpr uint8_t IS_INVERTED_FLAG = 1 << 0;
+
+public:
+	ContainerMetadataCollection();
+
+public:
+	void AddMetadata(ContainerMetadata metadata);
+	idx_t GetMetadataSizeForSegment() const;
+	idx_t GetMetadataSize(idx_t container_count, idx_t run_containers, idx_t array_containers) const;
+	idx_t GetRunContainerCount() const;
+	idx_t GetArrayAndBitsetContainerCount() const;
+	void FlushSegment();
+	void Reset();
+	// Write the metadata for the current segment
+	idx_t Serialize(data_ptr_t dest) const;
+	void Deserialize(data_ptr_t src, idx_t container_count);
+
+private:
+	void AddBitsetContainer();
+	void AddArrayContainer(idx_t amount, bool is_inverted);
+	void AddRunContainer(idx_t amount, bool is_inverted);
+	void AddContainerType(bool is_run, bool is_inverted);
+
+public:
+	//! Encode for each container in the lower 2 bits if the container 'is_run' and 'is_inverted'
+	vector<uint8_t> container_type;
+	//! Encode for each run container the length
+	vector<uint8_t> number_of_runs;
+	//! Encode for each array/bitset container the length
+	vector<uint8_t> cardinality;
+
+	idx_t count_in_segment = 0;
+	idx_t runs_in_segment = 0;
+	idx_t arrays_in_segment = 0;
+};
+
+struct ContainerMetadataCollectionScanner {
+public:
+	explicit ContainerMetadataCollectionScanner(ContainerMetadataCollection &collection);
+
+public:
+	ContainerMetadata GetNext();
+
+public:
+	const ContainerMetadataCollection &collection;
+	idx_t array_idx = 0;
+	idx_t run_idx = 0;
+	idx_t idx = 0;
+};
+
+struct BitmaskTableEntry {
+	uint8_t first_bit_set : 1;
+	uint8_t last_bit_set : 1;
+	uint8_t valid_count : 6;
+	uint8_t run_count;
+};
+
+//===--------------------------------------------------------------------===//
+// Analyze
+//===--------------------------------------------------------------------===//
+struct RoaringAnalyzeState : public AnalyzeState {
+public:
+	explicit RoaringAnalyzeState(const CompressionInfo &info);
+
+public:
+	// RoaringStateAppender interface:
+
+	static void HandleByte(RoaringAnalyzeState &state, uint8_t array_index);
+	static void HandleRaggedByte(RoaringAnalyzeState &state, uint8_t array_index, idx_t relevant_bits);
+	static void HandleAllValid(RoaringAnalyzeState &state, idx_t amount);
+	static void HandleNoneValid(RoaringAnalyzeState &state, idx_t amount);
+	static idx_t Count(RoaringAnalyzeState &state);
+	static void Flush(RoaringAnalyzeState &state);
+
+public:
+	ContainerMetadata GetResult();
+	bool HasEnoughSpaceInSegment(idx_t required_space);
+	void FlushSegment();
+	void FlushContainer();
+	void Analyze(Vector &input, idx_t count);
+
+public:
+	unsafe_unique_array<BitmaskTableEntry> bitmask_table;
+
+	//! Analyze phase
+
+	//! The amount of set bits found
+	uint16_t one_count = 0;
+	//! The amount of unset bits found
+	uint16_t zero_count = 0;
+	//! The amount of runs (of 0's) so far
+	uint16_t run_count = 0;
+	//! Whether the last bit was set or not
+	bool last_bit_set;
+	//! The total amount of bits covered (one_count + zero_count)
+	uint16_t count = 0;
+
+	//! Flushed analyze data
+
+	//! The space used by the current segment
+	idx_t data_size = 0;
+	idx_t metadata_size = 0;
+
+	//! The total amount of segments to write
+	idx_t segment_count = 0;
+	//! The amount of values in the current segment;
+	idx_t current_count = 0;
+	//! The total amount of data to serialize
+	idx_t total_count = 0;
+
+	//! The total amount of bytes used to compress the whole segment
+	idx_t total_size = 0;
+	//! The container metadata, determining the type of each container to use during compression
+	ContainerMetadataCollection metadata_collection;
+	vector<ContainerMetadata> container_metadata;
+};
+
+//===--------------------------------------------------------------------===//
+// Compress
+//===--------------------------------------------------------------------===//
+
+struct ContainerCompressionState {
+	using append_func_t = void (*)(ContainerCompressionState &, bool, uint16_t);
+
+public:
+	ContainerCompressionState();
+
+public:
+	void Append(bool null, uint16_t amount = 1);
+	void OverrideArray(data_ptr_t &destination, bool nulls, idx_t count);
+	void OverrideRun(data_ptr_t &destination, idx_t count);
+	void OverrideUncompressed(data_ptr_t &destination);
+	void Finalize();
+	ContainerMetadata GetResult();
+	void Reset();
+
+public:
+	//! Buffered append state (we don't want to append every bit separately)
+	uint16_t length = 0;
+	bool last_bit_set;
+
+	//! Total amount of values covered by the container
+	uint16_t appended_count = 0;
+	//! How many of the total are null
+	uint16_t null_count = 0;
+	bool last_is_null = false;
+
+	RunContainerRLEPair *runs;
+	uint8_t *compressed_runs;
+	uint8_t *compressed_arrays[2];
+	uint16_t *arrays[2];
+
+	//! The runs (for sequential nulls)
+	RunContainerRLEPair base_runs[COMPRESSED_RUN_THRESHOLD];
+	//! The indices (for nulls | non-nulls)
+	uint16_t base_arrays[2][COMPRESSED_ARRAY_THRESHOLD];
+
+	uint16_t run_idx;
+	uint16_t array_idx[2];
+
+	uint8_t *array_counts[2];
+	uint8_t *run_counts;
+
+	uint8_t base_compressed_arrays[2][MAX_ARRAY_IDX];
+	uint8_t base_compressed_runs[MAX_RUN_IDX * 2];
+	uint8_t base_array_counts[2][COMPRESSED_SEGMENT_COUNT];
+	uint8_t base_run_counts[COMPRESSED_SEGMENT_COUNT];
+
+	validity_t *uncompressed = nullptr;
+	//! Whether the state has been finalized
+	bool finalized = false;
+	//! The function called to append to the state
+	append_func_t append_function;
+};
+
+struct RoaringCompressState : public CompressionState {
+public:
+	explicit RoaringCompressState(ColumnDataCheckpointData &checkpoint_data, unique_ptr<AnalyzeState> analyze_state_p);
+
+public:
+	//! RoaringStateAppender interface
+	static void HandleByte(RoaringCompressState &state, uint8_t array_index);
+	static void HandleRaggedByte(RoaringCompressState &state, uint8_t array_index, idx_t relevant_bits);
+	static void HandleAllValid(RoaringCompressState &state, idx_t amount);
+	static void HandleNoneValid(RoaringCompressState &state, idx_t amount);
+	static idx_t Count(RoaringCompressState &state);
+	static void Flush(RoaringCompressState &state);
+
+public:
+	idx_t GetContainerIndex();
+	idx_t GetRemainingSpace();
+	bool CanStore(idx_t container_size, const ContainerMetadata &metadata);
+	void InitializeContainer();
+	void CreateEmptySegment(idx_t row_start);
+	void FlushSegment();
+	void Finalize();
+	void FlushContainer();
+	void NextContainer();
+	void Compress(Vector &input, idx_t count);
+
+public:
+	unique_ptr<AnalyzeState> owned_analyze_state;
+	RoaringAnalyzeState &analyze_state;
+
+	ContainerCompressionState container_state;
+	ContainerMetadataCollection metadata_collection;
+	vector<ContainerMetadata> &container_metadata;
+
+	ColumnDataCheckpointData &checkpoint_data;
+	CompressionFunction &function;
+	unique_ptr<ColumnSegment> current_segment;
+	BufferHandle handle;
+
+	// Ptr to next free spot in segment;
+	data_ptr_t data_ptr;
+	// Ptr to next free spot for storing
+	data_ptr_t metadata_ptr;
+	//! The amount of values already compressed
+	idx_t total_count = 0;
+};
+
+//===--------------------------------------------------------------------===//
+// Scan
+//===--------------------------------------------------------------------===//
+
+struct ContainerScanState {
+public:
+	ContainerScanState(idx_t container_index_p, idx_t container_size)
+	    : container_index(container_index_p), container_size(container_size) {
+	}
+	virtual ~ContainerScanState() {
+	}
+
+public:
+	virtual void ScanPartial(Vector &result, idx_t result_offset, idx_t to_scan) = 0;
+	virtual void Skip(idx_t count) = 0;
+	virtual void Verify() const = 0;
+
+public:
+	//! The index of the container
+	idx_t container_index;
+	//! The size of the container (how many values does it hold)
+	idx_t container_size;
+	//! How much of the container is already consumed
+	idx_t scanned_count = 0;
+};
+
+struct ContainerSegmentScan {
+public:
+	explicit ContainerSegmentScan(data_ptr_t data);
+	ContainerSegmentScan(const ContainerSegmentScan &other) = delete;
+	ContainerSegmentScan(ContainerSegmentScan &&other) = delete;
+	ContainerSegmentScan &operator=(const ContainerSegmentScan &other) = delete;
+	ContainerSegmentScan &operator=(ContainerSegmentScan &&other) = delete;
+
+public:
+	// Returns the base of the current segment, forwarding the index if the segment is depleted of values
+	uint16_t operator++(int);
+
+private:
+	//! The COMPRESSED_SEGMENT_COUNT unsigned bytes indicating for each segment (256 bytes) of the container how many
+	//! values are in the segment
+	uint8_t *segments;
+	uint8_t index;
+	uint8_t count;
+};
+
+//! RUN Container
+
+struct RunContainerScanState : public ContainerScanState {
+public:
+	RunContainerScanState(idx_t container_index, idx_t container_size, idx_t count, data_ptr_t data_p);
+
+public:
+	void ScanPartial(Vector &result, idx_t result_offset, idx_t to_scan) override;
+	void Skip(idx_t to_skip) override;
+	void Verify() const override;
+
+protected:
+	virtual void LoadNextRun();
+
+protected:
+	RunContainerRLEPair run;
+	bool finished = false;
+	idx_t run_index = 0;
+	idx_t count;
+	data_ptr_t data;
+};
+
+struct CompressedRunContainerScanState : public RunContainerScanState {
+public:
+	CompressedRunContainerScanState(idx_t container_index, idx_t container_size, idx_t count, data_ptr_t segments,
+	                                data_ptr_t data);
+
+protected:
+	void LoadNextRun() override;
+	void Verify() const override;
+
+private:
+	data_ptr_t segments;
+	ContainerSegmentScan segment;
+};
+
+//! ARRAY Container
+
+template <bool INVERTED>
+struct ArrayContainerScanState : public ContainerScanState {
+public:
+	ArrayContainerScanState(idx_t container_index, idx_t container_size, idx_t count, data_ptr_t data_p)
+	    : ContainerScanState(container_index, container_size), data(data_p), count(count) {
+	}
+
+public:
+	virtual void LoadNextValue() {
+		if (array_index >= count) {
+			finished = true;
+			return;
+		}
+		value = reinterpret_cast<uint16_t *>(data)[array_index];
+		array_index++;
+	}
+
+public:
+	void ScanPartial(Vector &result, idx_t result_offset, idx_t to_scan) override {
+		auto &result_mask = FlatVector::Validity(result);
+
+		// This method assumes that the validity mask starts off as having all bits set for the entries that are being
+		// scanned.
+
+		if (!INVERTED) {
+			// If we are mapping valid entries, that means the majority of the bits are invalid
+			// so we set everything to invalid and only flip the bits that are present in the array
+			SetInvalidRange(result_mask, result_offset, result_offset + to_scan);
+		}
+
+		if (!array_index) {
+			LoadNextValue();
+		}
+		// At least one of the entries to scan is set
+		while (!finished) {
+			if (value >= scanned_count + to_scan) {
+				break;
+			}
+			if (value < scanned_count) {
+				LoadNextValue();
+				continue;
+			}
+			auto index = value - scanned_count;
+			if (INVERTED) {
+				result_mask.SetInvalid(result_offset + index);
+			} else {
+				result_mask.SetValid(result_offset + index);
+			}
+			LoadNextValue();
+		}
+		scanned_count += to_scan;
+	}
+
+	void Skip(idx_t to_skip) override {
+		idx_t end = scanned_count + to_skip;
+		if (!array_index) {
+			LoadNextValue();
+		}
+		while (!finished && value < end) {
+			LoadNextValue();
+		}
+		// In case array_index has already reached count
+		scanned_count = end;
+	}
+
+	void Verify() const override {
+#ifdef DEBUG
+		uint16_t index = 0;
+		auto array = reinterpret_cast<uint16_t *>(data);
+		for (uint16_t i = 0; i < count; i++) {
+			D_ASSERT(!i || array[i] > index);
+			index = array[i];
+		}
+#endif
+	}
+
+protected:
+	uint16_t value;
+	data_ptr_t data;
+	bool finished = false;
+	const idx_t count;
+	idx_t array_index = 0;
+};
+
+template <bool INVERTED>
+struct CompressedArrayContainerScanState : public ArrayContainerScanState<INVERTED> {
+public:
+	CompressedArrayContainerScanState(idx_t container_index, idx_t container_size, idx_t count, data_ptr_t segments,
+	                                  data_ptr_t data)
+	    : ArrayContainerScanState<INVERTED>(container_index, container_size, count, data), segments(segments),
+	      segment(segments) {
+		D_ASSERT(count >= COMPRESSED_ARRAY_THRESHOLD);
+	}
+
+public:
+	void LoadNextValue() override {
+		if (this->array_index >= this->count) {
+			this->finished = true;
+			return;
+		}
+		this->value = segment++;
+		this->value += reinterpret_cast<uint8_t *>(this->data)[this->array_index];
+		this->array_index++;
+	}
+	void Verify() const override {
+#ifdef DEBUG
+		uint16_t index = 0;
+		ContainerSegmentScan verify_segment(segments);
+		for (uint16_t i = 0; i < this->count; i++) {
+			// Get the value
+			uint16_t new_index = verify_segment++;
+			new_index += reinterpret_cast<uint8_t *>(this->data)[i];
+
+			D_ASSERT(!i || new_index > index);
+			index = new_index;
+		}
+#endif
+	}
+
+private:
+	data_ptr_t segments;
+	ContainerSegmentScan segment;
+};
+
+//! BITSET Container
+
+struct BitsetContainerScanState : public ContainerScanState {
+public:
+	BitsetContainerScanState(idx_t container_index, idx_t count, validity_t *bitset);
+
+public:
+	void ScanPartial(Vector &result, idx_t result_offset, idx_t to_scan) override;
+	void Skip(idx_t to_skip) override;
+	void Verify() const override;
+
+public:
+	validity_t *bitset;
+};
+
+struct RoaringScanState : public SegmentScanState {
+public:
+	explicit RoaringScanState(ColumnSegment &segment);
+
+public:
+	idx_t SkipVector(const ContainerMetadata &metadata);
+	bool UseContainerStateCache(idx_t container_index, idx_t internal_offset);
+	ContainerMetadata GetContainerMetadata(idx_t container_index);
+	data_ptr_t GetStartOfContainerData(idx_t container_index);
+	ContainerScanState &LoadContainer(idx_t container_index, idx_t internal_offset);
+	void ScanInternal(ContainerScanState &scan_state, idx_t to_scan, Vector &result, idx_t offset);
+	idx_t GetContainerIndex(idx_t start_index, idx_t &offset);
+	void ScanPartial(idx_t start_idx, Vector &result, idx_t offset, idx_t count);
+	void Skip(ContainerScanState &scan_state, idx_t skip_count);
+
+public:
+	BufferHandle handle;
+	ColumnSegment &segment;
+	unique_ptr<ContainerScanState> current_container;
+	data_ptr_t data_ptr;
+	ContainerMetadataCollection metadata_collection;
+	vector<ContainerMetadata> container_metadata;
+	vector<idx_t> data_start_position;
+};
+
+} // namespace roaring
+
+} // namespace duckdb