Repository: duckdb/duckdb
File: src/storage/compression/zstd.cpp (C++)

Code before:
	}

	bool UseVectorStateCache(idx_t vector_idx, idx_t internal_offset) {
		if (!current_vector) {
			// No vector loaded yet
			return false;
		}
		if (current_vector->metadata.vector_idx != vector_idx) {
			// Not the same vector
			return false;
		}
		if (current_vector->scanned_count != internal_offset) {
			// Not the same scan offset
			return false;
		}
		return true;
	}

	ZSTDVectorScanState &LoadVector(idx_t vector_idx, idx_t internal_offset) {
		if (UseVectorStateCache(vector_idx, internal_offset)) {
			return *current_vector;
		}
		current_vector = make_uniq<ZSTDVectorScanState>();
		current_vector->metadata = GetVectorMetadata(vector_idx);
		auto &metadata = current_vector->metadata;
		auto &scan_state = *current_vector;
		data_ptr_t handle_start;
		idx_t ptr_offset = 0;
		if (metadata.block_id == INVALID_BLOCK) {
			// Data lives on the segment's page
			handle_start = segment_handle.Ptr();
			ptr_offset += segment_block_offset;
		} else {
			// Data lives on an extra page, have to load the block first
			auto block = LoadPage(metadata.block_id);
			auto data_handle = buffer_manager.Pin(block);
			handle_start = data_handle.Ptr();
			scan_state.buffer_handles.push_back(std::move(data_handle));
		}

		ptr_offset += metadata.block_offset;
		ptr_offset = AlignValue<idx_t, sizeof(string_length_t)>(ptr_offset);
		scan_state.current_buffer_ptr = handle_start + ptr_offset;

		auto vector_size = metadata.count;

		auto string_lengths_size = (sizeof(string_length_t) * vector_size);
		scan_state.string_lengths = reinterpret_cast<string_length_t *>(scan_state.current_buffer_ptr);
		scan_state.current_buffer_ptr += string_lengths_size;

		// Update the in_buffer to point to the start of the compressed data frame

Review comment (bug):
Perhaps this should take the metadata.block_id into account?

I'll need to investigate the CI failure

Code after:
			// Data lives on the segment's page
			handle_start = segment_handle.Ptr();
			ptr_offset += segment_block_offset;
		} else {
			// Data lives on an extra page, have to load the block first
			auto block = LoadPage(metadata.block_id);
			auto data_handle = buffer_manager.Pin(block);
			handle_start = data_handle.Ptr();
			scan_state.buffer_handles.push_back(std::move(data_handle));
		}

		ptr_offset += metadata.block_offset;
		ptr_offset = AlignValue<idx_t, sizeof(string_length_t)>(ptr_offset);
		scan_state.current_buffer_ptr = handle_start + ptr_offset;

		auto vector_size = metadata.count;

		auto string_lengths_size = (sizeof(string_length_t) * vector_size);
		scan_state.string_lengths = reinterpret_cast<string_length_t *>(scan_state.current_buffer_ptr);
		scan_state.current_buffer_ptr += string_lengths_size;

		// Update the in_buffer to point to the start of the compressed data frame
		idx_t current_offset = UnsafeNumericCast<idx_t>(scan_state.current_buffer_ptr - handle_start);
		scan_state.in_buffer.src = scan_state.current_buffer_ptr;
		scan_state.in_buffer.pos = 0;
		if (scan_state.metadata.block_offset + string_lengths_size + scan_state.metadata.compressed_size >
		    (segment.SegmentSize() - sizeof(block_id_t))) {
			//! We know that the compressed size is too big to fit on the current page
			scan_state.in_buffer.size =
			    MinValue(metadata.compressed_size, block_manager.GetBlockSize() - sizeof(block_id_t) - current_offset);
		} else {
			scan_state.in_buffer.size =
			    MinValue(metadata.compressed_size, block_manager.GetBlockSize() - current_offset);
		}

		// Initialize the context for streaming decompression
		duckdb_zstd::ZSTD_DCtx_reset(decompression_context, duckdb_zstd::ZSTD_reset_session_only);
		duckdb_zstd::ZSTD_DCtx_refDDict(decompression_context, nullptr);

		if (internal_offset) {
			Skip(scan_state, internal_offset);
		}
		return scan_state;
	}

	void LoadNextPageForVector(ZSTDVectorScanState &scan_state) {
		if (scan_state.in_buffer.pos != scan_state.in_buffer.size) {
			throw InternalException(
			    "(ZSTDScanState::LoadNextPageForVector) Trying to load the next page before consuming the current one");
		}
		// Read the next block id from the end of the page