duckdb 0.7.2-dev1901.0 → 0.7.2-dev2233.0

package/src/duckdb/src/common/types/row/tuple_data_collection.cpp

@@ -0,0 +1,511 @@
+#include "duckdb/common/types/row/tuple_data_collection.hpp"
+
+#include "duckdb/common/fast_mem.hpp"
+#include "duckdb/common/printer.hpp"
+#include "duckdb/common/row_operations/row_operations.hpp"
+#include "duckdb/common/types/row/tuple_data_allocator.hpp"
+
+#include <algorithm>
+
+namespace duckdb {
+
+using ValidityBytes = TupleDataLayout::ValidityBytes;
+
+TupleDataCollection::TupleDataCollection(BufferManager &buffer_manager, const TupleDataLayout &layout_p)
+    : layout(layout_p.Copy()), allocator(make_shared<TupleDataAllocator>(buffer_manager, layout)) {
+	Initialize();
+}
+
+TupleDataCollection::TupleDataCollection(shared_ptr<TupleDataAllocator> allocator)
+    : layout(allocator->GetLayout().Copy()), allocator(std::move(allocator)) {
+	Initialize();
+}
+
+TupleDataCollection::~TupleDataCollection() {
+}
+
+void TupleDataCollection::Initialize() {
+	D_ASSERT(!layout.GetTypes().empty());
+	this->count = 0;
+	scatter_functions.reserve(layout.ColumnCount());
+	gather_functions.reserve(layout.ColumnCount());
+	for (idx_t col_idx = 0; col_idx < layout.ColumnCount(); col_idx++) {
+		auto &type = layout.GetTypes()[col_idx];
+		scatter_functions.emplace_back(GetScatterFunction(type));
+		gather_functions.emplace_back(GetGatherFunction(type));
+	}
+}
+
+void TupleDataCollection::GetAllColumnIDs(vector<column_t> &column_ids) {
+	column_ids.reserve(layout.ColumnCount());
+	for (idx_t col_idx = 0; col_idx < layout.ColumnCount(); col_idx++) {
+		column_ids.emplace_back(col_idx);
+	}
+}
+
+const TupleDataLayout &TupleDataCollection::GetLayout() const {
+	return layout;
+}
+
+const idx_t &TupleDataCollection::Count() const {
+	return count;
+}
+
+idx_t TupleDataCollection::ChunkCount() const {
+	idx_t total_chunk_count = 0;
+	for (const auto &segment : segments) {
+		total_chunk_count += segment.ChunkCount();
+	}
+	return total_chunk_count;
+}
+
+idx_t TupleDataCollection::SizeInBytes() const {
+	idx_t total_size = 0;
+	for (const auto &segment : segments) {
+		total_size += segment.SizeInBytes();
+	}
+	return total_size;
+}
+
+void TupleDataCollection::GetBlockPointers(vector<data_ptr_t> &block_pointers) const {
+	D_ASSERT(segments.size() == 1);
+	const auto &segment = segments[0];
+	const auto block_count = segment.allocator->RowBlockCount();
+	D_ASSERT(segment.pinned_row_handles.size() == block_count);
+	block_pointers.resize(block_count);
+	for (idx_t block_idx = 0; block_idx < block_count; block_idx++) {
+		block_pointers[block_idx] = segment.pinned_row_handles[block_idx].Ptr();
+	}
+}
+
+void TupleDataCollection::Unpin() {
+	for (auto &segment : segments) {
+		segment.Unpin();
+	}
+}
+
+void VerifyAppendColumns(const TupleDataLayout &layout, const vector<column_t> &column_ids) {
+#ifdef DEBUG
+	for (idx_t col_idx = 0; col_idx < layout.ColumnCount(); col_idx++) {
+		if (std::find(column_ids.begin(), column_ids.end(), col_idx) != column_ids.end()) {
+			continue;
+		}
+		// This column will not be appended in the first go - verify that it is fixed-size - we cannot resize heap after
+		const auto physical_type = layout.GetTypes()[col_idx].InternalType();
+		D_ASSERT(physical_type != PhysicalType::VARCHAR && physical_type != PhysicalType::LIST);
+		if (physical_type == PhysicalType::STRUCT) {
+			const auto &struct_layout = layout.GetStructLayout(col_idx);
+			vector<column_t> struct_column_ids;
+			struct_column_ids.reserve(struct_layout.ColumnCount());
+			for (idx_t struct_col_idx = 0; struct_col_idx < struct_layout.ColumnCount(); struct_col_idx++) {
+				struct_column_ids.emplace_back(struct_col_idx);
+			}
+			VerifyAppendColumns(struct_layout, struct_column_ids);
+		}
+	}
+#endif
+}
+
+void TupleDataCollection::InitializeAppend(TupleDataAppendState &append_state, TupleDataPinProperties properties) {
+	vector<column_t> column_ids;
+	GetAllColumnIDs(column_ids);
+	InitializeAppend(append_state, std::move(column_ids), properties);
+}
+
+void TupleDataCollection::InitializeAppend(TupleDataAppendState &append_state, vector<column_t> column_ids,
+                                           TupleDataPinProperties properties) {
+	VerifyAppendColumns(layout, column_ids);
+	InitializeAppend(append_state.pin_state, properties);
+	InitializeAppend(append_state.chunk_state, std::move(column_ids));
+}
+
+void TupleDataCollection::InitializeAppend(TupleDataPinState &pin_state, TupleDataPinProperties properties) {
+	pin_state.properties = properties;
+	if (segments.empty()) {
+		segments.emplace_back(allocator);
+	}
+}
+
+static void InitializeVectorFormat(vector<TupleDataVectorFormat> &vector_data, const vector<LogicalType> &types) {
+	vector_data.resize(types.size());
+	for (idx_t col_idx = 0; col_idx < types.size(); col_idx++) {
+		const auto &type = types[col_idx];
+		switch (type.InternalType()) {
+		case PhysicalType::STRUCT: {
+			const auto &child_list = StructType::GetChildTypes(type);
+			vector<LogicalType> child_types;
+			child_types.reserve(child_list.size());
+			for (const auto &child_entry : child_list) {
+				child_types.emplace_back(child_entry.second);
+			}
+			InitializeVectorFormat(vector_data[col_idx].child_formats, child_types);
+			break;
+		}
+		case PhysicalType::LIST:
+			InitializeVectorFormat(vector_data[col_idx].child_formats, {ListType::GetChildType(type)});
+			break;
+		default:
+			break;
+		}
+	}
+}
+
+void TupleDataCollection::InitializeAppend(TupleDataChunkState &chunk_state, vector<column_t> column_ids) {
+	if (column_ids.empty()) {
+		GetAllColumnIDs(column_ids);
+	}
+	InitializeVectorFormat(chunk_state.vector_data, layout.GetTypes());
+	chunk_state.column_ids = std::move(column_ids);
+}
+
+void TupleDataCollection::Append(DataChunk &new_chunk, const SelectionVector &append_sel, idx_t append_count) {
+	TupleDataAppendState append_state;
+	InitializeAppend(append_state);
+	Append(append_state, new_chunk, append_sel, append_count);
+}
+
+void TupleDataCollection::Append(DataChunk &new_chunk, vector<column_t> column_ids, const SelectionVector &append_sel,
+                                 const idx_t append_count) {
+	TupleDataAppendState append_state;
+	InitializeAppend(append_state, std::move(column_ids));
+	Append(append_state, new_chunk, append_sel, append_count);
+}
+
+void TupleDataCollection::Append(TupleDataAppendState &append_state, DataChunk &new_chunk,
+                                 const SelectionVector &append_sel, const idx_t append_count) {
+	Append(append_state.pin_state, append_state.chunk_state, new_chunk, append_sel, append_count);
+}
+
+void TupleDataCollection::Append(TupleDataPinState &pin_state, TupleDataChunkState &chunk_state, DataChunk &new_chunk,
+                                 const SelectionVector &append_sel, const idx_t append_count) {
+	TupleDataCollection::ToUnifiedFormat(chunk_state, new_chunk);
+	AppendUnified(pin_state, chunk_state, new_chunk, append_sel, append_count);
+}
+
+void TupleDataCollection::AppendUnified(TupleDataPinState &pin_state, TupleDataChunkState &chunk_state,
+                                        DataChunk &new_chunk, const SelectionVector &append_sel,
+                                        const idx_t append_count) {
+	const idx_t actual_append_count = append_count == DConstants::INVALID_INDEX ? new_chunk.size() : append_count;
+	if (actual_append_count == 0) {
+		return;
+	}
+
+	if (!layout.AllConstant()) {
+		TupleDataCollection::ComputeHeapSizes(chunk_state, new_chunk, append_sel, actual_append_count);
+	}
+
+	Build(pin_state, chunk_state, 0, actual_append_count);
+
+#ifdef DEBUG
+	Vector heap_locations_copy(LogicalType::POINTER);
+	if (!layout.AllConstant()) {
+		VectorOperations::Copy(chunk_state.heap_locations, heap_locations_copy, actual_append_count, 0, 0);
+	}
+#endif
+
+	Scatter(chunk_state, new_chunk, append_sel, actual_append_count);
+
+#ifdef DEBUG
+	// Verify that the size of the data written to the heap is the same as the size we computed it would be
+	if (!layout.AllConstant()) {
+		const auto original_heap_locations = FlatVector::GetData<data_ptr_t>(heap_locations_copy);
+		const auto heap_sizes = FlatVector::GetData<idx_t>(chunk_state.heap_sizes);
+		const auto offset_heap_locations = FlatVector::GetData<data_ptr_t>(chunk_state.heap_locations);
+		for (idx_t i = 0; i < actual_append_count; i++) {
+			D_ASSERT(offset_heap_locations[i] == original_heap_locations[i] + heap_sizes[i]);
+		}
+	}
+#endif
+}
+
+static inline void ToUnifiedFormatInternal(TupleDataVectorFormat &format, Vector &vector, const idx_t count) {
+	vector.ToUnifiedFormat(count, format.data);
+	switch (vector.GetType().InternalType()) {
+	case PhysicalType::STRUCT: {
+		auto &entries = StructVector::GetEntries(vector);
+		D_ASSERT(format.child_formats.size() == entries.size());
+		for (idx_t struct_col_idx = 0; struct_col_idx < entries.size(); struct_col_idx++) {
+			ToUnifiedFormatInternal(format.child_formats[struct_col_idx], *entries[struct_col_idx], count);
+		}
+		break;
+	}
+	case PhysicalType::LIST:
+		D_ASSERT(format.child_formats.size() == 1);
+		ToUnifiedFormatInternal(format.child_formats[0], ListVector::GetEntry(vector), ListVector::GetListSize(vector));
+		break;
+	default:
+		break;
+	}
+}
+
+void TupleDataCollection::ToUnifiedFormat(TupleDataChunkState &chunk_state, DataChunk &new_chunk) {
+	D_ASSERT(chunk_state.vector_data.size() >= chunk_state.column_ids.size()); // Needs InitializeAppend
+	for (const auto &col_idx : chunk_state.column_ids) {
+		ToUnifiedFormatInternal(chunk_state.vector_data[col_idx], new_chunk.data[col_idx], new_chunk.size());
+	}
+}
+
+void TupleDataCollection::GetVectorData(const TupleDataChunkState &chunk_state, UnifiedVectorFormat result[]) {
+	const auto &vector_data = chunk_state.vector_data;
+	for (idx_t i = 0; i < vector_data.size(); i++) {
+		const auto &source = vector_data[i].data;
+		auto &target = result[i];
+		target.sel = source.sel;
+		target.data = source.data;
+		target.validity = source.validity;
+	}
+}
+
+void TupleDataCollection::Build(TupleDataPinState &pin_state, TupleDataChunkState &chunk_state,
+                                const idx_t append_offset, const idx_t append_count) {
+	segments.back().allocator->Build(segments.back(), pin_state, chunk_state, append_offset, append_count);
+	count += append_count;
+	Verify();
+}
+
+void VerifyHeapSizes(const data_ptr_t source_locations[], const idx_t heap_sizes[], const SelectionVector &append_sel,
+                     const idx_t append_count, const idx_t heap_size_offset) {
+#ifdef DEBUG
+	for (idx_t i = 0; i < append_count; i++) {
+		auto idx = append_sel.get_index(i);
+		const auto stored_heap_size = Load<uint32_t>(source_locations[idx] + heap_size_offset);
+		D_ASSERT(stored_heap_size == heap_sizes[idx]);
+	}
+#endif
+}
+
+void TupleDataCollection::CopyRows(TupleDataChunkState &chunk_state, TupleDataChunkState &input,
+                                   const SelectionVector &append_sel, const idx_t append_count) const {
+	const auto source_locations = FlatVector::GetData<data_ptr_t>(input.row_locations);
+	const auto target_locations = FlatVector::GetData<data_ptr_t>(chunk_state.row_locations);
+
+	// Copy rows
+	const auto row_width = layout.GetRowWidth();
+	for (idx_t i = 0; i < append_count; i++) {
+		auto idx = append_sel.get_index(i);
+		FastMemcpy(target_locations[i], source_locations[idx], row_width);
+	}
+
+	// Copy heap if we need to
+	if (!layout.AllConstant()) {
+		const auto source_heap_locations = FlatVector::GetData<data_ptr_t>(input.heap_locations);
+		const auto target_heap_locations = FlatVector::GetData<data_ptr_t>(chunk_state.heap_locations);
+		const auto heap_sizes = FlatVector::GetData<idx_t>(input.heap_sizes);
+		VerifyHeapSizes(source_locations, heap_sizes, append_sel, append_count, layout.GetHeapSizeOffset());
+
+		// Check if we need to copy anything at all
+		idx_t total_heap_size = 0;
+		for (idx_t i = 0; i < append_count; i++) {
+			auto idx = append_sel.get_index(i);
+			total_heap_size += heap_sizes[idx];
+		}
+		if (total_heap_size == 0) {
+			return;
+		}
+
+		// Copy heap
+		for (idx_t i = 0; i < append_count; i++) {
+			auto idx = append_sel.get_index(i);
+			FastMemcpy(target_heap_locations[i], source_heap_locations[idx], heap_sizes[idx]);
+		}
+
+		// Recompute pointers after copying the data
+		TupleDataAllocator::RecomputeHeapPointers(input.heap_locations, append_sel, target_locations,
+		                                          chunk_state.heap_locations, 0, append_count, layout, 0);
+	}
+}
+
+void TupleDataCollection::Combine(TupleDataCollection &other) {
+	if (other.count == 0) {
+		return;
+	}
+	if (this->layout.GetTypes() != other.GetLayout().GetTypes()) {
+		throw InternalException("Attempting to combine TupleDataCollection with mismatching types");
+	}
+	this->count += other.count;
+	this->segments.reserve(this->segments.size() + other.segments.size());
+	for (auto &other_seg : other.segments) {
+		this->segments.emplace_back(std::move(other_seg));
+	}
+	other.Reset();
+	Verify();
+}
+
+void TupleDataCollection::Combine(unique_ptr<TupleDataCollection> other) {
+	Combine(*other);
+}
+
+void TupleDataCollection::Reset() {
+	count = 0;
+	segments.clear();
+
+	// Refreshes the TupleDataAllocator to prevent holding on to allocated data unnecessarily
+	allocator = make_shared<TupleDataAllocator>(*allocator);
+}
+
+void TupleDataCollection::InitializeChunk(DataChunk &chunk) const {
+	chunk.Initialize(allocator->GetAllocator(), layout.GetTypes());
+}
+
+void TupleDataCollection::InitializeScanChunk(TupleDataScanState &state, DataChunk &chunk) const {
+	auto &column_ids = state.chunk_state.column_ids;
+	D_ASSERT(!column_ids.empty());
+	vector<LogicalType> chunk_types;
+	chunk_types.reserve(column_ids.size());
+	for (idx_t i = 0; i < column_ids.size(); i++) {
+		auto column_idx = column_ids[i];
+		D_ASSERT(column_idx < layout.ColumnCount());
+		chunk_types.push_back(layout.GetTypes()[column_idx]);
+	}
+	chunk.Initialize(allocator->GetAllocator(), chunk_types);
+}
+
+void TupleDataCollection::InitializeScan(TupleDataScanState &state, TupleDataPinProperties properties) const {
+	vector<column_t> column_ids;
+	column_ids.reserve(layout.ColumnCount());
+	for (idx_t i = 0; i < layout.ColumnCount(); i++) {
+		column_ids.push_back(i);
+	}
+	InitializeScan(state, std::move(column_ids), properties);
+}
+
+void TupleDataCollection::InitializeScan(TupleDataScanState &state, vector<column_t> column_ids,
+                                         TupleDataPinProperties properties) const {
+	state.pin_state.row_handles.clear();
+	state.pin_state.heap_handles.clear();
+	state.pin_state.properties = properties;
+	state.segment_index = 0;
+	state.chunk_index = 0;
+	state.chunk_state.column_ids = std::move(column_ids);
+}
+
+void TupleDataCollection::InitializeScan(TupleDataParallelScanState &gstate, TupleDataPinProperties properties) const {
+	InitializeScan(gstate.scan_state, properties);
+}
+
+void TupleDataCollection::InitializeScan(TupleDataParallelScanState &state, vector<column_t> column_ids,
+                                         TupleDataPinProperties properties) const {
+	InitializeScan(state.scan_state, std::move(column_ids), properties);
+}
+
+bool TupleDataCollection::Scan(TupleDataScanState &state, DataChunk &result) {
+	const auto segment_index_before = state.segment_index;
+	idx_t segment_index;
+	idx_t chunk_index;
+	if (!NextScanIndex(state, segment_index, chunk_index)) {
+		return false;
+	}
+	if (segment_index_before != DConstants::INVALID_INDEX && segment_index != segment_index_before) {
+		FinalizePinState(state.pin_state, segments[segment_index_before]);
+	}
+	ScanAtIndex(state.pin_state, state.chunk_state, state.chunk_state.column_ids, segment_index, chunk_index, result);
+	return true;
+}
+
+bool TupleDataCollection::Scan(TupleDataParallelScanState &gstate, TupleDataLocalScanState &lstate, DataChunk &result) {
+	lstate.pin_state.properties = gstate.scan_state.pin_state.properties;
+
+	const auto segment_index_before = lstate.segment_index;
+	idx_t segment_index;
+	idx_t chunk_index;
+	{
+		lock_guard<mutex> guard(gstate.lock);
+		if (!NextScanIndex(gstate.scan_state, segment_index, chunk_index)) {
+			return false;
+		}
+	}
+	if (segment_index_before != DConstants::INVALID_INDEX && segment_index_before != segment_index) {
+		FinalizePinState(lstate.pin_state, segments[lstate.segment_index]);
+		lstate.segment_index = segment_index;
+	}
+	ScanAtIndex(lstate.pin_state, lstate.chunk_state, gstate.scan_state.chunk_state.column_ids, segment_index,
+	            chunk_index, result);
+	return true;
+}
+
+void TupleDataCollection::FinalizePinState(TupleDataPinState &pin_state, TupleDataSegment &segment) {
+	segment.allocator->ReleaseOrStoreHandles(pin_state, segment);
+}
+
+void TupleDataCollection::FinalizePinState(TupleDataPinState &pin_state) {
+	D_ASSERT(segments.size() == 1);
+	allocator->ReleaseOrStoreHandles(pin_state, segments.back());
+}
+
+bool TupleDataCollection::NextScanIndex(TupleDataScanState &state, idx_t &segment_index, idx_t &chunk_index) {
+	// Check if we still have segments to scan
+	if (state.segment_index >= segments.size()) {
+		// No more data left in the scan
+		return false;
+	}
+	// Check within the current segment if we still have chunks to scan
+	while (state.chunk_index >= segments[state.segment_index].ChunkCount()) {
+		// Exhausted all chunks for this segment: Move to the next one
+		state.segment_index++;
+		state.chunk_index = 0;
+		if (state.segment_index >= segments.size()) {
+			return false;
+		}
+	}
+	segment_index = state.segment_index;
+	chunk_index = state.chunk_index++;
+	return true;
+}
+
+void TupleDataCollection::ScanAtIndex(TupleDataPinState &pin_state, TupleDataChunkState &chunk_state,
+                                      const vector<column_t> &column_ids, idx_t segment_index, idx_t chunk_index,
+                                      DataChunk &result) {
+	auto &segment = segments[segment_index];
+	auto &chunk = segment.chunks[chunk_index];
+	segment.allocator->InitializeChunkState(segment, pin_state, chunk_state, chunk_index, false);
+	result.Reset();
+	Gather(chunk_state.row_locations, *FlatVector::IncrementalSelectionVector(), chunk.count, column_ids, result,
+	       *FlatVector::IncrementalSelectionVector());
+	result.SetCardinality(chunk.count);
+}
+
+string TupleDataCollection::ToString() {
+	DataChunk chunk;
+	InitializeChunk(chunk);
+
+	TupleDataScanState scan_state;
+	InitializeScan(scan_state);
+
+	string result = StringUtil::Format("TupleDataCollection - [%llu Chunks, %llu Rows]\n", ChunkCount(), Count());
+	idx_t chunk_idx = 0;
+	idx_t row_count = 0;
+	while (Scan(scan_state, chunk)) {
+		result +=
+		    StringUtil::Format("Chunk %llu - [Rows %llu - %llu]\n", chunk_idx, row_count, row_count + chunk.size()) +
+		    chunk.ToString();
+		chunk_idx++;
+		row_count += chunk.size();
+	}
+
+	return result;
+}
+
+void TupleDataCollection::Print() {
+	Printer::Print(ToString());
+}
+
+void TupleDataCollection::Verify() const {
+#ifdef DEBUG
+	idx_t total_segment_count = 0;
+	for (const auto &segment : segments) {
+		segment.Verify();
+		total_segment_count += segment.count;
+	}
+	D_ASSERT(total_segment_count == this->count);
+#endif
+}
+
+void TupleDataCollection::VerifyEverythingPinned() const {
+#ifdef DEBUG
+	for (const auto &segment : segments) {
+		segment.VerifyEverythingPinned();
+	}
+#endif
+}
+
+} // namespace duckdb

package/src/duckdb/src/common/types/row/tuple_data_iterator.cpp

@@ -0,0 +1,96 @@
+#include "duckdb/common/types/row/tuple_data_iterator.hpp"
+
+#include "duckdb/common/types/row/tuple_data_allocator.hpp"
+
+namespace duckdb {
+
+TupleDataChunkIterator::TupleDataChunkIterator(TupleDataCollection &collection_p, TupleDataPinProperties properties_p,
+                                               bool init_heap)
+    : TupleDataChunkIterator(collection_p, properties_p, 0, collection_p.ChunkCount(), init_heap) {
+}
+
+TupleDataChunkIterator::TupleDataChunkIterator(TupleDataCollection &collection_p, TupleDataPinProperties properties,
+                                               idx_t chunk_idx_from, idx_t chunk_idx_to, bool init_heap_p)
+    : collection(collection_p), init_heap(init_heap_p) {
+	state.pin_state.properties = properties;
+	D_ASSERT(chunk_idx_from < chunk_idx_to);
+	D_ASSERT(chunk_idx_to <= collection.ChunkCount());
+	idx_t overall_chunk_index = 0;
+	for (idx_t segment_idx = 0; segment_idx < collection.segments.size(); segment_idx++) {
+		const auto &segment = collection.segments[segment_idx];
+		if (chunk_idx_from >= overall_chunk_index && chunk_idx_from <= overall_chunk_index + segment.ChunkCount()) {
+			// We start in this segment
+			start_segment_idx = segment_idx;
+			start_chunk_idx = chunk_idx_from - overall_chunk_index;
+		}
+		if (chunk_idx_to >= overall_chunk_index && chunk_idx_to <= overall_chunk_index + segment.ChunkCount()) {
+			// We end in this segment
+			end_segment_idx = segment_idx;
+			end_chunk_idx = chunk_idx_to - overall_chunk_index;
+		}
+		overall_chunk_index += segment.ChunkCount();
+	}
+
+	Reset();
+}
+
+void TupleDataChunkIterator::InitializeCurrentChunk() {
+	auto &segment = collection.segments[current_segment_idx];
+	segment.allocator->InitializeChunkState(segment, state.pin_state, state.chunk_state, current_chunk_idx, init_heap);
+}
+
+bool TupleDataChunkIterator::Done() const {
+	return current_segment_idx == end_segment_idx && current_chunk_idx == end_chunk_idx;
+}
+
+bool TupleDataChunkIterator::Next() {
+	D_ASSERT(!Done()); // Check if called after already done
+
+	// Set the next indices and checks if we're at the end of the collection
+	// NextScanIndex can go past this iterators 'end', so we have to check the indices again
+	const auto segment_idx_before = current_segment_idx;
+	if (!collection.NextScanIndex(state, current_segment_idx, current_chunk_idx) || Done()) {
+		// Drop pins / stores them if TupleDataPinProperties::KEEP_EVERYTHING_PINNED
+		collection.FinalizePinState(state.pin_state, collection.segments[segment_idx_before]);
+		current_segment_idx = end_segment_idx;
+		current_chunk_idx = end_chunk_idx;
+		return false;
+	}
+
+	// Finalize pin state when moving from one segment to the next
+	if (current_segment_idx != segment_idx_before) {
+		collection.FinalizePinState(state.pin_state, collection.segments[segment_idx_before]);
+	}
+
+	InitializeCurrentChunk();
+	return true;
+}
+
+void TupleDataChunkIterator::Reset() {
+	state.segment_index = start_segment_idx;
+	state.chunk_index = start_chunk_idx;
+	collection.NextScanIndex(state, current_segment_idx, current_chunk_idx);
+	InitializeCurrentChunk();
+}
+
+idx_t TupleDataChunkIterator::GetCurrentChunkCount() const {
+	return collection.segments[current_segment_idx].chunks[current_chunk_idx].count;
+}
+
+TupleDataChunkState &TupleDataChunkIterator::GetChunkState() {
+	return state.chunk_state;
+}
+
+data_ptr_t *TupleDataChunkIterator::GetRowLocations() {
+	return FlatVector::GetData<data_ptr_t>(state.chunk_state.row_locations);
+}
+
+data_ptr_t *TupleDataChunkIterator::GetHeapLocations() {
+	return FlatVector::GetData<data_ptr_t>(state.chunk_state.heap_locations);
+}
+
+idx_t *TupleDataChunkIterator::GetHeapSizes() {
+	return FlatVector::GetData<idx_t>(state.chunk_state.heap_sizes);
+}
+
+} // namespace duckdb