Repository: duckdb/duckdb
File: src/function/scalar/system/aggregate_export.cpp (C++)

Code before:
		const bool is_null1 = !state1_data.validity.RowIsValid(state1_data.sel->get_index(i));

		if (is_null0 && is_null1) {
			FlatVector::SetNull(result, i, true);
			continue;
		}

		if (is_null0 || is_null1) {
			auto &non_null_vec = is_null0 ? input.data[1] : input.data[0];
			auto &non_null_state_data = is_null0 ? state1_data : state0_data;

			layout.Load(non_null_vec, non_null_state_data, i, local_state.state_buffer1.get());
			layout.Store(result, i, local_state.state_buffer1.get());
			continue;
		}

		// Both are non-NULL, combine them
		layout.Load(input.data[0], state0_data, i, local_state.state_buffer0.get());
		layout.Load(input.data[1], state1_data, i, local_state.state_buffer1.get());

		AggregateInputData aggr_input_data(nullptr, local_state.allocator, AggregateCombineType::ALLOW_DESTRUCTIVE);
		bind_data.aggr.GetStateCombineCallback()(local_state.state_vector0, local_state.state_vector1, aggr_input_data,
		                                         1);

		layout.Store(result, i, local_state.state_buffer1.get());
	}
}

unique_ptr<FunctionData> BindAggregateState(ClientContext &context, ScalarFunction &bound_function,
                                            vector<unique_ptr<Expression>> &arguments) {
	// grab the aggregate type and bind the aggregate again

	// the aggregate name and types are in the logical type of the aggregate state, make sure its sane
	auto &arg_return_type = arguments[0]->return_type;
	for (auto &arg_type : bound_function.arguments) {
		arg_type = arg_return_type;
	}

	if (arg_return_type.id() != LogicalTypeId::LEGACY_AGGREGATE_STATE &&
	    arg_return_type.id() != LogicalTypeId::AGGREGATE_STATE) {
		throw BinderException("Can only FINALIZE aggregate state, not %s", arg_return_type.ToString());
	}
	// combine
	if (arguments.size() == 2 && arguments[0]->return_type != arguments[1]->return_type &&
	    arguments[1]->return_type.id() != LogicalTypeId::BLOB) {
		throw BinderException("Cannot COMBINE aggregate states from different functions, %s <> %s",
		                      arguments[0]->return_type.ToString(), arguments[1]->return_type.ToString());
	}

	// following error states are only reachable when someone messes up creating the state_type
	// which is impossible from SQL

Review comment (suggestion):
The logic seems correct, but we're looping over the data 5 times:
1. Fill null selection vectors
2. Set both null rows using `FlatVector::SetNull`
3. Copy from input 0
4. Copy from input 1
5. Combine both inputs

I think it's OK for now especially since we're using the selection vectors for both the struct and the non-struct combine, and it also makes the code very readable. I think it should be possible to reduce the number of iterations, e.g., by initializing the NULL states to "empty" (maybe using `AggregateStateInitialize`) once when get rid of the legacy stuff - but that's a worry for later :)

Code after:
	input.data[1].ToUnifiedFormat(input.size(), state1_data);

	// Partition rows by NULL using SelectionVector
	SelectionVector both_null_sel(STANDARD_VECTOR_SIZE);
	// input1 is null
	SelectionVector copy_from_0_sel(STANDARD_VECTOR_SIZE);
	// input0 is null
	SelectionVector copy_from_1_sel(STANDARD_VECTOR_SIZE);
	SelectionVector both_valid_sel(STANDARD_VECTOR_SIZE);

	idx_t both_null_count = 0, copy_from_0_count = 0, copy_from_1_count = 0, both_valid_count = 0;

	for (idx_t i = 0; i < input.size(); i++) {
		const bool is_null0 = !state0_data.validity.RowIsValid(state0_data.sel->get_index(i));
		const bool is_null1 = !state1_data.validity.RowIsValid(state1_data.sel->get_index(i));

		if (is_null0 && is_null1) {
			both_null_sel.set_index(both_null_count++, i);
		} else if (is_null0) {
			copy_from_1_sel.set_index(copy_from_1_count++, i);
		} else if (is_null1) {
			copy_from_0_sel.set_index(copy_from_0_count++, i);
		} else {
			both_valid_sel.set_index(both_valid_count++, i);
		}
	}

	D_ASSERT(both_null_count + copy_from_0_count + copy_from_1_count + both_valid_count == input.size());

	// Handle both-null rows
	for (idx_t i = 0; i < both_null_count; i++) {
		FlatVector::SetNull(result, both_null_sel.get_index(i), true);
	}

	// Handle one-null rows - copy non-null input directly to result
	// copy_from_0: input1 is null, copy input0
	if (copy_from_0_count > 0) {
		if (layout.is_struct) {
			idx_t offset_in_state = 0;
			for (idx_t field_idx = 0; field_idx < layout.child_types->size(); field_idx++) {
				auto &field_type = layout.child_types->at(field_idx).second;
				auto physical = field_type.InternalType();
				idx_t field_size = GetTypeIdSize(physical);
				idx_t alignment = MinValue<idx_t>(field_size, 8);
				offset_in_state = AlignValue(offset_in_state, alignment);
				TemplateDispatch<CopyFromInputFieldOp>(physical, input.data[0], result, field_idx, copy_from_0_sel,
				                                       copy_from_0_count, state0_data);
				offset_in_state += field_size;
			}
		} else {
			for (idx_t i = 0; i < copy_from_0_count; i++) {