Repository: duckdb/duckdb
File: src/common/types/row/tuple_data_collection.cpp (C++)

Code before:
	for (idx_t seg_idx = 0; seg_idx < segments.size(); seg_idx++) {
		auto &segment = segments[seg_idx];
		idx_t chunk_count = segment->ChunkCount();

		if (current_chunk + chunk_count <= target_chunk) {
			total_rows += segment->count;
			current_chunk += chunk_count;
		} else {
			idx_t chunk_idx_in_segment = target_chunk - current_chunk;
			for (idx_t chunk_idx = 0; chunk_idx < chunk_idx_in_segment; chunk_idx++) {
				total_rows += segment->chunks[chunk_idx]->count;
			}

			// reset scan state to target segment
			state.segment_index = seg_idx;
			state.chunk_index = chunk_idx_in_segment;
			break;
		}
	}

	D_ASSERT(target_chunk < current_chunk);
	return total_rows;
}

bool TupleDataCollection::ScanComplete(const TupleDataScanState &state) const {
	if (Count() == 0) {
		return true;
	}
	return state.segment_index == segments.size() - 1 && state.chunk_index == segments.back()->ChunkCount();
}

void TupleDataCollection::FinalizePinState(TupleDataPinState &pin_state, TupleDataSegment &segment) {
	segment.allocator->ReleaseOrStoreHandles(pin_state, segment);
}

void TupleDataCollection::FinalizePinState(TupleDataPinState &pin_state) {
	D_ASSERT(!segments.empty());
	FinalizePinState(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;

Review comment (suggestion):
While this is correct here (because `IEJoinUnion::AppendKey` initializes a new scan state every time it calls `Seek`), it won't be correct when reusing the same scan state to `Seek` multiple times, as the buffer handles need to be cleared when crossing into a new segment.

I think it's safest if we add this to the top of the function:
```C++
D_ASSERT(state.pin_state.properties == TupleDataPinProperties::UNPIN_AFTER_DONE);
state.pin_state.row_handles.clear();
state.pin_state.heap_handles.clear();
```

Then, future use of this function is less likely to cause issues.

Code after:
	state.pin_state.heap_handles.clear();

	// early return for empty collection
	if (segments.empty()) {
		return 0;
	}

	idx_t current_chunk = 0;
	idx_t total_rows = 0;
	for (idx_t seg_idx = 0; seg_idx < segments.size(); seg_idx++) {
		auto &segment = segments[seg_idx];
		idx_t chunk_count = segment->ChunkCount();

		if (current_chunk + chunk_count <= target_chunk) {
			total_rows += segment->count;
			current_chunk += chunk_count;
		} else {
			idx_t chunk_idx_in_segment = target_chunk - current_chunk;
			for (idx_t chunk_idx = 0; chunk_idx < chunk_idx_in_segment; chunk_idx++) {
				total_rows += segment->chunks[chunk_idx]->count;
			}
			current_chunk += chunk_count;

			// reset scan state to target segment
			state.segment_index = seg_idx;
			state.chunk_index = chunk_idx_in_segment;
			break;
		}
	}

	D_ASSERT(target_chunk < current_chunk);
	return total_rows;
}

bool TupleDataCollection::ScanComplete(const TupleDataScanState &state) const {
	if (Count() == 0) {
		return true;
	}
	return state.segment_index == segments.size() - 1 && state.chunk_index == segments.back()->ChunkCount();
}

void TupleDataCollection::FinalizePinState(TupleDataPinState &pin_state, TupleDataSegment &segment) {
	segment.allocator->ReleaseOrStoreHandles(pin_state, segment);
}

void TupleDataCollection::FinalizePinState(TupleDataPinState &pin_state) {
	D_ASSERT(!segments.empty());
	FinalizePinState(pin_state, *segments.back());
}

bool TupleDataCollection::NextScanIndex(TupleDataScanState &state, idx_t &segment_index, idx_t &chunk_index) {