duckdb 0.8.0 → 0.8.1-dev111.0

package/src/duckdb/src/execution/operator/join/physical_index_join.cpp

@@ -67,7 +67,7 @@ PhysicalIndexJoin::PhysicalIndexJoin(LogicalOperator &op, unique_ptr<PhysicalOpe
       left_projection_map(left_projection_map_p), right_projection_map(std::move(right_projection_map_p)),
       index(index_p), conditions(std::move(cond)), join_type(join_type), lhs_first(lhs_first) {
 	D_ASSERT(right->type == PhysicalOperatorType::TABLE_SCAN);
-	auto &tbl_scan = (PhysicalTableScan &)*right;
+	auto &tbl_scan = right->Cast<PhysicalTableScan>();
 	column_ids = std::move(column_ids_p);
 	children.push_back(std::move(left));
 	children.push_back(std::move(right));
@@ -109,7 +109,7 @@ unique_ptr<OperatorState> PhysicalIndexJoin::GetOperatorState(ExecutionContext &
 
 void PhysicalIndexJoin::Output(ExecutionContext &context, DataChunk &input, DataChunk &chunk,
                                OperatorState &state_p) const {
-	auto &phy_tbl_scan = (PhysicalTableScan &)*children[1];
+	auto &phy_tbl_scan = children[1]->Cast<PhysicalTableScan>();
 	auto &bind_tbl = phy_tbl_scan.bind_data->Cast<TableScanBindData>();
 	auto &transaction = DuckTransaction::Get(context.client, bind_tbl.table.catalog);
 	auto &state = state_p.Cast<IndexJoinOperatorState>();
@@ -139,7 +139,7 @@ void PhysicalIndexJoin::Output(ExecutionContext &context, DataChunk &input, Data
 		}
 		state.rhs_chunk.Reset();
 		state.fetch_state = make_uniq<ColumnFetchState>();
-		Vector row_ids(LogicalType::ROW_TYPE, (data_ptr_t)&fetch_rows[0]);
+		Vector row_ids(LogicalType::ROW_TYPE, data_ptr_cast(&fetch_rows[0]));
 		tbl.Fetch(transaction, state.rhs_chunk, fetch_ids, row_ids, output_sel_idx, *state.fetch_state);
 	}
 

package/src/duckdb/src/execution/operator/join/physical_nested_loop_join.cpp

@@ -443,7 +443,7 @@ unique_ptr<GlobalSourceState> PhysicalNestedLoopJoin::GetGlobalSourceState(Clien
 
 unique_ptr<LocalSourceState> PhysicalNestedLoopJoin::GetLocalSourceState(ExecutionContext &context,
                                                                          GlobalSourceState &gstate) const {
-	return make_uniq<NestedLoopJoinLocalScanState>(*this, (NestedLoopJoinGlobalScanState &)gstate);
+	return make_uniq<NestedLoopJoinLocalScanState>(*this, gstate.Cast<NestedLoopJoinGlobalScanState>());
 }
 
 SourceResultType PhysicalNestedLoopJoin::GetData(ExecutionContext &context, DataChunk &chunk,
@@ -451,8 +451,8 @@ SourceResultType PhysicalNestedLoopJoin::GetData(ExecutionContext &context, Data
 	D_ASSERT(IsRightOuterJoin(join_type));
 	// check if we need to scan any unmatched tuples from the RHS for the full/right outer join
 	auto &sink = sink_state->Cast<NestedLoopJoinGlobalState>();
-	auto &gstate = (NestedLoopJoinGlobalScanState &)input.global_state;
-	auto &lstate = (NestedLoopJoinLocalScanState &)input.local_state;
+	auto &gstate = input.global_state.Cast<NestedLoopJoinGlobalScanState>();
+	auto &lstate = input.local_state.Cast<NestedLoopJoinLocalScanState>();
 
 	// if the LHS is exhausted in a FULL/RIGHT OUTER JOIN, we scan chunks we still need to output
 	sink.right_outer.Scan(gstate.scan_state, lstate.scan_state, chunk);

package/src/duckdb/src/execution/operator/join/physical_piecewise_merge_join.cpp

@@ -385,7 +385,7 @@ static idx_t MergeJoinSimpleBlocks(PiecewiseMergeJoinState &lstate, MergeJoinGlo
 
 void PhysicalPiecewiseMergeJoin::ResolveSimpleJoin(ExecutionContext &context, DataChunk &input, DataChunk &chunk,
                                                    OperatorState &state_p) const {
-	auto &state = (PiecewiseMergeJoinState &)state_p;
+	auto &state = state_p.Cast<PiecewiseMergeJoinState>();
 	auto &gstate = sink_state->Cast<MergeJoinGlobalState>();
 
 	state.ResolveJoinKeys(input);
@@ -507,7 +507,7 @@ static idx_t MergeJoinComplexBlocks(BlockMergeInfo &l, BlockMergeInfo &r, const
 
 OperatorResultType PhysicalPiecewiseMergeJoin::ResolveComplexJoin(ExecutionContext &context, DataChunk &input,
                                                                   DataChunk &chunk, OperatorState &state_p) const {
-	auto &state = (PiecewiseMergeJoinState &)state_p;
+	auto &state = state_p.Cast<PiecewiseMergeJoinState>();
 	auto &gstate = sink_state->Cast<MergeJoinGlobalState>();
 	auto &rsorted = *gstate.table->global_sort_state.sorted_blocks[0];
 	const auto left_cols = input.ColumnCount();
@@ -682,7 +682,7 @@ SourceResultType PhysicalPiecewiseMergeJoin::GetData(ExecutionContext &context,
 	D_ASSERT(IsRightOuterJoin(join_type));
 	// check if we need to scan any unmatched tuples from the RHS for the full/right outer join
 	auto &sink = sink_state->Cast<MergeJoinGlobalState>();
-	auto &state = (PiecewiseJoinScanState &)input.global_state;
+	auto &state = input.global_state.Cast<PiecewiseJoinScanState>();
 
 	lock_guard<mutex> l(state.lock);
 	if (!state.scanner) {

package/src/duckdb/src/execution/operator/order/physical_top_n.cpp

@@ -276,7 +276,7 @@ void TopNHeap::Finalize() {
 }
 
 void TopNHeap::Reduce() {
-	idx_t min_sort_threshold = MaxValue<idx_t>(STANDARD_VECTOR_SIZE * 5, 2 * (limit + offset));
+	idx_t min_sort_threshold = MaxValue<idx_t>(STANDARD_VECTOR_SIZE * 5ULL, 2ULL * (limit + offset));
 	if (sort_state.count < min_sort_threshold) {
 		// only reduce when we pass two times the limit + offset, or 5 vectors (whichever comes first)
 		return;

package/src/duckdb/src/execution/operator/persistent/buffered_csv_reader.cpp

@@ -186,8 +186,8 @@ static string TrimWhitespace(const string &col_name) {
 
 static string NormalizeColumnName(const string &col_name) {
 	// normalize UTF8 characters to NFKD
-	auto nfkd = utf8proc_NFKD((const utf8proc_uint8_t *)col_name.c_str(), col_name.size());
-	const string col_name_nfkd = string((const char *)nfkd, strlen((const char *)nfkd));
+	auto nfkd = utf8proc_NFKD(reinterpret_cast<const utf8proc_uint8_t *>(col_name.c_str()), col_name.size());
+	const string col_name_nfkd = string(const_char_ptr_cast(nfkd), strlen(const_char_ptr_cast(nfkd)));
 	free(nfkd);
 
 	// only keep ASCII characters 0-9 a-z A-Z and replace spaces with regular whitespace

package/src/duckdb/src/execution/operator/persistent/csv_file_handle.cpp

@@ -99,7 +99,7 @@ idx_t CSVFileHandle::Read(void *buffer, idx_t nr_bytes) {
 	}
 	// we have data left to read from the file
 	// read directly into the buffer
-	auto bytes_read = file_handle->Read((char *)buffer + result_offset, nr_bytes - result_offset);
+	auto bytes_read = file_handle->Read(char_ptr_cast(buffer) + result_offset, nr_bytes - result_offset);
 	file_size = file_handle->GetFileSize();
 	read_position += bytes_read;
 	if (reset_enabled) {
@@ -114,7 +114,7 @@ idx_t CSVFileHandle::Read(void *buffer, idx_t nr_bytes) {
 			}
 			cached_buffer = std::move(new_buffer);
 		}
-		memcpy(cached_buffer.get() + buffer_size, (char *)buffer + result_offset, bytes_read);
+		memcpy(cached_buffer.get() + buffer_size, char_ptr_cast(buffer) + result_offset, bytes_read);
 		buffer_size += bytes_read;
 	}
 

package/src/duckdb/src/execution/operator/persistent/csv_reader_options.cpp

@@ -177,6 +177,8 @@ void BufferedCSVReaderOptions::SetReadOption(const string &loption, const Value
 		null_padding = ParseBoolean(value, loption);
 	} else if (loption == "allow_quoted_nulls") {
 		allow_quoted_nulls = ParseBoolean(value, loption);
+	} else if (loption == "parallel") {
+		parallel_mode = ParseBoolean(value, loption) ? ParallelMode::PARALLEL : ParallelMode::SINGLE_THREADED;
 	} else {
 		throw BinderException("Unrecognized option for CSV reader \"%s\"", loption);
 	}

package/src/duckdb/src/execution/operator/persistent/physical_batch_insert.cpp

@@ -279,9 +279,8 @@ unique_ptr<GlobalSinkState> PhysicalBatchInsert::GetGlobalSinkState(ClientContex
 		// CREATE TABLE AS
 		D_ASSERT(!insert_table);
 		auto &catalog = schema->catalog;
-		table = (TableCatalogEntry *)catalog
-		            .CreateTable(catalog.GetCatalogTransaction(context), *schema.get_mutable(), *info)
-		            .get();
+		auto created_table = catalog.CreateTable(catalog.GetCatalogTransaction(context), *schema.get_mutable(), *info);
+		table = &created_table->Cast<TableCatalogEntry>();
 	} else {
 		D_ASSERT(insert_table);
 		D_ASSERT(insert_table->IsDuckTable());

package/src/duckdb/src/execution/operator/persistent/physical_insert.cpp

@@ -63,7 +63,7 @@ PhysicalInsert::PhysicalInsert(LogicalOperator &op, SchemaCatalogEntry &schema,
 
 void PhysicalInsert::GetInsertInfo(const BoundCreateTableInfo &info, vector<LogicalType> &insert_types,
                                    vector<unique_ptr<Expression>> &bound_defaults) {
-	auto &create_info = (CreateTableInfo &)*info.base;
+	auto &create_info = info.base->Cast<CreateTableInfo>();
 	for (auto &col : create_info.columns.Physical()) {
 		insert_types.push_back(col.GetType());
 		bound_defaults.push_back(make_uniq<BoundConstantExpression>(Value(col.GetType())));

package/src/duckdb/src/execution/operator/projection/physical_pivot.cpp

@@ -17,11 +17,11 @@ PhysicalPivot::PhysicalPivot(vector<LogicalType> types_p, unique_ptr<PhysicalOpe
 	}
 	// extract the empty aggregate expressions
 	for (auto &aggr_expr : bound_pivot.aggregates) {
-		auto &aggr = (BoundAggregateExpression &)*aggr_expr;
+		auto &aggr = aggr_expr->Cast<BoundAggregateExpression>();
 		// for each aggregate, initialize an empty aggregate state and finalize it immediately
 		auto state = make_unsafe_uniq_array<data_t>(aggr.function.state_size());
 		aggr.function.initialize(state.get());
-		Vector state_vector(Value::POINTER((uintptr_t)state.get()));
+		Vector state_vector(Value::POINTER(CastPointerToValue(state.get())));
 		Vector result_vector(aggr_expr->return_type);
 		AggregateInputData aggr_input_data(aggr.bind_info.get(), Allocator::DefaultAllocator());
 		aggr.function.finalize(state_vector, aggr_input_data, result_vector, 1, 0);

package/src/duckdb/src/execution/operator/projection/physical_projection.cpp

@@ -27,7 +27,7 @@ PhysicalProjection::PhysicalProjection(vector<LogicalType> types, vector<unique_
 
 OperatorResultType PhysicalProjection::Execute(ExecutionContext &context, DataChunk &input, DataChunk &chunk,
                                                GlobalOperatorState &gstate, OperatorState &state_p) const {
-	auto &state = (ProjectionState &)state_p;
+	auto &state = state_p.Cast<ProjectionState>();
 	state.executor.Execute(input, chunk);
 	return OperatorResultType::NEED_MORE_INPUT;
 }

package/src/duckdb/src/execution/operator/projection/physical_unnest.cpp

@@ -65,7 +65,7 @@ void UnnestOperatorState::SetLongestListLength() {
 		if (vector_data.validity.RowIsValid(current_idx)) {
 
 			// check if this list is longer
-			auto list_data = (list_entry_t *)vector_data.data;
+			auto list_data = UnifiedVectorFormat::GetData<list_entry_t>(vector_data);
 			auto list_entry = list_data[current_idx];
 			if (list_entry.length > longest_list_length) {
 				longest_list_length = list_entry.length;
@@ -98,7 +98,7 @@ static void UnnestNull(idx_t start, idx_t end, Vector &result) {
 template <class T>
 static void TemplatedUnnest(UnifiedVectorFormat &vector_data, idx_t start, idx_t end, Vector &result) {
 
-	auto source_data = (T *)vector_data.data;
+	auto source_data = UnifiedVectorFormat::GetData<T>(vector_data);
 	auto &source_mask = vector_data.validity;
 
 	D_ASSERT(result.GetVectorType() == VectorType::FLAT_VECTOR);
@@ -314,7 +314,7 @@ OperatorResultType PhysicalUnnest::ExecuteInternal(ExecutionContext &context, Da
 
 				} else {
 
-					auto list_data = (list_entry_t *)vector_data.data;
+					auto list_data = UnifiedVectorFormat::GetData<list_entry_t>(vector_data);
 					auto list_entry = list_data[current_idx];
 
 					idx_t list_count = 0;

package/src/duckdb/src/execution/operator/scan/physical_column_data_scan.cpp

@@ -1,6 +1,7 @@
 #include "duckdb/execution/operator/scan/physical_column_data_scan.hpp"
 
 #include "duckdb/execution/operator/join/physical_delim_join.hpp"
+#include "duckdb/execution/operator/aggregate/physical_hash_aggregate.hpp"
 #include "duckdb/parallel/meta_pipeline.hpp"
 #include "duckdb/parallel/pipeline.hpp"
 
@@ -60,7 +61,7 @@ void PhysicalColumnDataScan::BuildPipelines(Pipeline &current, MetaPipeline &met
 		D_ASSERT(delim_sink->type == PhysicalOperatorType::DELIM_JOIN);
 		auto &delim_join = delim_sink->Cast<PhysicalDelimJoin>();
 		current.AddDependency(delim_dependency);
-		state.SetPipelineSource(current, (PhysicalOperator &)*delim_join.distinct);
+		state.SetPipelineSource(current, delim_join.distinct->Cast<PhysicalOperator>());
 		return;
 	}
 	case PhysicalOperatorType::RECURSIVE_CTE_SCAN:

package/src/duckdb/src/execution/operator/scan/physical_expression_scan.cpp

@@ -22,7 +22,7 @@ unique_ptr<OperatorState> PhysicalExpressionScan::GetOperatorState(ExecutionCont
 
 OperatorResultType PhysicalExpressionScan::Execute(ExecutionContext &context, DataChunk &input, DataChunk &chunk,
                                                    GlobalOperatorState &gstate, OperatorState &state_p) const {
-	auto &state = (ExpressionScanState &)state_p;
+	auto &state = state_p.Cast<ExpressionScanState>();
 
 	for (; chunk.size() + input.size() <= STANDARD_VECTOR_SIZE && state.expression_index < expressions.size();
 	     state.expression_index++) {

package/src/duckdb/src/execution/operator/scan/physical_positional_scan.cpp

@@ -19,7 +19,7 @@ PhysicalPositionalScan::PhysicalPositionalScan(vector<LogicalType> types, unique
 	if (left->type == PhysicalOperatorType::TABLE_SCAN) {
 		child_tables.emplace_back(std::move(left));
 	} else if (left->type == PhysicalOperatorType::POSITIONAL_SCAN) {
-		auto &left_scan = (PhysicalPositionalScan &)*left;
+		auto &left_scan = left->Cast<PhysicalPositionalScan>();
 		child_tables = std::move(left_scan.child_tables);
 	} else {
 		throw InternalException("Invalid left input for PhysicalPositionalScan");
@@ -28,7 +28,7 @@ PhysicalPositionalScan::PhysicalPositionalScan(vector<LogicalType> types, unique
 	if (right->type == PhysicalOperatorType::TABLE_SCAN) {
 		child_tables.emplace_back(std::move(right));
 	} else if (right->type == PhysicalOperatorType::POSITIONAL_SCAN) {
-		auto &right_scan = (PhysicalPositionalScan &)*right;
+		auto &right_scan = right->Cast<PhysicalPositionalScan>();
 		auto &right_tables = right_scan.child_tables;
 		child_tables.reserve(child_tables.size() + right_tables.size());
 		std::move(right_tables.begin(), right_tables.end(), std::back_inserter(child_tables));

package/src/duckdb/src/execution/operator/schema/physical_create_type.cpp

@@ -45,7 +45,7 @@ SinkResultType PhysicalCreateType::Sink(ExecutionContext &context, DataChunk &ch
 		gstate.capacity *= 2;
 	}
 
-	auto src_ptr = (string_t *)sdata.data;
+	auto src_ptr = UnifiedVectorFormat::GetData<string_t>(sdata);
 	auto result_ptr = FlatVector::GetData<string_t>(gstate.result);
 	// Input vector has NULL value, we just throw an exception
 	for (idx_t i = 0; i < chunk.size(); i++) {

package/src/duckdb/src/execution/perfect_aggregate_hashtable.cpp

@@ -51,7 +51,7 @@ PerfectAggregateHashTable::~PerfectAggregateHashTable() {
 template <class T>
 static void ComputeGroupLocationTemplated(UnifiedVectorFormat &group_data, Value &min, uintptr_t *address_data,
                                           idx_t current_shift, idx_t count) {
-	auto data = (T *)group_data.data;
+	auto data = UnifiedVectorFormat::GetData<T>(group_data);
 	auto min_val = min.GetValueUnsafe<T>();
 	if (!group_data.validity.AllValid()) {
 		for (idx_t i = 0; i < count; i++) {

package/src/duckdb/src/execution/physical_plan/plan_comparison_join.cpp

@@ -186,7 +186,7 @@ static bool PlanIndexJoin(ClientContext &context, LogicalComparisonJoin &op, uni
 	}
 	// index joins are not supported if there are pushed down table filters
 	D_ASSERT(right->type == PhysicalOperatorType::TABLE_SCAN);
-	auto &tbl_scan = (PhysicalTableScan &)*right;
+	auto &tbl_scan = right->Cast<PhysicalTableScan>();
 	//	if (tbl_scan.table_filters && !tbl_scan.table_filters->filters.empty()) {
 	//		return false;
 	//	}

package/src/duckdb/src/execution/physical_plan/plan_create_table.cpp

@@ -34,7 +34,7 @@ unique_ptr<PhysicalOperator> DuckCatalog::PlanCreateTableAs(ClientContext &conte
 }
 
 unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalCreateTable &op) {
-	const auto &create_info = (CreateTableInfo &)*op.info->base;
+	const auto &create_info = op.info->base->Cast<CreateTableInfo>();
 	auto &catalog = op.info->schema.catalog;
 	auto existing_entry = catalog.GetEntry<TableCatalogEntry>(context, create_info.schema, create_info.table,
 	                                                          OnEntryNotFound::RETURN_NULL);

package/src/duckdb/src/execution/physical_plan_generator.cpp

@@ -207,10 +207,10 @@ unique_ptr<PhysicalOperator> PhysicalPlanGenerator::CreatePlan(LogicalOperator &
 		plan = CreatePlan(op.Cast<LogicalReset>());
 		break;
 	case LogicalOperatorType::LOGICAL_PIVOT:
-		plan = CreatePlan((LogicalPivot &)op);
+		plan = CreatePlan(op.Cast<LogicalPivot>());
 		break;
 	case LogicalOperatorType::LOGICAL_EXTENSION_OPERATOR:
-		plan = ((LogicalExtensionOperator &)op).CreatePlan(context, *this);
+		plan = op.Cast<LogicalExtensionOperator>().CreatePlan(context, *this);
 
 		if (!plan) {
 			throw InternalException("Missing PhysicalOperator for Extension Operator");

package/src/duckdb/src/execution/radix_partitioned_hashtable.cpp

@@ -408,7 +408,7 @@ SourceResultType RadixPartitionedHashTable::GetData(ExecutionContext &context, D
 			aggr.function.initialize(aggr_state.get());
 
 			AggregateInputData aggr_input_data(aggr.bind_info.get(), Allocator::DefaultAllocator());
-			Vector state_vector(Value::POINTER((uintptr_t)aggr_state.get()));
+			Vector state_vector(Value::POINTER(CastPointerToValue(aggr_state.get())));
 			aggr.function.finalize(state_vector, aggr_input_data, chunk.data[null_groups.size() + i], 1, 0);
 			if (aggr.function.destructor) {
 				aggr.function.destructor(state_vector, aggr_input_data, 1);

package/src/duckdb/src/execution/reservoir_sample.cpp

@@ -186,7 +186,7 @@ void BaseReservoirSampling::InitializeReservoir(idx_t cur_size, idx_t sample_siz
 		//! we use a priority queue to extract the minimum key in O(1) time
 		for (idx_t i = 0; i < sample_size; i++) {
 			double k_i = random.NextRandom();
-			reservoir_weights.push(std::make_pair(-k_i, i));
+			reservoir_weights.emplace(-k_i, i);
 		}
 		SetNextEntry();
 	}
@@ -217,7 +217,7 @@ void BaseReservoirSampling::ReplaceElement() {
 	//! we generate a random number between (min_threshold, 1)
 	double r2 = random.NextRandom(min_threshold, 1);
 	//! now we insert the new weight into the reservoir
-	reservoir_weights.push(std::make_pair(-r2, min_entry));
+	reservoir_weights.emplace(-r2, min_entry);
 	//! we update the min entry with the new min entry in the reservoir
 	SetNextEntry();
 }

package/src/duckdb/src/execution/window_segment_tree.cpp

@@ -12,7 +12,8 @@ namespace duckdb {
 
 WindowAggregateState::WindowAggregateState(AggregateObject aggr, const LogicalType &result_type_p)
     : aggr(std::move(aggr)), result_type(result_type_p), state(aggr.function.state_size()),
-      statev(Value::POINTER((idx_t)state.data())), statep(Value::POINTER((idx_t)state.data())) {
+      statev(Value::POINTER(CastPointerToValue(state.data()))),
+      statep(Value::POINTER(CastPointerToValue(state.data()))) {
 	statev.SetVectorType(VectorType::FLAT_VECTOR); // Prevent conversion of results to constants
 }
 
@@ -180,11 +181,12 @@ void WindowConstantAggregate::Compute(Vector &target, idx_t rid, idx_t start, id
 //===--------------------------------------------------------------------===//
 // WindowSegmentTree
 //===--------------------------------------------------------------------===//
-WindowSegmentTree::WindowSegmentTree(AggregateObject aggr, const LogicalType &result_type_p, DataChunk *input,
+WindowSegmentTree::WindowSegmentTree(AggregateObject aggr_p, const LogicalType &result_type_p, DataChunk *input,
                                      const ValidityMask &filter_mask_p, WindowAggregationMode mode_p)
-    : aggr(std::move(aggr)), result_type(result_type_p), state(aggr.function.state_size()),
-      statep(Value::POINTER((idx_t)state.data())), frame(0, 0), statev(Value::POINTER((idx_t)state.data())),
-      internal_nodes(0), input_ref(input), filter_mask(filter_mask_p), mode(mode_p) {
+    : aggr(std::move(aggr_p)), result_type(result_type_p), state(aggr.function.state_size()),
+      statep(Value::POINTER(CastPointerToValue(state.data()))), frame(0, 0),
+      statev(Value::POINTER(CastPointerToValue(state.data()))), internal_nodes(0), input_ref(input),
+      filter_mask(filter_mask_p), mode(mode_p) {
 	statep.Flatten(input->size());
 	statev.SetVectorType(VectorType::FLAT_VECTOR); // Prevent conversion of results to constants
 
@@ -212,7 +214,7 @@ WindowSegmentTree::~WindowSegmentTree() {
 	AggregateInputData aggr_input_data(aggr.GetFunctionData(), Allocator::DefaultAllocator());
 	// call the destructor for all the intermediate states
 	data_ptr_t address_data[STANDARD_VECTOR_SIZE];
-	Vector addresses(LogicalType::POINTER, (data_ptr_t)address_data);
+	Vector addresses(LogicalType::POINTER, data_ptr_cast(address_data));
 	idx_t count = 0;
 	for (idx_t i = 0; i < internal_nodes; i++) {
 		address_data[count++] = data_ptr_t(levels_flat_native.get() + i * state.size());

package/src/duckdb/src/function/aggregate/distributive/count.cpp

@@ -7,30 +7,30 @@ namespace duckdb {
 
 struct BaseCountFunction {
 	template <class STATE>
-	static void Initialize(STATE *state) {
-		*state = 0;
+	static void Initialize(STATE &state) {
+		state = 0;
 	}
 
 	template <class STATE, class OP>
-	static void Combine(const STATE &source, STATE *target, AggregateInputData &) {
-		*target += source;
+	static void Combine(const STATE &source, STATE &target, AggregateInputData &) {
+		target += source;
 	}
 
 	template <class T, class STATE>
-	static void Finalize(Vector &result, AggregateInputData &, STATE *state, T *target, ValidityMask &mask, idx_t idx) {
-		target[idx] = *state;
+	static void Finalize(STATE &state, T &target, AggregateFinalizeData &finalize_data) {
+		target = state;
 	}
 };
 
 struct CountStarFunction : public BaseCountFunction {
 	template <class STATE, class OP>
-	static void Operation(STATE *state, AggregateInputData &, idx_t idx) {
-		*state += 1;
+	static void Operation(STATE &state, AggregateInputData &, idx_t idx) {
+		state += 1;
 	}
 
 	template <class STATE, class OP>
-	static void ConstantOperation(STATE *state, AggregateInputData &, idx_t count) {
-		*state += count;
+	static void ConstantOperation(STATE &state, AggregateInputData &, idx_t count) {
+		state += count;
 	}
 
 	template <typename RESULT_TYPE>
@@ -55,27 +55,165 @@ struct CountStarFunction : public BaseCountFunction {
 };
 
 struct CountFunction : public BaseCountFunction {
-	template <class INPUT_TYPE, class STATE, class OP>
-	static void Operation(STATE *state, AggregateInputData &, INPUT_TYPE *input, ValidityMask &mask, idx_t idx) {
-		*state += 1;
+	using STATE = int64_t;
+
+	static void Operation(STATE &state) {
+		state += 1;
 	}
 
-	template <class INPUT_TYPE, class STATE, class OP>
-	static void ConstantOperation(STATE *state, AggregateInputData &, INPUT_TYPE *input, ValidityMask &mask,
-	                              idx_t count) {
-		*state += count;
+	static void ConstantOperation(STATE &state, idx_t count) {
+		state += count;
 	}
 
 	static bool IgnoreNull() {
 		return true;
 	}
+
+	static inline void CountFlatLoop(STATE **__restrict states, ValidityMask &mask, idx_t count) {
+		if (!mask.AllValid()) {
+			idx_t base_idx = 0;
+			auto entry_count = ValidityMask::EntryCount(count);
+			for (idx_t entry_idx = 0; entry_idx < entry_count; entry_idx++) {
+				auto validity_entry = mask.GetValidityEntry(entry_idx);
+				idx_t next = MinValue<idx_t>(base_idx + ValidityMask::BITS_PER_VALUE, count);
+				if (ValidityMask::AllValid(validity_entry)) {
+					// all valid: perform operation
+					for (; base_idx < next; base_idx++) {
+						CountFunction::Operation(*states[base_idx]);
+					}
+				} else if (ValidityMask::NoneValid(validity_entry)) {
+					// nothing valid: skip all
+					base_idx = next;
+					continue;
+				} else {
+					// partially valid: need to check individual elements for validity
+					idx_t start = base_idx;
+					for (; base_idx < next; base_idx++) {
+						if (ValidityMask::RowIsValid(validity_entry, base_idx - start)) {
+							CountFunction::Operation(*states[base_idx]);
+						}
+					}
+				}
+			}
+		} else {
+			for (idx_t i = 0; i < count; i++) {
+				CountFunction::Operation(*states[i]);
+			}
+		}
+	}
+
+	static inline void CountScatterLoop(STATE **__restrict states, const SelectionVector &isel,
+	                                    const SelectionVector &ssel, ValidityMask &mask, idx_t count) {
+		if (!mask.AllValid()) {
+			// potential NULL values
+			for (idx_t i = 0; i < count; i++) {
+				auto idx = isel.get_index(i);
+				auto sidx = ssel.get_index(i);
+				if (mask.RowIsValid(idx)) {
+					CountFunction::Operation(*states[sidx]);
+				}
+			}
+		} else {
+			// quick path: no NULL values
+			for (idx_t i = 0; i < count; i++) {
+				auto sidx = ssel.get_index(i);
+				CountFunction::Operation(*states[sidx]);
+			}
+		}
+	}
+
+	static void CountScatter(Vector inputs[], AggregateInputData &aggr_input_data, idx_t input_count, Vector &states,
+	                         idx_t count) {
+		auto &input = inputs[0];
+		if (input.GetVectorType() == VectorType::FLAT_VECTOR && states.GetVectorType() == VectorType::FLAT_VECTOR) {
+			auto sdata = FlatVector::GetData<STATE *>(states);
+			CountFlatLoop(sdata, FlatVector::Validity(input), count);
+		} else {
+			UnifiedVectorFormat idata, sdata;
+			input.ToUnifiedFormat(count, idata);
+			states.ToUnifiedFormat(count, sdata);
+			CountScatterLoop(reinterpret_cast<STATE **>(sdata.data), *idata.sel, *sdata.sel, idata.validity, count);
+		}
+	}
+
+	static inline void CountFlatUpdateLoop(STATE &result, ValidityMask &mask, idx_t count) {
+		idx_t base_idx = 0;
+		auto entry_count = ValidityMask::EntryCount(count);
+		for (idx_t entry_idx = 0; entry_idx < entry_count; entry_idx++) {
+			auto validity_entry = mask.GetValidityEntry(entry_idx);
+			idx_t next = MinValue<idx_t>(base_idx + ValidityMask::BITS_PER_VALUE, count);
+			if (ValidityMask::AllValid(validity_entry)) {
+				// all valid
+				result += next - base_idx;
+				base_idx = next;
+			} else if (ValidityMask::NoneValid(validity_entry)) {
+				// nothing valid: skip all
+				base_idx = next;
+				continue;
+			} else {
+				// partially valid: need to check individual elements for validity
+				idx_t start = base_idx;
+				for (; base_idx < next; base_idx++) {
+					if (ValidityMask::RowIsValid(validity_entry, base_idx - start)) {
+						result++;
+					}
+				}
+			}
+		}
+	}
+
+	static inline void CountUpdateLoop(STATE &result, ValidityMask &mask, idx_t count,
+	                                   const SelectionVector &sel_vector) {
+		if (mask.AllValid()) {
+			// no NULL values
+			result += count;
+			return;
+		}
+		for (idx_t i = 0; i < count; i++) {
+			auto idx = sel_vector.get_index(i);
+			if (mask.RowIsValid(idx)) {
+				result++;
+			}
+		}
+	}
+
+	static void CountUpdate(Vector inputs[], AggregateInputData &, idx_t input_count, data_ptr_t state_p, idx_t count) {
+		auto &input = inputs[0];
+		auto &result = *reinterpret_cast<STATE *>(state_p);
+		switch (input.GetVectorType()) {
+		case VectorType::CONSTANT_VECTOR: {
+			if (!ConstantVector::IsNull(input)) {
+				// if the constant is not null increment the state
+				result += count;
+			}
+			break;
+		}
+		case VectorType::FLAT_VECTOR: {
+			CountFlatUpdateLoop(result, FlatVector::Validity(input), count);
+			break;
+		}
+		case VectorType::SEQUENCE_VECTOR: {
+			// sequence vectors cannot have NULL values
+			result += count;
+			break;
+		}
+		default: {
+			UnifiedVectorFormat idata;
+			input.ToUnifiedFormat(count, idata);
+			CountUpdateLoop(result, idata.validity, count, *idata.sel);
+			break;
+		}
+		}
+	}
 };
 
 AggregateFunction CountFun::GetFunction() {
-	auto fun = AggregateFunction::UnaryAggregate<int64_t, int64_t, int64_t, CountFunction>(
-	    LogicalType(LogicalTypeId::ANY), LogicalType::BIGINT);
+	AggregateFunction fun({LogicalType(LogicalTypeId::ANY)}, LogicalType::BIGINT, AggregateFunction::StateSize<int64_t>,
+	                      AggregateFunction::StateInitialize<int64_t, CountFunction>, CountFunction::CountScatter,
+	                      AggregateFunction::StateCombine<int64_t, CountFunction>,
+	                      AggregateFunction::StateFinalize<int64_t, int64_t, CountFunction>,
+	                      FunctionNullHandling::SPECIAL_HANDLING, CountFunction::CountUpdate);
 	fun.name = "count";
-	fun.null_handling = FunctionNullHandling::SPECIAL_HANDLING;
 	fun.order_dependent = AggregateOrderDependent::NOT_ORDER_DEPENDENT;
 	return fun;
 }