diff --git a/src/duckdb/extension/parquet/parquet_reader.cpp b/src/duckdb/extension/parquet/parquet_reader.cpp index ea47aa207..70c88f097 100644 --- a/src/duckdb/extension/parquet/parquet_reader.cpp +++ b/src/duckdb/extension/parquet/parquet_reader.cpp @@ -1282,9 +1282,9 @@ ParquetReader::ParquetReader(ClientContext &context_p, OpenFileInfo file_p, Parq } } else { metadata = std::move(metadata_p); - if (parquet_options.encryption_config) { - encryption_util = context_p.db->GetEncryptionUtil(true); - } + } + if (parquet_options.encryption_config && !encryption_util) { + encryption_util = context_p.db->GetEncryptionUtil(true); } InitializeSchema(context_p); // Length-pushdown rewrites these columns to BIGINT, update the local schema diff --git a/src/duckdb/src/catalog/catalog_entry/duck_table_entry.cpp b/src/duckdb/src/catalog/catalog_entry/duck_table_entry.cpp index d6b0fb105..ebf418354 100644 --- a/src/duckdb/src/catalog/catalog_entry/duck_table_entry.cpp +++ b/src/duckdb/src/catalog/catalog_entry/duck_table_entry.cpp @@ -10,6 +10,7 @@ #include "duckdb/function/table/table_scan.hpp" #include "duckdb/main/database.hpp" #include "duckdb/parser/constraints/list.hpp" +#include "duckdb/parser/expression/constant_expression.hpp" #include "duckdb/parser/parsed_data/comment_on_column_info.hpp" #include "duckdb/parser/parsed_expression_iterator.hpp" #include "duckdb/planner/binder.hpp" @@ -333,7 +334,7 @@ unique_ptr DuckTableEntry::AlterEntry(ClientContext &context, Alte } case AlterTableType::ALTER_COLUMN_TYPE: { auto &change_type_info = table_info.Cast(); - return ChangeColumnType(context, change_type_info); + return ChangeColumnType(context, change_type_info, AlterTableType::ALTER_COLUMN_TYPE); } case AlterTableType::FOREIGN_KEY_CONSTRAINT: { auto &foreign_key_constraint_info = table_info.Cast(); @@ -391,6 +392,29 @@ static void RenameExpression(ParsedExpression &root_expr, RenameColumnInfo &info }); } +// Keep struct literal defaults aligned with nested field renames. +static unique_ptr RemapStructDefault(unique_ptr default_value, + const LogicalType &new_type, const Value &mapping) { + vector> children; + children.push_back(std::move(default_value)); + children.push_back(make_uniq(Value(new_type))); + children.push_back(make_uniq(mapping.Copy())); + children.push_back(make_uniq(Value())); + return make_uniq("remap_struct", std::move(children)); +} + +static const Value &GetRemapStructMapping(ChangeColumnTypeInfo &info) { + D_ASSERT(info.expression); + D_ASSERT(info.expression->GetExpressionClass() == ExpressionClass::FUNCTION); + auto &function = info.expression->Cast(); + D_ASSERT(function.FunctionName() == "remap_struct"); + auto &arguments = function.GetArguments(); + D_ASSERT(arguments.size() == 4); + auto &mapping = arguments[2].GetExpression(); + D_ASSERT(mapping.GetExpressionClass() == ExpressionClass::CONSTANT); + return mapping.Cast().GetValue(); +} + unique_ptr DuckTableEntry::RenameColumn(ClientContext &context, RenameColumnInfo &info) { auto rename_idx = GetColumnIndex(info.old_name); if (rename_idx.index == COLUMN_IDENTIFIER_ROW_ID) { @@ -674,7 +698,7 @@ unique_ptr DuckTableEntry::AddField(ClientContext &context, AddFie ChangeColumnTypeInfo change_column_type(info.GetAlterEntryData(), info.column_path[0], std::move(res.new_type), std::move(function)); - return ChangeColumnType(context, change_column_type); + return ChangeColumnType(context, change_column_type, AlterTableType::ADD_FIELD); } void DuckTableEntry::UpdateConstraintsOnColumnDrop(const LogicalIndex &removed_index, @@ -930,7 +954,7 @@ unique_ptr DuckTableEntry::RemoveField(ClientContext &context, Rem ChangeColumnTypeInfo change_column_type(info.GetAlterEntryData(), info.column_path[0], std::move(res.new_type), std::move(function)); - return ChangeColumnType(context, change_column_type); + return ChangeColumnType(context, change_column_type, AlterTableType::REMOVE_FIELD); } DroppedFieldMapping RenameFieldFromStruct(const LogicalType &type, const vector &column_path, @@ -1024,7 +1048,7 @@ unique_ptr DuckTableEntry::RenameField(ClientContext &context, Ren auto function = make_uniq("remap_struct", std::move(children)); ChangeColumnTypeInfo change_column_type(info.GetAlterEntryData(), info.column_path[0], std::move(res.new_type), std::move(function)); - return ChangeColumnType(context, change_column_type); + return ChangeColumnType(context, change_column_type, AlterTableType::RENAME_FIELD); } unique_ptr DuckTableEntry::SetDefault(ClientContext &context, SetDefaultInfo &info) { @@ -1112,7 +1136,8 @@ unique_ptr DuckTableEntry::DropNotNull(ClientContext &context, Dro return make_uniq(catalog, schema, *bound_create_info, storage, triggers); } -unique_ptr DuckTableEntry::ChangeColumnType(ClientContext &context, ChangeColumnTypeInfo &info) { +unique_ptr DuckTableEntry::ChangeColumnType(ClientContext &context, ChangeColumnTypeInfo &info, + AlterTableType alter_table_type) { // Bind type auto type_binder = Binder::CreateBinder(context); type_binder->SetSearchPath(catalog, schema.name); @@ -1148,6 +1173,10 @@ unique_ptr DuckTableEntry::ChangeColumnType(ClientContext &context throw NotImplementedException("Changing types of generated columns is not supported yet"); } copy.SetType(info.target_type); + if (alter_table_type == AlterTableType::RENAME_FIELD && copy.HasDefaultValue()) { + copy.SetDefaultValue( + RemapStructDefault(copy.DefaultValue().Copy(), info.target_type, GetRemapStructMapping(info))); + } } // TODO: check if the generated_expression breaks, only delete it if it does if (copy.Generated() && column_dependency_manager.IsDependencyOf(col.Logical(), change_idx)) { diff --git a/src/duckdb/src/catalog/default/default_functions.cpp b/src/duckdb/src/catalog/default/default_functions.cpp index 67ce7f3c0..1724b4abd 100644 --- a/src/duckdb/src/catalog/default/default_functions.cpp +++ b/src/duckdb/src/catalog/default/default_functions.cpp @@ -147,7 +147,7 @@ static const DefaultMacro internal_macros[] = { "(arr, sep := ',') AS case len(arr::varchar[]) when 0 then '' else list_aggr(arr::varchar[], 'string_agg', sep) " "end"}, - {DEFAULT_SCHEMA, "generate_subscripts", "(arr, dim) AS unnest(generate_series(1, array_length(arr, dim)))"}, + {DEFAULT_SCHEMA, "generate_subscripts", "(arr, dim := 1) AS unnest(generate_series(1, array_length(arr, dim)))"}, {DEFAULT_SCHEMA, "fdiv", "(x, y) AS floor(x/y)"}, {DEFAULT_SCHEMA, "fmod", "(x, y) AS (x-y*floor(x/y))"}, {DEFAULT_SCHEMA, "split_part", diff --git a/src/duckdb/src/common/enum_util.cpp b/src/duckdb/src/common/enum_util.cpp index 4aa8b5017..9b34b9455 100644 --- a/src/duckdb/src/common/enum_util.cpp +++ b/src/duckdb/src/common/enum_util.cpp @@ -4412,6 +4412,25 @@ PipelineBroadcastExchangeConsumerMode EnumUtil::FromString(StringUtil::StringToEnum(GetPipelineBroadcastExchangeConsumerModeValues(), 4, "PipelineBroadcastExchangeConsumerMode", value)); } +const StringUtil::EnumStringLiteral *GetPipelineBroadcastExchangeOrderModeValues() { + static constexpr StringUtil::EnumStringLiteral values[] { + { static_cast(PipelineBroadcastExchangeOrderMode::UNORDERED), "UNORDERED" }, + { static_cast(PipelineBroadcastExchangeOrderMode::SEQUENTIAL), "SEQUENTIAL" }, + { static_cast(PipelineBroadcastExchangeOrderMode::BATCH_INDEX), "BATCH_INDEX" } + }; + return values; +} + +template<> +const char* EnumUtil::ToChars(PipelineBroadcastExchangeOrderMode value) { + return StringUtil::EnumToString(GetPipelineBroadcastExchangeOrderModeValues(), 3, "PipelineBroadcastExchangeOrderMode", static_cast(value)); +} + +template<> +PipelineBroadcastExchangeOrderMode EnumUtil::FromString(const char *value) { + return static_cast(StringUtil::StringToEnum(GetPipelineBroadcastExchangeOrderModeValues(), 3, "PipelineBroadcastExchangeOrderMode", value)); +} + const StringUtil::EnumStringLiteral *GetPipelineInputModeValues() { static constexpr StringUtil::EnumStringLiteral values[] { { static_cast(PipelineInputMode::SCHEDULED_SOURCE), "SCHEDULED_SOURCE" }, diff --git a/src/duckdb/src/execution/index/fixed_size_allocator.cpp b/src/duckdb/src/execution/index/fixed_size_allocator.cpp index 5e9efbf7f..0dfc77d17 100644 --- a/src/duckdb/src/execution/index/fixed_size_allocator.cpp +++ b/src/duckdb/src/execution/index/fixed_size_allocator.cpp @@ -338,8 +338,14 @@ void FixedSizeAllocator::Init(const FixedSizeAllocatorInfo &info) { auto allocation_size = info.allocation_sizes[i]; // create the FixedSizeBuffer - buffers[buffer_id] = - make_uniq(block_manager, segment_count, allocation_size, buffer_block_pointer); + if (info.transient_block_handles) { + D_ASSERT(info.transient_block_handles->size() == info.buffer_ids.size()); + buffers[buffer_id] = make_uniq(block_manager, segment_count, allocation_size, + std::move((*info.transient_block_handles)[i])); + } else { + buffers[buffer_id] = + make_uniq(block_manager, segment_count, allocation_size, buffer_block_pointer); + } total_segment_count += segment_count; } diff --git a/src/duckdb/src/execution/index/fixed_size_buffer.cpp b/src/duckdb/src/execution/index/fixed_size_buffer.cpp index 1655d41e5..51f0470cc 100644 --- a/src/duckdb/src/execution/index/fixed_size_buffer.cpp +++ b/src/duckdb/src/execution/index/fixed_size_buffer.cpp @@ -47,6 +47,15 @@ FixedSizeBuffer::FixedSizeBuffer(BlockManager &block_manager, MemoryTag memory_t memset(buffer_handle.GetDataMutable(), 0, block_size); } +FixedSizeBuffer::FixedSizeBuffer(BlockManager &block_manager, const idx_t segment_count, const idx_t allocation_size, + shared_ptr block_handle_p) + : block_manager(block_manager), readers(0), segment_count(segment_count), allocation_size(allocation_size), + dirty(false), vacuum(false), loaded(false), block_pointer(), block_handle(std::move(block_handle_p)) { + D_ASSERT(block_handle); + buffer_handle = block_manager.buffer_manager.Pin(block_handle); + D_ASSERT(buffer_handle.IsValid()); +} + FixedSizeBuffer::FixedSizeBuffer(BlockManager &block_manager, const idx_t segment_count, const idx_t allocation_size, const BlockPointer &block_pointer) : block_manager(block_manager), readers(0), segment_count(segment_count), allocation_size(allocation_size), diff --git a/src/duckdb/src/execution/operator/persistent/physical_batch_copy_to_file.cpp b/src/duckdb/src/execution/operator/persistent/physical_batch_copy_to_file.cpp index 947d62690..6b94ad9b5 100644 --- a/src/duckdb/src/execution/operator/persistent/physical_batch_copy_to_file.cpp +++ b/src/duckdb/src/execution/operator/persistent/physical_batch_copy_to_file.cpp @@ -11,6 +11,7 @@ #include "duckdb/parallel/base_pipeline_event.hpp" #include "duckdb/parallel/executor_task.hpp" #include "duckdb/storage/buffer_manager.hpp" +#include "duckdb/storage/storage_info.hpp" #include "duckdb/logging/logger.hpp" #include "duckdb/logging/log_type.hpp" @@ -47,6 +48,10 @@ InsertionOrderPreservingMap PhysicalBatchCopyToFile::ParamsToString() co return result; } +OperatorPartitionInfo PhysicalBatchCopyToFile::RequiredPartitionInfo() const { + return OperatorPartitionInfo::BatchIndex(batch_size.IsValid() ? batch_size : optional_idx(DEFAULT_ROW_GROUP_SIZE)); +} + //===--------------------------------------------------------------------===// // States //===--------------------------------------------------------------------===// diff --git a/src/duckdb/src/execution/operator/persistent/physical_batch_insert.cpp b/src/duckdb/src/execution/operator/persistent/physical_batch_insert.cpp index b841ba268..43d8f4f30 100644 --- a/src/duckdb/src/execution/operator/persistent/physical_batch_insert.cpp +++ b/src/duckdb/src/execution/operator/persistent/physical_batch_insert.cpp @@ -9,6 +9,7 @@ #include "duckdb/storage/table/row_group_collection.hpp" #include "duckdb/storage/table/scan_state.hpp" #include "duckdb/storage/table_io_manager.hpp" +#include "duckdb/storage/storage_info.hpp" #include "duckdb/transaction/duck_transaction.hpp" #include "duckdb/transaction/local_storage.hpp" #include "duckdb/common/types/column/column_data_collection.hpp" @@ -19,16 +20,21 @@ PhysicalBatchInsert::PhysicalBatchInsert(PhysicalPlan &physical_plan, vector> bound_constraints_p, idx_t estimated_cardinality) : PhysicalOperator(physical_plan, PhysicalOperatorType::BATCH_INSERT, std::move(types_p), estimated_cardinality), - insert_table(&table), insert_types(table.GetTypes()), bound_constraints(std::move(bound_constraints_p)) { + insert_table(&table), insert_types(table.GetTypes()), bound_constraints(std::move(bound_constraints_p)), + preferred_batch_size(table.GetStorage().GetRowGroupSize()) { } PhysicalBatchInsert::PhysicalBatchInsert(PhysicalPlan &physical_plan, LogicalOperator &op, SchemaCatalogEntry &schema, unique_ptr info_p, idx_t estimated_cardinality) : PhysicalOperator(physical_plan, PhysicalOperatorType::BATCH_CREATE_TABLE_AS, op.types, estimated_cardinality), - insert_table(nullptr), schema(&schema), info(std::move(info_p)) { + insert_table(nullptr), schema(&schema), info(std::move(info_p)), preferred_batch_size(DEFAULT_ROW_GROUP_SIZE) { PhysicalInsert::GetInsertInfo(*info, insert_types); } +OperatorPartitionInfo PhysicalBatchInsert::RequiredPartitionInfo() const { + return OperatorPartitionInfo::BatchIndex(preferred_batch_size); +} + //===--------------------------------------------------------------------===// // CollectionMerger //===--------------------------------------------------------------------===// diff --git a/src/duckdb/src/execution/operator/scan/physical_column_data_scan.cpp b/src/duckdb/src/execution/operator/scan/physical_column_data_scan.cpp index 3cea47eaa..7694b300d 100644 --- a/src/duckdb/src/execution/operator/scan/physical_column_data_scan.cpp +++ b/src/duckdb/src/execution/operator/scan/physical_column_data_scan.cpp @@ -4,33 +4,79 @@ #include "duckdb/logging/logger.hpp" #include "duckdb/common/types/column/column_data_collection.hpp" +#include "duckdb/common/types/column/column_data_collection_segment.hpp" #include "duckdb/execution/operator/aggregate/physical_hash_aggregate.hpp" #include "duckdb/execution/operator/join/physical_delim_join.hpp" #include "duckdb/execution/operator/set/physical_cte.hpp" +#include "duckdb/main/client_config.hpp" #include "duckdb/parallel/meta_pipeline.hpp" #include "duckdb/parallel/pipeline.hpp" +#include "duckdb/parallel/task_scheduler.hpp" +#include "duckdb/storage/storage_info.hpp" namespace duckdb { +static idx_t ColumnDataScanBatchCount(idx_t count, idx_t batch_size) { + return MaxValue(count / batch_size + (count % batch_size != 0), 1); +} + +static idx_t GetColumnDataScanBatchSize(ClientContext &context, idx_t collection_count, + const OperatorPartitionInfo &partition_info) { + if (ClientConfig::GetConfig(context).verify_parallelism) { + return STANDARD_VECTOR_SIZE; + } + if (!partition_info.RequiresBatchIndex()) { + return STANDARD_VECTOR_SIZE; + } + idx_t preferred_batch_size = DEFAULT_ROW_GROUP_SIZE; + if (partition_info.preferred_batch_size.IsValid() && partition_info.preferred_batch_size.GetIndex() > 0) { + preferred_batch_size = partition_info.preferred_batch_size.GetIndex(); + } + auto thread_count = MaxValue(TaskScheduler::GetScheduler(context).NumberOfThreads(), 1); + auto rows_per_thread = collection_count / thread_count + (collection_count % thread_count != 0); + auto parallelism_cap = MaxValue(STANDARD_VECTOR_SIZE, rows_per_thread); + return MinValue(preferred_batch_size, parallelism_cap); +} + +static vector GenerateColumnDataColumnIds(const vector &types) { + vector column_ids; + column_ids.reserve(types.size()); + for (idx_t i = 0; i < types.size(); i++) { + column_ids.push_back(i); + } + return column_ids; +} + PhysicalColumnDataScan::PhysicalColumnDataScan(PhysicalPlan &physical_plan, vector types, PhysicalOperatorType op_type, idx_t estimated_cardinality, optionally_owned_ptr collection_p) : PhysicalOperator(physical_plan, op_type, std::move(types), estimated_cardinality), - collection(std::move(collection_p)) { + collection(std::move(collection_p)), column_ids(GenerateColumnDataColumnIds(this->types)) { +} + +PhysicalColumnDataScan::PhysicalColumnDataScan(PhysicalPlan &physical_plan, vector types, + PhysicalOperatorType op_type, idx_t estimated_cardinality, + optionally_owned_ptr collection_p, + vector column_ids_p) + : PhysicalOperator(physical_plan, op_type, std::move(types), estimated_cardinality), + collection(std::move(collection_p)), column_ids(std::move(column_ids_p)) { + D_ASSERT(this->types.size() == column_ids.size()); } PhysicalColumnDataScan::PhysicalColumnDataScan(PhysicalPlan &physical_plan, vector types, PhysicalOperatorType op_type, idx_t estimated_cardinality, TableIndex cte_index) : PhysicalOperator(physical_plan, op_type, std::move(types), estimated_cardinality), collection(nullptr), - cte_index(cte_index) { + column_ids(GenerateColumnDataColumnIds(this->types)), cte_index(cte_index) { } class PhysicalColumnDataGlobalScanState : public GlobalSourceState { public: - explicit PhysicalColumnDataGlobalScanState(const ColumnDataCollection &collection) - : max_threads(MaxValue(collection.ChunkCount(), 1)) { - collection.InitializeScan(global_scan_state); + PhysicalColumnDataGlobalScanState(ClientContext &context, const ColumnDataCollection &collection, + const vector &column_ids, const OperatorPartitionInfo &partition_info) + : batch_size(GetColumnDataScanBatchSize(context, collection.Count(), partition_info)), + max_threads(ColumnDataScanBatchCount(collection.Count(), batch_size)) { + collection.InitializeScan(global_scan_state, column_ids); } idx_t MaxThreads() override { @@ -40,19 +86,68 @@ class PhysicalColumnDataGlobalScanState : public GlobalSourceState { public: ColumnDataParallelScanState global_scan_state; + const idx_t batch_size; const idx_t max_threads; + idx_t next_batch_index = 0; +}; + +struct ColumnDataScanEntry { + ColumnDataScanEntry(idx_t chunk_index_p, idx_t segment_index_p, idx_t row_index_p) + : chunk_index(chunk_index_p), segment_index(segment_index_p), row_index(row_index_p) { + } + + idx_t chunk_index; + idx_t segment_index; + idx_t row_index; }; class PhysicalColumnDataLocalScanState : public LocalSourceState { public: + bool AssignTask(const ColumnDataCollection &collection, PhysicalColumnDataGlobalScanState &gstate) { + entries.clear(); + entry_index = 0; + batch_index = DConstants::INVALID_INDEX; + + idx_t task_count = 0; + lock_guard l(gstate.global_scan_state.lock); + while (task_count < gstate.batch_size) { + idx_t chunk_index; + idx_t segment_index; + idx_t row_index; + if (!collection.NextScanIndex(gstate.global_scan_state.scan_state, chunk_index, segment_index, row_index)) { + break; + } + entries.emplace_back(chunk_index, segment_index, row_index); + task_count += collection.GetSegments()[segment_index]->chunk_data[chunk_index].count; + } + + if (entries.empty()) { + return false; + } + batch_index = gstate.next_batch_index++; + return true; + } + ColumnDataLocalScanState local_scan_state; + vector entries; + idx_t entry_index = 0; + idx_t batch_index = DConstants::INVALID_INDEX; }; unique_ptr PhysicalColumnDataScan::GetGlobalSourceState(ClientContext &context) const { if (!collection) { return make_uniq(); } - return make_uniq(*collection); + return GetGlobalSourceState(context, OperatorPartitionInfo::NoPartitionInfo()); +} + +unique_ptr +PhysicalColumnDataScan::GetGlobalSourceState(ClientContext &context, + const OperatorPartitionInfo &partition_info) const { + if (!collection) { + return make_uniq(); + } + return make_uniq(context, *collection, column_ids, partition_info); } unique_ptr PhysicalColumnDataScan::GetLocalSourceState(ExecutionContext &, @@ -64,8 +159,42 @@ SourceResultType PhysicalColumnDataScan::GetDataInternal(ExecutionContext &conte OperatorSourceInput &input) const { auto &gstate = input.global_state.Cast(); auto &lstate = input.local_state.Cast(); - collection->Scan(gstate.global_scan_state, lstate.local_scan_state, chunk); - return chunk.size() == 0 ? SourceResultType::FINISHED : SourceResultType::HAVE_MORE_OUTPUT; + if (lstate.entry_index >= lstate.entries.size() && !lstate.AssignTask(*collection, gstate)) { + chunk.Reset(); + return SourceResultType::FINISHED; + } + + auto &entry = lstate.entries[lstate.entry_index++]; + if (column_ids.empty()) { + chunk.Reset(); + auto &chunk_data = collection->GetSegments()[entry.segment_index]->chunk_data[entry.chunk_index]; + chunk.SetChildCardinality(chunk_data.count); + } else { + collection->ScanAtIndex(gstate.global_scan_state, lstate.local_scan_state, chunk, entry.chunk_index, + entry.segment_index, entry.row_index); + } + return SourceResultType::HAVE_MORE_OUTPUT; +} + +bool PhysicalColumnDataScan::SupportsPartitioning(const OperatorPartitionInfo &partition_info) const { + if (partition_info.RequiresPartitionColumns()) { + return false; + } + if (!partition_info.RequiresBatchIndex()) { + return false; + } + if (type == PhysicalOperatorType::CTE_SCAN && cte_source) { + return cte_source->SupportsPartitioning(partition_info); + } + return type == PhysicalOperatorType::COLUMN_DATA_SCAN || type == PhysicalOperatorType::CTE_SCAN; +} + +OperatorPartitionData PhysicalColumnDataScan::GetPartitionData(ExecutionContext &context, DataChunk &chunk, + GlobalSourceState &gstate, LocalSourceState &lstate_p, + const OperatorPartitionInfo &partition_info) const { + D_ASSERT(SupportsPartitioning(partition_info)); + auto &lstate = lstate_p.Cast(); + return OperatorPartitionData(lstate.batch_index); } ProgressData PhysicalColumnDataScan::GetProgress(ClientContext &context, GlobalSourceState &gstate) const { @@ -128,7 +257,7 @@ void PhysicalColumnDataScan::BuildPipelines(Pipeline ¤t, MetaPipeline &met return; } if (cte.ShouldUseBufferedConsumer(current)) { - cte.RegisterBufferedConsumer(source.consumer_idx); + cte.RegisterBufferedConsumer(current, source.consumer_idx); current.AddDataflowDependency(cte_dependency); DUCKDB_LOG(current.GetClientContext(), PhysicalOperatorLogType, cte, "PhysicalCTE", "SelectConsumer", {{"consumer", to_string(source.consumer_idx)}, {"mode", "BUFFERED"}}); diff --git a/src/duckdb/src/execution/operator/set/physical_cte.cpp b/src/duckdb/src/execution/operator/set/physical_cte.cpp index 40669ed46..99145c00c 100644 --- a/src/duckdb/src/execution/operator/set/physical_cte.cpp +++ b/src/duckdb/src/execution/operator/set/physical_cte.cpp @@ -2,6 +2,7 @@ #include "duckdb/common/atomic.hpp" #include "duckdb/common/enum_util.hpp" +#include "duckdb/common/types/batched_data_collection.hpp" #include "duckdb/common/types/column/column_data_collection.hpp" #include "duckdb/logging/log_type.hpp" #include "duckdb/logging/logger.hpp" @@ -17,7 +18,8 @@ enum class CTECombineState : uint8_t { PENDING, COMBINED }; class CTEConsumerGlobalSourceState : public GlobalSourceState { public: CTEConsumerGlobalSourceState(shared_ptr exchange_p, idx_t consumer_idx_p) - : exchange(std::move(exchange_p)), consumer_idx(consumer_idx_p) { + : exchange(std::move(exchange_p)), consumer_idx(consumer_idx_p), + max_threads(exchange->PreservesOrder() ? 1 : exchange->MaxThreads()) { } ~CTEConsumerGlobalSourceState() override { @@ -25,7 +27,7 @@ class CTEConsumerGlobalSourceState : public GlobalSourceState { } idx_t MaxThreads() override { - return exchange->MaxThreads(); + return max_threads; } void Unregister() { @@ -37,12 +39,18 @@ class CTEConsumerGlobalSourceState : public GlobalSourceState { shared_ptr exchange; idx_t consumer_idx; + idx_t max_threads; atomic unregistered {false}; }; class CTEConsumerLocalSourceState : public LocalSourceState { public: - shared_ptr current_chunk; + explicit CTEConsumerLocalSourceState(shared_ptr scan_state_p) + : scan_state(std::move(scan_state_p)) { + } + + shared_ptr scan_state; + optional_idx exchange_batch_index; }; PhysicalCTEConsumerSource::PhysicalCTEConsumerSource(PhysicalPlan &physical_plan, vector types, @@ -53,19 +61,48 @@ PhysicalCTEConsumerSource::PhysicalCTEConsumerSource(PhysicalPlan &physical_plan } unique_ptr PhysicalCTEConsumerSource::GetGlobalSourceState(ClientContext &context) const { + return GetGlobalSourceState(context, OperatorPartitionInfo::NoPartitionInfo()); +} + +unique_ptr +PhysicalCTEConsumerSource::GetGlobalSourceState(ClientContext &context, + const OperatorPartitionInfo &partition_info) const { return make_uniq(exchange, consumer_idx); } unique_ptr PhysicalCTEConsumerSource::GetLocalSourceState(ExecutionContext &context, GlobalSourceState &gstate) const { - return make_uniq(); + return make_uniq(exchange->GetScanState()); } SourceResultType PhysicalCTEConsumerSource::GetDataInternal(ExecutionContext &context, DataChunk &chunk, OperatorSourceInput &input) const { auto &gstate = input.global_state.Cast(); auto &lstate = input.local_state.Cast(); - return gstate.exchange->Scan(gstate.consumer_idx, chunk, lstate.current_chunk, input.interrupt_state); + return gstate.exchange->Scan(gstate.consumer_idx, chunk, *lstate.scan_state, lstate.exchange_batch_index, + input.interrupt_state); +} + +OperatorPartitionData PhysicalCTEConsumerSource::GetPartitionData(ExecutionContext &context, DataChunk &chunk, + GlobalSourceState &gstate, LocalSourceState &lstate_p, + const OperatorPartitionInfo &partition_info) const { + D_ASSERT(SupportsPartitioning(partition_info)); + auto &lstate = lstate_p.Cast(); + D_ASSERT(lstate.exchange_batch_index.IsValid()); + return OperatorPartitionData(lstate.exchange_batch_index.GetIndex()); +} + +bool PhysicalCTEConsumerSource::SupportsPartitioning(const OperatorPartitionInfo &partition_info) const { + return exchange->SupportsBatchIndex() && partition_info.RequiresBatchIndex() && + !partition_info.RequiresPartitionColumns(); +} + +OrderPreservationType PhysicalCTEConsumerSource::SourceOrder() const { + return exchange->SourceOrder(); +} + +bool PhysicalCTEConsumerSource::ParallelSource() const { + return exchange->OrderMode() != PipelineBroadcastExchangeOrderMode::SEQUENTIAL; } ProgressData PhysicalCTEConsumerSource::GetProgress(ClientContext &context, GlobalSourceState &gstate) const { @@ -118,6 +155,7 @@ class CTEGlobalState : public GlobalSinkState { optional_ptr working_table_ref; shared_ptr exchange; CTEExecutionMode execution_mode; + unique_ptr ordered_data; annotated_mutex lhs_lock; @@ -130,8 +168,15 @@ class CTEGlobalState : public GlobalSinkState { D_ASSERT(op.working_table); op.working_table->Reset(); working_table_ref = op.working_table.get(); + if (op.use_batch_index) { + ordered_data = make_uniq(context, op.working_table->Types(), + op.working_table->GetAllocatorType()); + } else { + ordered_data.reset(); + } } else { working_table_ref = nullptr; + ordered_data.reset(); } GlobalSinkState::Reset(context); } @@ -149,26 +194,52 @@ class CTEGlobalState : public GlobalSinkState { annotated_lock_guard guard(lhs_lock); working_table_ref->Combine(input); } + + void MergeBatches(BatchedDataCollection &input) { + annotated_lock_guard guard(lhs_lock); + D_ASSERT(ordered_data); + ordered_data->Merge(input); + } + + void FinalizeBatches() { + annotated_lock_guard guard(lhs_lock); + if (!ordered_data) { + return; + } + D_ASSERT(working_table_ref); + auto collection = ordered_data->FetchCollection(); + working_table_ref->Combine(*collection); + ordered_data.reset(); + } }; class CTELocalState : public LocalSinkState { public: - explicit CTELocalState(ClientContext &context, const PhysicalCTE &op) : execution_mode(op.GetExecutionMode()) { + explicit CTELocalState(ExecutionContext &context, const PhysicalCTE &op) + : execution_mode(op.GetExecutionMode()), use_batch_index(op.use_batch_index) { if (execution_mode != CTEExecutionMode::STREAMING_FANOUT) { D_ASSERT(op.working_table); - lhs_data = make_uniq(context, op.working_table->Types()); - lhs_data->InitializeAppend(append_state); + if (use_batch_index) { + lhs_batches = make_uniq(context.client, op.working_table->Types(), + op.working_table->GetAllocatorType()); + } else { + lhs_data = make_uniq(context.client, op.working_table->Types()); + lhs_data->InitializeAppend(append_state); + } } if (execution_mode != CTEExecutionMode::MATERIALIZED) { D_ASSERT(op.exchange); - exchange_state = op.exchange->GetLocalState(context); + D_ASSERT(context.pipeline); + exchange_state = op.exchange->GetLocalState(context.client, context.pipeline->GetBaseBatchIndex()); } } unique_ptr distinct_state; unique_ptr lhs_data; + unique_ptr lhs_batches; ColumnDataAppendState append_state; CTEExecutionMode execution_mode; + bool use_batch_index; CTESinkExecutionState sink_execution_state = CTESinkExecutionState::ACTIVE; // Combine can be retried after a blocked exchange finish. CTECombineState combine_state = CTECombineState::PENDING; @@ -176,8 +247,14 @@ class CTELocalState : public LocalSinkState { void Append(DataChunk &input) { D_ASSERT(execution_mode != CTEExecutionMode::STREAMING_FANOUT); - D_ASSERT(lhs_data); - lhs_data->Append(append_state, input); + if (use_batch_index) { + D_ASSERT(lhs_batches); + D_ASSERT(partition_info.batch_index.IsValid()); + lhs_batches->Append(input, partition_info.batch_index.GetIndex()); + } else { + D_ASSERT(lhs_data); + lhs_data->Append(append_state, input); + } } }; @@ -186,7 +263,7 @@ unique_ptr PhysicalCTE::GetGlobalSinkState(ClientContext &conte } unique_ptr PhysicalCTE::GetLocalSinkState(ExecutionContext &context) const { - auto state = make_uniq(context.client, *this); + auto state = make_uniq(context, *this); return std::move(state); } @@ -197,7 +274,7 @@ SinkResultType PhysicalCTE::Sink(ExecutionContext &context, DataChunk &chunk, Op if (lstate.execution_mode != CTEExecutionMode::MATERIALIZED) { D_ASSERT(gstate.exchange); D_ASSERT(lstate.exchange_state); - result = gstate.exchange->Push(chunk, *lstate.exchange_state, input.interrupt_state); + result = gstate.exchange->Push(chunk, *lstate.exchange_state, lstate.partition_info, input.interrupt_state); if (result == SinkResultType::BLOCKED) { if (lstate.sink_execution_state != CTESinkExecutionState::BLOCKED) { DUCKDB_LOG(context.client, PhysicalOperatorLogType, *this, "PhysicalCTE", "ExchangeBlocked", @@ -223,11 +300,22 @@ SinkResultType PhysicalCTE::Sink(ExecutionContext &context, DataChunk &chunk, Op return result; } +SinkNextBatchType PhysicalCTE::NextBatch(ExecutionContext &context, OperatorSinkNextBatchInput &input) const { + auto &gstate = input.global_state.Cast(); + auto &lstate = input.local_state.Cast(); + if (lstate.execution_mode == CTEExecutionMode::MATERIALIZED) { + return SinkNextBatchType::READY; + } + D_ASSERT(gstate.exchange); + D_ASSERT(lstate.exchange_state); + return gstate.exchange->NextBatch(*lstate.exchange_state, lstate.partition_info, input.interrupt_state); +} + SinkCombineResultType PhysicalCTE::Combine(ExecutionContext &context, OperatorSinkCombineInput &input) const { auto &lstate = input.local_state.Cast(); if (lstate.execution_mode != CTEExecutionMode::MATERIALIZED) { D_ASSERT(lstate.exchange_state); - auto result = exchange->FinishLocal(*lstate.exchange_state, input.interrupt_state); + auto result = exchange->FinishLocal(*lstate.exchange_state, lstate.partition_info, input.interrupt_state); if (result == SinkCombineResultType::BLOCKED) { return result; } @@ -235,8 +323,13 @@ SinkCombineResultType PhysicalCTE::Combine(ExecutionContext &context, OperatorSi if (lstate.execution_mode != CTEExecutionMode::STREAMING_FANOUT && lstate.combine_state == CTECombineState::PENDING) { auto &gstate = input.global_state.Cast(); - D_ASSERT(lstate.lhs_data); - gstate.MergeIT(*lstate.lhs_data); + if (lstate.use_batch_index) { + D_ASSERT(lstate.lhs_batches); + gstate.MergeBatches(*lstate.lhs_batches); + } else { + D_ASSERT(lstate.lhs_data); + gstate.MergeIT(*lstate.lhs_data); + } lstate.combine_state = CTECombineState::COMBINED; } @@ -245,8 +338,9 @@ SinkCombineResultType PhysicalCTE::Combine(ExecutionContext &context, OperatorSi SinkFinalizeType PhysicalCTE::Finalize(Pipeline &pipeline, Event &event, ClientContext &context, OperatorSinkFinalizeInput &input) const { + auto &gstate = input.global_state.Cast(); + gstate.FinalizeBatches(); if (exchange && UseStreamingExchange()) { - auto &gstate = input.global_state.Cast(); gstate.exchange->Finish(); gstate.exchange->FinishDirectConsumers(); DUCKDB_LOG(context, PhysicalOperatorLogType, *this, "PhysicalCTE", "ProducerFinished", @@ -261,6 +355,8 @@ SinkFinalizeType PhysicalCTE::Finalize(Pipeline &pipeline, Event &event, ClientC void PhysicalCTE::BuildPipelines(Pipeline ¤t, MetaPipeline &meta_pipeline) { D_ASSERT(children.size() == 2); pipeline_selection_state = CTEPipelineSelectionState::UNRESOLVED; + preferred_batch_size = optional_idx(); + conflicting_batch_sizes = false; op_state.reset(); sink_state.reset(); if (exchange) { @@ -269,12 +365,22 @@ void PhysicalCTE::BuildPipelines(Pipeline ¤t, MetaPipeline &meta_pipeline) exchange->ResetConsumerRegistrations(); } + if (meta_pipeline.HasRecursiveCTE() && use_batch_index) { + use_batch_index = false; + parallel = false; + } + auto &state = meta_pipeline.GetState(); auto &child_meta_pipeline = meta_pipeline.CreateChildMetaPipeline( current, *this, MetaPipelineType::REGULAR, exchange ? MetaPipelineDependencyMode::NO_DEPENDENCY : MetaPipelineDependencyMode::ADD_DEPENDENCY); child_meta_pipeline.Build(children[0]); + if (exchange) { + vector> producer_pipelines; + child_meta_pipeline.GetPipelines(producer_pipelines, false); + exchange->SetProducerPipelines(producer_pipelines); + } for (auto &cte_scan : cte_scans) { state.cte_dependencies.insert(make_pair(cte_scan, reference(*child_meta_pipeline.GetBasePipeline()))); @@ -308,7 +414,7 @@ void PhysicalCTE::BuildPipelines(Pipeline ¤t, MetaPipeline &meta_pipeline) } if (last_child_ptr) { meta_pipeline.AddRecursiveDependencies(side_effect_pipelines, *last_child_ptr, RecursiveDependencyMode::FORCE, - exchange ? DataflowDependencyMode::SKIP + exchange ? DataflowDependencyMode::SKIP_CONFLICTING : DataflowDependencyMode::INCLUDE); } } @@ -317,7 +423,11 @@ bool PhysicalCTE::TryRegisterDirectConsumer(Pipeline &pipeline, idx_t consumer_i if (!exchange) { return false; } - return exchange->TryRegisterDirectConsumer(pipeline, consumer_idx); + if (!exchange->TryRegisterDirectConsumer(pipeline, consumer_idx)) { + return false; + } + RegisterBatchPreference(pipeline); + return true; } bool PhysicalCTE::ShouldUseBufferedConsumer(Pipeline &pipeline) const { @@ -332,9 +442,10 @@ bool PhysicalCTE::ShouldUseBufferedConsumer(Pipeline &pipeline) const { return pipeline.CanStopSourceEarly(); } -void PhysicalCTE::RegisterBufferedConsumer(idx_t consumer_idx) { +void PhysicalCTE::RegisterBufferedConsumer(Pipeline &pipeline, idx_t consumer_idx) { D_ASSERT(exchange); exchange->SelectBufferedConsumer(consumer_idx); + RegisterBatchPreference(pipeline); } void PhysicalCTE::RegisterMaterializedConsumer(idx_t consumer_idx) { @@ -357,6 +468,25 @@ CTEExecutionMode PhysicalCTE::GetExecutionMode() const { return CTEExecutionMode::STREAMING_FANOUT; } +void PhysicalCTE::RegisterBatchPreference(Pipeline &pipeline) { + auto sink = pipeline.GetSink(); + if (!sink) { + return; + } + auto partition_info = sink->RequiredPartitionInfo(); + if (!partition_info.RequiresBatchIndex() || !partition_info.preferred_batch_size.IsValid() || + conflicting_batch_sizes) { + return; + } + auto batch_size = partition_info.preferred_batch_size.GetIndex(); + if (!preferred_batch_size.IsValid()) { + preferred_batch_size = batch_size; + } else if (preferred_batch_size.GetIndex() != batch_size) { + preferred_batch_size = optional_idx(); + conflicting_batch_sizes = true; + } +} + bool PhysicalCTE::UseStreamingExchange() const { return GetExecutionMode() != CTEExecutionMode::MATERIALIZED; } @@ -402,6 +532,10 @@ InsertionOrderPreservingMap PhysicalCTE::ParamsToString() const { } if (exchange) { auto summary = exchange->GetConsumerSummary(); + result["Order Mode"] = EnumUtil::ToString(exchange->OrderMode()); + if (preferred_batch_size.IsValid()) { + result["Preferred Batch Size"] = StringUtil::Format("%llu", preferred_batch_size.GetIndex()); + } if (summary.materialized > 0) { result["Materialized Consumers"] = StringUtil::Format("%llu", summary.materialized); } @@ -420,8 +554,7 @@ ProgressData PhysicalCTE::GetSinkProgress(ClientContext &context, GlobalSinkStat if (!state.working_table_ref) { return ProgressData {0, 1, true}; } - auto &working_table = *state.working_table_ref; - auto count = double(working_table.Count()); + auto count = double(state.ordered_data ? state.ordered_data->Count() : state.working_table_ref->Count()); ProgressData progress; progress.done = count; progress.total = count + source_progress.total; diff --git a/src/duckdb/src/execution/operator/set/physical_recursive_cte_runtime.cpp b/src/duckdb/src/execution/operator/set/physical_recursive_cte_runtime.cpp index b0fcd0764..efe3ad8e4 100644 --- a/src/duckdb/src/execution/operator/set/physical_recursive_cte_runtime.cpp +++ b/src/duckdb/src/execution/operator/set/physical_recursive_cte_runtime.cpp @@ -170,7 +170,7 @@ class RecursiveCTETask : public ExecutorTask { } bool TaskBlockedOnResult() const override { - return pipeline_executor.RemainingSinkChunk(); + return pipeline.IsStreamingResultPipeline() && pipeline_executor.RemainingSinkChunk(); } private: diff --git a/src/duckdb/src/execution/physical_operator.cpp b/src/duckdb/src/execution/physical_operator.cpp index 63c62417a..6f606322b 100644 --- a/src/duckdb/src/execution/physical_operator.cpp +++ b/src/duckdb/src/execution/physical_operator.cpp @@ -133,6 +133,12 @@ unique_ptr PhysicalOperator::GetGlobalSourceState(ClientConte return make_uniq(); } +unique_ptr +PhysicalOperator::GetGlobalSourceState(ClientContext &context, const OperatorPartitionInfo &partition_info) const { + (void)partition_info; + return GetGlobalSourceState(context); +} + // LCOV_EXCL_START SourceResultType PhysicalOperator::GetData(ExecutionContext &context, DataChunk &chunk, OperatorSourceInput &input) const { diff --git a/src/duckdb/src/execution/physical_plan/plan_column_data_get.cpp b/src/duckdb/src/execution/physical_plan/plan_column_data_get.cpp index 6712ab5a1..501843bbf 100644 --- a/src/duckdb/src/execution/physical_plan/plan_column_data_get.cpp +++ b/src/duckdb/src/execution/physical_plan/plan_column_data_get.cpp @@ -7,8 +7,9 @@ namespace duckdb { PhysicalOperator &PhysicalPlanGenerator::CreatePlan(LogicalColumnDataGet &op) { D_ASSERT(op.children.empty()); D_ASSERT(op.collection); + op.ResolveOperatorTypes(); return Make(op.types, PhysicalOperatorType::COLUMN_DATA_SCAN, op.estimated_cardinality, - std::move(op.collection)); + std::move(op.collection), op.GetColumnIds()); } } // namespace duckdb diff --git a/src/duckdb/src/execution/physical_plan/plan_cte.cpp b/src/duckdb/src/execution/physical_plan/plan_cte.cpp index 82e39ae8f..ec270e45a 100644 --- a/src/duckdb/src/execution/physical_plan/plan_cte.cpp +++ b/src/duckdb/src/execution/physical_plan/plan_cte.cpp @@ -28,23 +28,29 @@ PhysicalOperator &PhysicalPlanGenerator::CreatePlan(LogicalMaterializedCTE &op) // Create the working_table that the PhysicalCTE will use for evaluation. auto working_table = make_shared_ptr(context, op.children[0]->types); + + // Add the ColumnDataCollection to the context of this PhysicalPlanGenerator + recursive_cte_tables[op.table_index] = working_table; + materialized_ctes[op.table_index] = vector>(); + + // Create the plan for the left side. This is the materialization. + auto &left = CreatePlan(*op.children[0]); + const auto cte_body_order = OrderPreservationRecursive(left); + const auto preserve_cte_order = PreserveInsertionOrder(left); + const auto use_batch_index = preserve_cte_order && UseBatchIndex(left); + const auto source_order = preserve_cte_order ? cte_body_order : OrderPreservationType::NO_ORDER; + materialized_cte_orders[op.table_index] = source_order; + shared_ptr exchange; if (use_exchange) { auto completion_mode = cte_body_has_side_effects ? PipelineBroadcastExchangeCompletionMode::RUN_TO_COMPLETION : PipelineBroadcastExchangeCompletionMode::STOP_WHEN_UNCONSUMED; - exchange = make_shared_ptr(context, op.children[0]->types, completion_mode); - } - - // Add the ColumnDataCollection to the context of this PhysicalPlanGenerator - recursive_cte_tables[op.table_index] = working_table; - if (exchange) { + exchange = make_shared_ptr(context, op.children[0]->types, completion_mode, + source_order, use_batch_index); materialized_cte_exchanges[op.table_index] = exchange; } - materialized_ctes[op.table_index] = vector>(); - // Create the plan for the left side. This is the materialization. - auto &left = CreatePlan(*op.children[0]); // Initialize an empty vector to collect the scan operators. auto &right = CreatePlan(*op.children[1]); @@ -54,6 +60,9 @@ PhysicalOperator &PhysicalPlanGenerator::CreatePlan(LogicalMaterializedCTE &op) cast_cte.exchange = exchange; cast_cte.cte_scans = materialized_ctes[op.table_index]; cast_cte.cte_body_has_side_effects = cte_body_has_side_effects; + cast_cte.preserve_order = preserve_cte_order; + cast_cte.use_batch_index = use_batch_index; + cast_cte.parallel = !preserve_cte_order || use_batch_index; return cte; } diff --git a/src/duckdb/src/execution/physical_plan/plan_recursive_cte.cpp b/src/duckdb/src/execution/physical_plan/plan_recursive_cte.cpp index 5df9212a7..a46fc1ab7 100644 --- a/src/duckdb/src/execution/physical_plan/plan_recursive_cte.cpp +++ b/src/duckdb/src/execution/physical_plan/plan_recursive_cte.cpp @@ -113,25 +113,25 @@ PhysicalOperator &PhysicalPlanGenerator::CreatePlan(LogicalCTERef &op) { auto &chunk_scan = Make(op.chunk_types, PhysicalOperatorType::CTE_SCAN, op.estimated_cardinality, op.cte_index); + auto cte = recursive_cte_tables.find(op.cte_index); + if (cte == recursive_cte_tables.end()) { + throw InvalidInputException("Referenced materialized CTE does not exist."); + } + auto &cast_chunk_scan = chunk_scan.Cast(); + cast_chunk_scan.collection = cte->second.get(); + auto cte_order = materialized_cte_orders.find(op.cte_index); + if (cte_order != materialized_cte_orders.end()) { + cast_chunk_scan.source_order = cte_order->second; + } + auto exchange = materialized_cte_exchanges.find(op.cte_index); if (exchange != materialized_cte_exchanges.end()) { - auto cte = recursive_cte_tables.find(op.cte_index); - if (cte == recursive_cte_tables.end()) { - throw InvalidInputException("Referenced materialized CTE does not exist."); - } // Exchange consumers can still be converted to materialized scans during pipeline construction. - cast_chunk_scan.collection = cte->second.get(); auto consumer_idx = exchange->second->RegisterConsumer(); auto &source = Make(op.chunk_types, op.estimated_cardinality, op.cte_index, exchange->second, consumer_idx); cast_chunk_scan.cte_source = source; - } else { - auto cte = recursive_cte_tables.find(op.cte_index); - if (cte == recursive_cte_tables.end()) { - throw InvalidInputException("Referenced materialized CTE does not exist."); - } - cast_chunk_scan.collection = cte->second.get(); } materialized_cte->second.push_back(cast_chunk_scan); return chunk_scan; diff --git a/src/duckdb/src/function/table/version/pragma_version.cpp b/src/duckdb/src/function/table/version/pragma_version.cpp index 34d2e77c4..def464979 100644 --- a/src/duckdb/src/function/table/version/pragma_version.cpp +++ b/src/duckdb/src/function/table/version/pragma_version.cpp @@ -1,5 +1,5 @@ #ifndef DUCKDB_PATCH_VERSION -#define DUCKDB_PATCH_VERSION "0-alpha36121" +#define DUCKDB_PATCH_VERSION "0-alpha36155" #endif #ifndef DUCKDB_MINOR_VERSION #define DUCKDB_MINOR_VERSION 0 @@ -8,10 +8,10 @@ #define DUCKDB_MAJOR_VERSION 2 #endif #ifndef DUCKDB_VERSION -#define DUCKDB_VERSION "v2.0.0-alpha36121" +#define DUCKDB_VERSION "v2.0.0-alpha36155" #endif #ifndef DUCKDB_SOURCE_ID -#define DUCKDB_SOURCE_ID "d54e4f6a20" +#define DUCKDB_SOURCE_ID "76361ce4fb" #endif #include "duckdb/function/table/system_functions.hpp" #include "duckdb/main/database.hpp" diff --git a/src/duckdb/src/include/duckdb/catalog/catalog_entry/duck_table_entry.hpp b/src/duckdb/src/include/duckdb/catalog/catalog_entry/duck_table_entry.hpp index 0ebc3f355..d17d42bae 100644 --- a/src/duckdb/src/include/duckdb/catalog/catalog_entry/duck_table_entry.hpp +++ b/src/duckdb/src/include/duckdb/catalog/catalog_entry/duck_table_entry.hpp @@ -85,7 +85,8 @@ class DuckTableEntry : public TableCatalogEntry { unique_ptr RemoveColumn(ClientContext &context, RemoveColumnInfo &info); unique_ptr RemoveField(ClientContext &context, RemoveFieldInfo &info); unique_ptr SetDefault(ClientContext &context, SetDefaultInfo &info); - unique_ptr ChangeColumnType(ClientContext &context, ChangeColumnTypeInfo &info); + unique_ptr ChangeColumnType(ClientContext &context, ChangeColumnTypeInfo &info, + AlterTableType alter_table_type); unique_ptr SetNotNull(ClientContext &context, SetNotNullInfo &info); unique_ptr DropNotNull(ClientContext &context, DropNotNullInfo &info); unique_ptr AddForeignKeyConstraint(AlterForeignKeyInfo &info); diff --git a/src/duckdb/src/include/duckdb/common/enum_util.hpp b/src/duckdb/src/include/duckdb/common/enum_util.hpp index 586fd5219..5ba1e0383 100644 --- a/src/duckdb/src/include/duckdb/common/enum_util.hpp +++ b/src/duckdb/src/include/duckdb/common/enum_util.hpp @@ -396,6 +396,8 @@ enum class PhysicalType : uint8_t; enum class PipelineBroadcastExchangeConsumerMode : uint8_t; +enum class PipelineBroadcastExchangeOrderMode : uint8_t; + enum class PipelineInputMode : uint8_t; enum class PragmaType : uint8_t; @@ -1141,6 +1143,9 @@ const char* EnumUtil::ToChars(PhysicalType value); template<> const char* EnumUtil::ToChars(PipelineBroadcastExchangeConsumerMode value); +template<> +const char* EnumUtil::ToChars(PipelineBroadcastExchangeOrderMode value); + template<> const char* EnumUtil::ToChars(PipelineInputMode value); @@ -1985,6 +1990,9 @@ PhysicalType EnumUtil::FromString(const char *value); template<> PipelineBroadcastExchangeConsumerMode EnumUtil::FromString(const char *value); +template<> +PipelineBroadcastExchangeOrderMode EnumUtil::FromString(const char *value); + template<> PipelineInputMode EnumUtil::FromString(const char *value); diff --git a/src/duckdb/src/include/duckdb/common/enums/row_id_handling.hpp b/src/duckdb/src/include/duckdb/common/enums/row_id_handling.hpp index 9533ce112..64b11f138 100644 --- a/src/duckdb/src/include/duckdb/common/enums/row_id_handling.hpp +++ b/src/duckdb/src/include/duckdb/common/enums/row_id_handling.hpp @@ -9,6 +9,7 @@ #pragma once #include "duckdb/common/constants.hpp" +#include "duckdb/common/windows_undefs.hpp" // test3 namespace duckdb { diff --git a/src/duckdb/src/include/duckdb/execution/index/fixed_size_buffer.hpp b/src/duckdb/src/include/duckdb/execution/index/fixed_size_buffer.hpp index e7ee6878a..a2119c50a 100644 --- a/src/duckdb/src/include/duckdb/execution/index/fixed_size_buffer.hpp +++ b/src/duckdb/src/include/duckdb/execution/index/fixed_size_buffer.hpp @@ -44,6 +44,8 @@ class FixedSizeBuffer { public: //! Constructor for a new in-memory buffer explicit FixedSizeBuffer(BlockManager &block_manager, MemoryTag memory_tag); + FixedSizeBuffer(BlockManager &block_manager, const idx_t segment_count, const idx_t allocation_size, + shared_ptr block_handle); //! Constructor for deserializing buffer metadata from disk FixedSizeBuffer(BlockManager &block_manager, const idx_t segment_count, const idx_t allocation_size, const BlockPointer &block_pointer); diff --git a/src/duckdb/src/include/duckdb/execution/operator/helper/physical_batch_collector.hpp b/src/duckdb/src/include/duckdb/execution/operator/helper/physical_batch_collector.hpp index ff6365f6b..47d329590 100644 --- a/src/duckdb/src/include/duckdb/execution/operator/helper/physical_batch_collector.hpp +++ b/src/duckdb/src/include/duckdb/execution/operator/helper/physical_batch_collector.hpp @@ -37,6 +37,10 @@ class PhysicalBatchCollector : public PhysicalResultCollector { bool ParallelSink() const override { return true; } + + PipelineExternalInputSupport GetExternalInputSupport() const override { + return PipelineExternalInputSupport::SUPPORTED; + } }; //===--------------------------------------------------------------------===// diff --git a/src/duckdb/src/include/duckdb/execution/operator/helper/physical_buffered_batch_collector.hpp b/src/duckdb/src/include/duckdb/execution/operator/helper/physical_buffered_batch_collector.hpp index cb5e8892f..7b5d19e87 100644 --- a/src/duckdb/src/include/duckdb/execution/operator/helper/physical_buffered_batch_collector.hpp +++ b/src/duckdb/src/include/duckdb/execution/operator/helper/physical_buffered_batch_collector.hpp @@ -45,6 +45,10 @@ class PhysicalBufferedBatchCollector : public PhysicalResultCollector { return true; } + PipelineExternalInputSupport GetExternalInputSupport() const override { + return PipelineExternalInputSupport::SUPPORTED; + } + bool IsStreaming() const override { return true; } diff --git a/src/duckdb/src/include/duckdb/execution/operator/persistent/physical_batch_copy_to_file.hpp b/src/duckdb/src/include/duckdb/execution/operator/persistent/physical_batch_copy_to_file.hpp index d37ffc94d..ded366210 100644 --- a/src/duckdb/src/include/duckdb/execution/operator/persistent/physical_batch_copy_to_file.hpp +++ b/src/duckdb/src/include/duckdb/execution/operator/persistent/physical_batch_copy_to_file.hpp @@ -58,9 +58,7 @@ class PhysicalBatchCopyToFile : public PhysicalOperator { unique_ptr GetGlobalSinkState(ClientContext &context) const override; SinkNextBatchType NextBatch(ExecutionContext &context, OperatorSinkNextBatchInput &input) const override; - OperatorPartitionInfo RequiredPartitionInfo() const override { - return OperatorPartitionInfo::BatchIndex(); - } + OperatorPartitionInfo RequiredPartitionInfo() const override; bool IsSink() const override { return true; @@ -70,6 +68,10 @@ class PhysicalBatchCopyToFile : public PhysicalOperator { return true; } + PipelineExternalInputSupport GetExternalInputSupport() const override { + return PipelineExternalInputSupport::SUPPORTED; + } + public: void AddLocalBatch(ClientContext &context, GlobalSinkState &gstate, LocalSinkState &state) const; void AddRawBatchData(ClientContext &context, GlobalSinkState &gstate_p, idx_t batch_index, diff --git a/src/duckdb/src/include/duckdb/execution/operator/persistent/physical_batch_insert.hpp b/src/duckdb/src/include/duckdb/execution/operator/persistent/physical_batch_insert.hpp index 047641c84..69b55ea0d 100644 --- a/src/duckdb/src/include/duckdb/execution/operator/persistent/physical_batch_insert.hpp +++ b/src/duckdb/src/include/duckdb/execution/operator/persistent/physical_batch_insert.hpp @@ -35,6 +35,8 @@ class PhysicalBatchInsert : public PhysicalOperator { optional_ptr schema; //! Create table info, in case of CREATE TABLE AS unique_ptr info; + //! Preferred input batch size + optional_idx preferred_batch_size; public: // Source interface @@ -55,9 +57,7 @@ class PhysicalBatchInsert : public PhysicalOperator { SinkFinalizeType Finalize(Pipeline &pipeline, Event &event, ClientContext &context, OperatorSinkFinalizeInput &input) const override; - OperatorPartitionInfo RequiredPartitionInfo() const override { - return OperatorPartitionInfo::BatchIndex(); - } + OperatorPartitionInfo RequiredPartitionInfo() const override; bool IsSink() const override { return true; @@ -67,6 +67,10 @@ class PhysicalBatchInsert : public PhysicalOperator { return true; } + PipelineExternalInputSupport GetExternalInputSupport() const override { + return PipelineExternalInputSupport::SUPPORTED; + } + private: bool ExecuteTask(ClientContext &context, GlobalSinkState &gstate_p, LocalSinkState &lstate_p) const; void ExecuteTasks(ClientContext &context, GlobalSinkState &gstate_p, LocalSinkState &lstate_p) const; diff --git a/src/duckdb/src/include/duckdb/execution/operator/scan/physical_column_data_scan.hpp b/src/duckdb/src/include/duckdb/execution/operator/scan/physical_column_data_scan.hpp index f3c7fadea..9f16822b2 100644 --- a/src/duckdb/src/include/duckdb/execution/operator/scan/physical_column_data_scan.hpp +++ b/src/duckdb/src/include/duckdb/execution/operator/scan/physical_column_data_scan.hpp @@ -22,6 +22,9 @@ class PhysicalColumnDataScan : public PhysicalOperator { public: PhysicalColumnDataScan(PhysicalPlan &physical_plan, vector types, PhysicalOperatorType op_type, idx_t estimated_cardinality, optionally_owned_ptr collection); + PhysicalColumnDataScan(PhysicalPlan &physical_plan, vector types, PhysicalOperatorType op_type, + idx_t estimated_cardinality, optionally_owned_ptr collection, + vector column_ids); PhysicalColumnDataScan(PhysicalPlan &physical_plan, vector types, PhysicalOperatorType op_type, idx_t estimated_cardinality, TableIndex cte_index); @@ -29,27 +32,41 @@ class PhysicalColumnDataScan : public PhysicalOperator { //! (optionally owned) column data collection to scan optionally_owned_ptr collection; optional_ptr cte_source; + //! Column ids scanned from the underlying collection + vector column_ids; TableIndex cte_index; optional_idx delim_index; public: unique_ptr GetGlobalSourceState(ClientContext &context) const override; + unique_ptr GetGlobalSourceState(ClientContext &context, + const OperatorPartitionInfo &partition_info) const override; unique_ptr GetLocalSourceState(ExecutionContext &context, GlobalSourceState &gstate) const override; SourceResultType GetDataInternal(ExecutionContext &context, DataChunk &chunk, OperatorSourceInput &input) const override; ProgressData GetProgress(ClientContext &context, GlobalSourceState &gstate) const override; + OperatorPartitionData GetPartitionData(ExecutionContext &context, DataChunk &chunk, GlobalSourceState &gstate, + LocalSourceState &lstate, + const OperatorPartitionInfo &partition_info) const override; bool IsSource() const override { return true; } + bool SupportsPartitioning(const OperatorPartitionInfo &partition_info) const override; + + OrderPreservationType SourceOrder() const override { + return source_order; + } InsertionOrderPreservingMap ParamsToString() const override; bool ParallelSource() const override { return true; } + OrderPreservationType source_order = OrderPreservationType::INSERTION_ORDER; + public: void BuildPipelines(Pipeline ¤t, MetaPipeline &meta_pipeline) override; }; diff --git a/src/duckdb/src/include/duckdb/execution/operator/set/physical_cte.hpp b/src/duckdb/src/include/duckdb/execution/operator/set/physical_cte.hpp index 9216a624f..74743b13b 100644 --- a/src/duckdb/src/include/duckdb/execution/operator/set/physical_cte.hpp +++ b/src/duckdb/src/include/duckdb/execution/operator/set/physical_cte.hpp @@ -29,20 +29,25 @@ class PhysicalCTEConsumerSource : public PhysicalOperator { TableIndex cte_index, shared_ptr exchange, idx_t consumer_idx); unique_ptr GetGlobalSourceState(ClientContext &context) const override; + unique_ptr GetGlobalSourceState(ClientContext &context, + const OperatorPartitionInfo &partition_info) const override; unique_ptr GetLocalSourceState(ExecutionContext &context, GlobalSourceState &gstate) const override; SourceResultType GetDataInternal(ExecutionContext &context, DataChunk &chunk, OperatorSourceInput &input) const override; + OperatorPartitionData GetPartitionData(ExecutionContext &context, DataChunk &chunk, GlobalSourceState &gstate, + LocalSourceState &lstate, + const OperatorPartitionInfo &partition_info) const override; ProgressData GetProgress(ClientContext &context, GlobalSourceState &gstate) const override; void SourceFinished(ClientContext &context, GlobalSourceState &gstate) const override; + bool SupportsPartitioning(const OperatorPartitionInfo &partition_info) const override; + OrderPreservationType SourceOrder() const override; bool IsSource() const override { return true; } - bool ParallelSource() const override { - return true; - } + bool ParallelSource() const override; InsertionOrderPreservingMap ParamsToString() const override; @@ -70,6 +75,11 @@ class PhysicalCTE : public PhysicalOperator { Identifier ctename; bool cte_body_has_side_effects = false; CTEPipelineSelectionState pipeline_selection_state = CTEPipelineSelectionState::UNRESOLVED; + bool preserve_order = false; + bool use_batch_index = false; + bool parallel = true; + optional_idx preferred_batch_size; + bool conflicting_batch_sizes = false; public: // Sink interface @@ -79,6 +89,7 @@ class PhysicalCTE : public PhysicalOperator { unique_ptr GetLocalSinkState(ExecutionContext &context) const override; SinkCombineResultType Combine(ExecutionContext &context, OperatorSinkCombineInput &input) const override; + SinkNextBatchType NextBatch(ExecutionContext &context, OperatorSinkNextBatchInput &input) const override; SinkFinalizeType Finalize(Pipeline &pipeline, Event &event, ClientContext &context, OperatorSinkFinalizeInput &input) const override; @@ -87,11 +98,16 @@ class PhysicalCTE : public PhysicalOperator { } bool ParallelSink() const override { - return true; + return parallel; + } + + OperatorPartitionInfo RequiredPartitionInfo() const override { + return use_batch_index ? OperatorPartitionInfo::BatchIndex(preferred_batch_size) + : OperatorPartitionInfo::NoPartitionInfo(); } bool SinkOrderDependent() const override { - return false; + return preserve_order; } InsertionOrderPreservingMap ParamsToString() const override; @@ -103,12 +119,15 @@ class PhysicalCTE : public PhysicalOperator { void BuildPipelines(Pipeline ¤t, MetaPipeline &meta_pipeline) override; bool TryRegisterDirectConsumer(Pipeline &pipeline, idx_t consumer_idx); bool ShouldUseBufferedConsumer(Pipeline &pipeline) const; - void RegisterBufferedConsumer(idx_t consumer_idx); + void RegisterBufferedConsumer(Pipeline &pipeline, idx_t consumer_idx); void RegisterMaterializedConsumer(idx_t consumer_idx); CTEExecutionMode GetExecutionMode() const; bool UseStreamingExchange() const; vector> GetSources() const override; + +private: + void RegisterBatchPreference(Pipeline &pipeline); }; } // namespace duckdb diff --git a/src/duckdb/src/include/duckdb/execution/partition_info.hpp b/src/duckdb/src/include/duckdb/execution/partition_info.hpp index 8a56f77fb..d56739fbc 100644 --- a/src/duckdb/src/include/duckdb/execution/partition_info.hpp +++ b/src/duckdb/src/include/duckdb/execution/partition_info.hpp @@ -41,18 +41,22 @@ struct OperatorPartitionInfo { OperatorPartitionInfo() = default; explicit OperatorPartitionInfo(bool batch_index) : batch_index(batch_index) { } + OperatorPartitionInfo(bool batch_index, optional_idx preferred_batch_size) + : batch_index(batch_index), preferred_batch_size(preferred_batch_size) { + } explicit OperatorPartitionInfo(vector partition_columns_p) : partition_columns(std::move(partition_columns_p)) { } bool batch_index = false; + optional_idx preferred_batch_size; vector partition_columns; static OperatorPartitionInfo NoPartitionInfo() { return OperatorPartitionInfo(false); } - static OperatorPartitionInfo BatchIndex() { - return OperatorPartitionInfo(true); + static OperatorPartitionInfo BatchIndex(optional_idx preferred_batch_size = optional_idx()) { + return OperatorPartitionInfo(true, preferred_batch_size); } static OperatorPartitionInfo PartitionColumns(vector columns) { return OperatorPartitionInfo(std::move(columns)); diff --git a/src/duckdb/src/include/duckdb/execution/physical_operator.hpp b/src/duckdb/src/include/duckdb/execution/physical_operator.hpp index b5a087e73..799ca47b9 100644 --- a/src/duckdb/src/include/duckdb/execution/physical_operator.hpp +++ b/src/duckdb/src/include/duckdb/execution/physical_operator.hpp @@ -138,6 +138,8 @@ class PhysicalOperator { virtual unique_ptr GetLocalSourceState(ExecutionContext &context, GlobalSourceState &gstate) const; virtual unique_ptr GetGlobalSourceState(ClientContext &context) const; + virtual unique_ptr GetGlobalSourceState(ClientContext &context, + const OperatorPartitionInfo &partition_info) const; protected: virtual SourceResultType GetDataInternal(ExecutionContext &context, DataChunk &chunk, diff --git a/src/duckdb/src/include/duckdb/execution/physical_plan_generator.hpp b/src/duckdb/src/include/duckdb/execution/physical_plan_generator.hpp index 2cd479d0b..c659c995b 100644 --- a/src/duckdb/src/include/duckdb/execution/physical_plan_generator.hpp +++ b/src/duckdb/src/include/duckdb/execution/physical_plan_generator.hpp @@ -81,6 +81,7 @@ class PhysicalPlanGenerator { unordered_map> recurring_cte_tables; //! Materialized CTE ids must be collected. unordered_map>> materialized_ctes; + unordered_map materialized_cte_orders; //! The index for duplicate eliminated joins. idx_t delim_index = 0; //! Tracks whether we are planning the recursive member of a recursive CTE. diff --git a/src/duckdb/src/include/duckdb/function/cast/variant/json_to_variant.hpp b/src/duckdb/src/include/duckdb/function/cast/variant/json_to_variant.hpp index e04d3b491..68c30b4e5 100644 --- a/src/duckdb/src/include/duckdb/function/cast/variant/json_to_variant.hpp +++ b/src/duckdb/src/include/duckdb/function/cast/variant/json_to_variant.hpp @@ -37,6 +37,10 @@ struct ReadJSONHolder { yyjson_doc *doc = nullptr; }; +static inline string_t GetString(yyjson_val *val) { + return string_t(unsafe_yyjson_get_str(val), NumericCast(unsafe_yyjson_get_len(val))); +} + } // namespace template @@ -86,6 +90,28 @@ static bool ConvertJSONObject(yyjson_val *obj, ToVariantGlobalResultData &result yyjson_obj_iter iter; yyjson_obj_iter_init(obj, &iter); + struct JSONObjectEntry { + yyjson_val *key; + yyjson_val *value; + }; + + // Maps object keys to their entry index. + string_map_t key_to_entry; + vector entries; + key_to_entry.reserve(iter.max); + entries.reserve(iter.max); + + while (auto key = yyjson_obj_iter_next(&iter)) { + auto key_string = GetString(key); + auto inserted_entry = key_to_entry.emplace(key_string, entries.size()); + auto val = yyjson_obj_iter_get_val(key); + if (inserted_entry.second) { + entries.push_back({key, val}); + } else { + entries[inserted_entry.first->second].value = val; + } + } + auto keys_offset_data = OffsetData::GetKeys(result.offsets); auto children_offset_data = OffsetData::GetChildren(result.offsets); auto values_offset_data = OffsetData::GetValues(result.offsets); @@ -97,7 +123,7 @@ static bool ConvertJSONObject(yyjson_val *obj, ToVariantGlobalResultData &result auto &variant = result.variant; auto &children_list_entry = variant.children_data[result_index]; auto &keys_list_entry = variant.keys_data[result_index]; - uint32_t count = NumericCast(iter.max); + uint32_t count = NumericCast(entries.size()); auto start_child_index = children_list_entry.offset + children_offset_data[result_index]; WriteContainerData(result.variant, result_index, blob_offset_data[result_index], count, children_offset_data[result_index]); @@ -108,15 +134,11 @@ static bool ConvertJSONObject(yyjson_val *obj, ToVariantGlobalResultData &result keys_offset_data[result_index] += count; //! Iterate over all the children in the Object - yyjson_val *key, *val; - while ((key = yyjson_obj_iter_next(&iter))) { - auto key_string = yyjson_get_str(key); - uint32_t key_string_len = NumericCast(unsafe_yyjson_get_len(key)); - + for (const auto &entry : entries) { if (WRITE_DATA) { - auto str = string_t(key_string, key_string_len); + auto key_string = GetString(entry.key); auto keys_index = start_key_index++; - auto dictionary_index = result.GetOrCreateIndex(str); + auto dictionary_index = result.GetOrCreateIndex(key_string); //! Set the keys_index variant.keys_index_data[start_child_index] = keys_index; @@ -125,7 +147,7 @@ static bool ConvertJSONObject(yyjson_val *obj, ToVariantGlobalResultData &result result.keys_selvec.set_index(keys_list_entry.offset + keys_index, dictionary_index); } - val = yyjson_obj_iter_get_val(key); + auto val = entry.value; if (!ConvertJSON(val, result, result_index, false)) { return false; } @@ -133,14 +155,6 @@ static bool ConvertJSONObject(yyjson_val *obj, ToVariantGlobalResultData &result return true; } -namespace { - -static inline string_t GetString(yyjson_val *val) { - return string_t(unsafe_yyjson_get_str(val), NumericCast(unsafe_yyjson_get_len(val))); -} - -} // namespace - template static bool ConvertJSONPrimitive(yyjson_val *val, ToVariantGlobalResultData &result, idx_t result_index, bool is_root) { auto json_tag = unsafe_yyjson_get_tag(val); diff --git a/src/duckdb/src/include/duckdb/main/extension_entries.hpp b/src/duckdb/src/include/duckdb/main/extension_entries.hpp index 58787140b..f7f2ccfc6 100644 --- a/src/duckdb/src/include/duckdb/main/extension_entries.hpp +++ b/src/duckdb/src/include/duckdb/main/extension_entries.hpp @@ -148,6 +148,7 @@ static constexpr ExtensionFunctionEntry EXTENSION_FUNCTIONS[] = { {"countif", "core_functions", CatalogType::AGGREGATE_FUNCTION_ENTRY}, {"covar_pop", "core_functions", CatalogType::AGGREGATE_FUNCTION_ENTRY}, {"covar_samp", "core_functions", CatalogType::AGGREGATE_FUNCTION_ENTRY}, + {"create_fts_boolean_query_macros", "fts", CatalogType::PRAGMA_FUNCTION_ENTRY}, {"create_fts_index", "fts", CatalogType::PRAGMA_FUNCTION_ENTRY}, {"current_database", "core_functions", CatalogType::SCALAR_FUNCTION_ENTRY}, {"current_date", "icu", CatalogType::SCALAR_FUNCTION_ENTRY}, @@ -221,6 +222,7 @@ static constexpr ExtensionFunctionEntry EXTENSION_FUNCTIONS[] = { {"from_json", "json", CatalogType::SCALAR_FUNCTION_ENTRY}, {"from_json_strict", "json", CatalogType::SCALAR_FUNCTION_ENTRY}, {"fsum", "core_functions", CatalogType::AGGREGATE_FUNCTION_ENTRY}, + {"fts_tokenize_opensearch_standard", "fts", CatalogType::SCALAR_FUNCTION_ENTRY}, {"fuzz_all_functions", "sqlsmith", CatalogType::TABLE_FUNCTION_ENTRY}, {"fuzzyduck", "sqlsmith", CatalogType::TABLE_FUNCTION_ENTRY}, {"gamma", "core_functions", CatalogType::SCALAR_FUNCTION_ENTRY}, diff --git a/src/duckdb/src/include/duckdb/optimizer/remove_unused_columns.hpp b/src/duckdb/src/include/duckdb/optimizer/remove_unused_columns.hpp index 7c0aa4c5a..9a5bff52c 100644 --- a/src/duckdb/src/include/duckdb/optimizer/remove_unused_columns.hpp +++ b/src/duckdb/src/include/duckdb/optimizer/remove_unused_columns.hpp @@ -22,6 +22,7 @@ namespace duckdb { class Binder; class BoundColumnRefExpression; class ClientContext; +class LogicalColumnDataGet; class Optimizer; struct ReferencedExtractComponent { @@ -148,6 +149,7 @@ class RemoveUnusedColumns : public BaseColumnPruner { private: template void ClearUnusedExpressions(vector &list, TableIndex table_idx, bool replace = true); + void RemoveColumnsFromLogicalColumnDataGet(LogicalColumnDataGet &get); void RemoveColumnsFromLogicalGet(LogicalGet &get, unique_ptr &op_ref); void CheckPushdownExtract(LogicalOperator &op); void RewriteExpressions(LogicalProjection &proj, idx_t expression_count); diff --git a/src/duckdb/src/include/duckdb/parallel/meta_pipeline.hpp b/src/duckdb/src/include/duckdb/parallel/meta_pipeline.hpp index c7db38511..f74dac103 100644 --- a/src/duckdb/src/include/duckdb/parallel/meta_pipeline.hpp +++ b/src/duckdb/src/include/duckdb/parallel/meta_pipeline.hpp @@ -20,7 +20,7 @@ enum class MetaPipelineType : uint8_t { enum class MetaPipelineDependencyMode : uint8_t { ADD_DEPENDENCY, NO_DEPENDENCY }; enum class RecursiveDependencyMode : uint8_t { RESPECT_PARALLELISM, FORCE }; -enum class DataflowDependencyMode : uint8_t { INCLUDE, SKIP }; +enum class DataflowDependencyMode : uint8_t { INCLUDE, SKIP_CONFLICTING }; //! MetaPipeline represents a set of pipelines that all have the same sink class MetaPipeline : public enable_shared_from_this { diff --git a/src/duckdb/src/include/duckdb/parallel/pipeline.hpp b/src/duckdb/src/include/duckdb/parallel/pipeline.hpp index 7381d3f19..fe5318760 100644 --- a/src/duckdb/src/include/duckdb/parallel/pipeline.hpp +++ b/src/duckdb/src/include/duckdb/parallel/pipeline.hpp @@ -42,6 +42,7 @@ class PipelineTask : public ExecutorTask { Pipeline &pipeline; unique_ptr pipeline_executor; + optional_idx reserved_batch_index; string TaskType() const override { return "PipelineTask"; @@ -148,11 +149,19 @@ class Pipeline : public enable_shared_from_this { bool IsExternalInput() const { return input_mode == PipelineInputMode::EXTERNAL_INPUT; } + void SetExternalStreamingResultProducer() { + external_streaming_result_producer = true; + } + bool IsStreamingResultPipeline() const; void SetExternalInputEvent(const shared_ptr &event); void CompleteExternalInput(); - bool CanUseExternalInput() const; + bool CanUseExternalInput(const OperatorPartitionInfo &source_partition_info) const; bool CanStopSourceEarly() const; + idx_t GetBaseBatchIndex() const { + return base_batch_index; + } + //! Registers a new batch index for a pipeline executor - returns the current minimum batch index idx_t RegisterNewBatchIndex(); @@ -189,6 +198,8 @@ class Pipeline : public enable_shared_from_this { idx_t base_batch_index = 0; //! How this pipeline receives input chunks PipelineInputMode input_mode = PipelineInputMode::SCHEDULED_SOURCE; + //! Whether this pipeline directly feeds a streaming result collector + bool external_streaming_result_producer = false; //! Event that represents execution of an externally fed pipeline weak_ptr external_input_event DUCKDB_GUARDED_BY(external_input_lock); ExternalInputEventState external_input_event_state DUCKDB_GUARDED_BY(external_input_lock) = diff --git a/src/duckdb/src/include/duckdb/parallel/pipeline_broadcast_exchange.hpp b/src/duckdb/src/include/duckdb/parallel/pipeline_broadcast_exchange.hpp index 91bdbb1bf..cc02f53e8 100644 --- a/src/duckdb/src/include/duckdb/parallel/pipeline_broadcast_exchange.hpp +++ b/src/duckdb/src/include/duckdb/parallel/pipeline_broadcast_exchange.hpp @@ -12,8 +12,11 @@ #include "duckdb/common/common.hpp" #include "duckdb/common/deque.hpp" #include "duckdb/common/enums/operator_result_type.hpp" +#include "duckdb/common/enums/order_preservation_type.hpp" #include "duckdb/common/mutex.hpp" +#include "duckdb/common/set.hpp" #include "duckdb/common/types/data_chunk.hpp" +#include "duckdb/execution/partition_info.hpp" #include "duckdb/execution/progress_data.hpp" #include "duckdb/parallel/interrupt.hpp" @@ -22,10 +25,12 @@ namespace duckdb { class ClientContext; class Pipeline; class PipelineBroadcastExchange; +class PipelineBroadcastExchangeScanState; class PipelineExecutor; class PhysicalOperator; enum class PipelineBroadcastExchangeCompletionMode : uint8_t { STOP_WHEN_UNCONSUMED, RUN_TO_COMPLETION }; +enum class PipelineBroadcastExchangeOrderMode : uint8_t { UNORDERED, SEQUENTIAL, BATCH_INDEX }; enum class PipelineBroadcastExchangeLocalMode : uint8_t { DIRECT_ONLY, BUFFERED }; enum class PipelineBroadcastExchangeDirectPushState : uint8_t { NOT_STARTED, RESUMING, ACTIVE, FINISHED }; enum class PipelineBroadcastExchangeConsumerMode : uint8_t { UNRESOLVED, BUFFERED, DIRECT, MATERIALIZED }; @@ -43,20 +48,28 @@ struct PipelineBroadcastExchangeConsumerSummary { class PipelineBroadcastExchangeLocalState { public: - PipelineBroadcastExchangeLocalState(ClientContext &context, const PipelineBroadcastExchange &exchange); + PipelineBroadcastExchangeLocalState(ClientContext &context, const PipelineBroadcastExchange &exchange, + idx_t producer_base_batch_index); ~PipelineBroadcastExchangeLocalState(); private: friend class PipelineBroadcastExchange; - SinkResultType Push(DataChunk &chunk, const InterruptState &interrupt_state); - SinkCombineResultType Finish(const InterruptState &interrupt_state); + SinkResultType Push(DataChunk &chunk, const SourcePartitionInfo &partition_info, + const InterruptState &interrupt_state); + SinkNextBatchType NextBatch(const SourcePartitionInfo &partition_info, const InterruptState &interrupt_state); + SinkCombineResultType Finish(const SourcePartitionInfo &partition_info, const InterruptState &interrupt_state); bool HasDirectConsumers() const; bool DirectConsumersFinished() const; void ResetPush(); + OperatorPartitionData GetSourcePartitionData(const SourcePartitionInfo &partition_info) const; + optional_idx GetSourceMinBatchIndex(const SourcePartitionInfo &partition_info) const; vector> direct_executors; idx_t direct_idx = 0; + idx_t direct_next_batch_idx = 0; idx_t direct_finalize_idx = 0; + idx_t producer_base_batch_index; + bool supports_batch_index; PipelineBroadcastExchangeLocalMode mode = PipelineBroadcastExchangeLocalMode::BUFFERED; PipelineBroadcastExchangeDirectPushState direct_push_state = PipelineBroadcastExchangeDirectPushState::NOT_STARTED; }; @@ -66,7 +79,8 @@ class PipelineBroadcastExchange { public: PipelineBroadcastExchange(ClientContext &context, vector types_p, - PipelineBroadcastExchangeCompletionMode completion_mode_p); + PipelineBroadcastExchangeCompletionMode completion_mode_p, + OrderPreservationType source_order_p, bool use_batch_index_p); const vector &Types() const { return types; @@ -74,9 +88,24 @@ class PipelineBroadcastExchange { bool RunToCompletion() const { return completion_mode == PipelineBroadcastExchangeCompletionMode::RUN_TO_COMPLETION; } + bool PreservesOrder() const { + return order_mode != PipelineBroadcastExchangeOrderMode::UNORDERED; + } + bool SupportsBatchIndex() const { + return order_mode == PipelineBroadcastExchangeOrderMode::BATCH_INDEX; + } + PipelineBroadcastExchangeOrderMode OrderMode() const { + return order_mode; + } + OrderPreservationType SourceOrder() const { + return source_order; + } - unique_ptr GetLocalState(ClientContext &context) const; + unique_ptr GetLocalState(ClientContext &context, + idx_t producer_base_batch_index) const; + shared_ptr GetScanState() const; + void SetProducerPipelines(const vector> &pipelines); idx_t RegisterConsumer(); bool TryRegisterDirectConsumer(Pipeline &pipeline, idx_t consumer_idx); void SelectBufferedConsumer(idx_t consumer_idx); @@ -86,15 +115,17 @@ class PipelineBroadcastExchange { void SetLogOperator(const PhysicalOperator &op); SinkResultType Push(DataChunk &chunk, PipelineBroadcastExchangeLocalState &lstate, - const InterruptState &interrupt_state); + const SourcePartitionInfo &partition_info, const InterruptState &interrupt_state); + SinkNextBatchType NextBatch(PipelineBroadcastExchangeLocalState &lstate, const SourcePartitionInfo &partition_info, + const InterruptState &interrupt_state); SinkCombineResultType FinishLocal(PipelineBroadcastExchangeLocalState &lstate, - const InterruptState &interrupt_state); + const SourcePartitionInfo &partition_info, const InterruptState &interrupt_state); void Finish(); void FinishDirectConsumers(); void Cancel(); - SourceResultType Scan(idx_t consumer_idx, DataChunk &chunk, shared_ptr ¤t_chunk, - const InterruptState &interrupt_state); + SourceResultType Scan(idx_t consumer_idx, DataChunk &chunk, PipelineBroadcastExchangeScanState &scan_state, + optional_idx &batch_index, const InterruptState &interrupt_state); void UnregisterConsumer(idx_t consumer_idx); ProgressData ScanProgress(idx_t consumer_idx, idx_t estimated_cardinality) const; @@ -105,6 +136,8 @@ class PipelineBroadcastExchange { PipelineBroadcastExchangeConsumerSummary GetConsumerSummary() const; private: + friend class PipelineBroadcastExchangeScanState; + struct ChunkPool; struct BroadcastSpool; struct BroadcastSpoolReader; @@ -114,13 +147,12 @@ class PipelineBroadcastExchange { struct BufferState; enum class ConsumerLifecycle : uint8_t { ACTIVE, INACTIVE }; - enum class ConsumerReadState : uint8_t { IDLE, READING }; enum class ProducerState : uint8_t { ACTIVE, FINISHED, CANCELLED }; enum class AppendReservationState : uint8_t { IDLE, RESERVED }; enum class WatermarkState : uint8_t { BELOW_HIGH_WATERMARK, ABOVE_HIGH_WATERMARK }; enum class WriterWakeMode : uint8_t { LOW_WATERMARK, FORCE }; enum class AppendAdmission : uint8_t { READY, BLOCKED, UNCONSUMED, CANCELLED }; - enum class BufferedPushState : uint8_t { NOT_REQUIRED, APPENDED, BLOCKED, UNCONSUMED, CANCELLED }; + enum class BufferedPushState : uint8_t { NOT_REQUIRED, APPENDED, STAGED, BLOCKED, UNCONSUMED, CANCELLED }; enum class ExchangeLogEvent : uint8_t { SPOOL_CREATED, HIGH_WATERMARK_BLOCKED, @@ -146,16 +178,17 @@ class PipelineBroadcastExchange { idx_t rows_read = 0; PipelineBroadcastExchangeConsumerMode mode = PipelineBroadcastExchangeConsumerMode::UNRESOLVED; ConsumerLifecycle lifecycle = ConsumerLifecycle::ACTIVE; - ConsumerReadState read_state = ConsumerReadState::IDLE; - idx_t read_position = 0; - shared_ptr shared_reader; + bool exhausted = false; + set in_flight_reads; }; public: ~PipelineBroadcastExchange(); private: - BufferedPushState Append(DataChunk &chunk, const InterruptState &interrupt_state); + BufferedPushState Append(DataChunk &chunk, idx_t batch_index, idx_t min_batch_index, + const InterruptState &interrupt_state); + BufferedPushState FlushReadyBatches(idx_t min_batch_index, const InterruptState &interrupt_state); SinkResultType CompletePush(DataChunk &chunk, PipelineBroadcastExchangeLocalState &lstate, BufferedPushState buffered_state); void RecordProducedRows(idx_t count); @@ -167,18 +200,22 @@ class PipelineBroadcastExchange { PipelineBroadcastExchangeConsumerSummary GetConsumerSummaryLocked() const DUCKDB_REQUIRES(lock); AppendAdmission PrepareAppendLocked(const InterruptState &interrupt_state, vector &log_entries) DUCKDB_REQUIRES(lock); - AppendAdmission ReserveAppendLocked(const InterruptState &interrupt_state, AppendReservation &reservation, - vector &log_entries) DUCKDB_REQUIRES(lock); + AppendAdmission ReserveAppendLocked(idx_t batch_index, const InterruptState &interrupt_state, + AppendReservation &reservation, vector &log_entries) + DUCKDB_REQUIRES(lock); + AppendAdmission PrepareStageLocked(idx_t batch_index, const InterruptState &interrupt_state, + vector &log_entries) DUCKDB_REQUIRES(lock); BufferedPushState CompleteAppendLocked(const AppendReservation &reservation, shared_ptr copy, - idx_t row_count, vector &readers, + idx_t row_count, bool record_produced_rows, vector &readers, vector &appenders, vector &log_entries) DUCKDB_REQUIRES(lock); void AbortAppendReservation(vector &readers, vector &writers, vector &appenders) DUCKDB_REQUIRES(lock); SourceResultType ReserveScanLocked(idx_t consumer_idx, const InterruptState &interrupt_state, - shared_ptr &next_chunk, SpoolReadReservation &spool_read, - vector &writers, vector &log_entries) - DUCKDB_REQUIRES(lock); + PipelineBroadcastExchangeScanState &scan_state, + shared_ptr &next_chunk, optional_idx &batch_index, + SpoolReadReservation &spool_read, vector &writers, + vector &log_entries) DUCKDB_REQUIRES(lock); void CompleteSpoolReadLocked(idx_t consumer_idx, const SpoolReadReservation &spool_read, DataChunk &chunk, vector &readers, vector &writers, vector &log_entries) DUCKDB_REQUIRES(lock); @@ -200,7 +237,10 @@ class PipelineBroadcastExchange { ClientContext &context; vector types; PipelineBroadcastExchangeCompletionMode completion_mode; + PipelineBroadcastExchangeOrderMode order_mode; + OrderPreservationType source_order; idx_t max_threads; + vector> producer_pipelines; unique_ptr buffer; mutable annotated_mutex lock; diff --git a/src/duckdb/src/include/duckdb/parallel/pipeline_executor.hpp b/src/duckdb/src/include/duckdb/parallel/pipeline_executor.hpp index facb6a6e3..2094969d3 100644 --- a/src/duckdb/src/include/duckdb/parallel/pipeline_executor.hpp +++ b/src/duckdb/src/include/duckdb/parallel/pipeline_executor.hpp @@ -57,7 +57,7 @@ class ExecutionBudget { //! The Pipeline class represents an execution pipeline class PipelineExecutor { public: - PipelineExecutor(ClientContext &context, Pipeline &pipeline); + PipelineExecutor(ClientContext &context, Pipeline &pipeline, optional_idx reserved_batch_index = optional_idx()); //! Fully execute a pipeline with a source and a sink until the source is completely exhausted PipelineExecuteResult Execute(); @@ -65,9 +65,15 @@ class PipelineExecutor { //! Returns true if execution is finished, false if Execute should be called again PipelineExecuteResult Execute(idx_t max_chunks); //! Pushes a chunk from an external producer through this pipeline into the sink - PipelineExecuteResult PushExternal(DataChunk &input); + PipelineExecuteResult PushExternal(DataChunk &input, const OperatorPartitionData &partition_data, + optional_idx source_min_batch_index); + //! Advances an externally-fed pipeline to the producer's next batch + PipelineExecuteResult NextBatchExternal(const OperatorPartitionData &partition_data, + optional_idx source_min_batch_index); + //! Completes the current batch of an externally-fed pipeline + PipelineExecuteResult FinishBatchExternal(optional_idx source_min_batch_index); //! Finalizes an externally-fed pipeline executor after the producer is exhausted - PipelineExecuteResult FinishExternal(); + PipelineExecuteResult FinishExternal(optional_idx source_min_batch_index); //! Called after depleting the source: finalizes the execution of this pipeline executor //! This should only be called once per PipelineExecutor. @@ -159,6 +165,8 @@ class PipelineExecutor { bool should_flush_current_idx = true; //! Whether this executor has already run at least once bool has_executed = false; + //! Whether an externally-fed sink has observed its initial batch + bool external_batch_initialized = false; private: void StartOperator(PhysicalOperator &op); @@ -185,6 +193,9 @@ class PipelineExecutor { //! Notifies the sink that a new batch has started SinkNextBatchType NextBatch(DataChunk &source_chunk, const bool have_more_output); + SinkNextBatchType NextBatch(OperatorPartitionData next_data, bool force = false, + optional_idx external_min_batch_index = optional_idx()); + OperatorPartitionData ToPipelinePartitionData(const OperatorPartitionData &source_data) const; //! Tries to flush all state from intermediate operators. Will return true if all state is flushed, false in the //! case of a blocked sink. diff --git a/src/duckdb/src/include/duckdb/planner/operator/logical_column_data_get.hpp b/src/duckdb/src/include/duckdb/planner/operator/logical_column_data_get.hpp index e05ad0efa..02621f16d 100644 --- a/src/duckdb/src/include/duckdb/planner/operator/logical_column_data_get.hpp +++ b/src/duckdb/src/include/duckdb/planner/operator/logical_column_data_get.hpp @@ -32,10 +32,14 @@ class LogicalColumnDataGet : public LogicalOperator { TableIndex table_index; //! The types of the chunk vector chunk_types; + //! Column ids that are scanned from the collection + vector column_ids; //! (optionally owned) column data collection optionally_owned_ptr collection; public: + void SetColumnIds(vector column_ids); + const vector &GetColumnIds() const; vector GetColumnBindings() override; void Serialize(Serializer &serializer) const override; @@ -46,8 +50,12 @@ class LogicalColumnDataGet : public LogicalOperator { protected: void ResolveTypes() override { - // types are resolved in the constructor - this->types = chunk_types; + types.clear(); + types.reserve(column_ids.size()); + for (auto column_id : column_ids) { + D_ASSERT(column_id < chunk_types.size()); + types.push_back(chunk_types[column_id]); + } } }; } // namespace duckdb diff --git a/src/duckdb/src/include/duckdb/storage/index_storage_info.hpp b/src/duckdb/src/include/duckdb/storage/index_storage_info.hpp index 704b6c809..3f29aebd7 100644 --- a/src/duckdb/src/include/duckdb/storage/index_storage_info.hpp +++ b/src/duckdb/src/include/duckdb/storage/index_storage_info.hpp @@ -10,9 +10,11 @@ #include "duckdb/common/case_insensitive_map.hpp" #include "duckdb/common/identifier.hpp" +#include "duckdb/common/shared_ptr.hpp" #include "duckdb/common/types/value.hpp" #include "duckdb/common/unordered_set.hpp" #include "duckdb/storage/block.hpp" +#include "duckdb/storage/buffer/block_handle.hpp" namespace duckdb { @@ -24,6 +26,9 @@ struct FixedSizeAllocatorInfo { vector segment_counts; vector allocation_sizes; vector buffers_with_free_space; + //! Transient block handles used to initialize read-only entities (i.e., index entries during WAL replay) + //! These are transient runtime state and are never serialized. + shared_ptr>> transient_block_handles; void Serialize(Serializer &serializer) const; static FixedSizeAllocatorInfo Deserialize(Deserializer &deserializer); diff --git a/src/duckdb/src/optimizer/common_subplan_optimizer.cpp b/src/duckdb/src/optimizer/common_subplan_optimizer.cpp index b8f4401c2..c842a755a 100644 --- a/src/duckdb/src/optimizer/common_subplan_optimizer.cpp +++ b/src/duckdb/src/optimizer/common_subplan_optimizer.cpp @@ -8,6 +8,7 @@ #include "duckdb/common/serializer/binary_serializer.hpp" #include "duckdb/common/arena_containers/arena_unordered_map.hpp" #include "duckdb/common/arena_containers/arena_vector.hpp" +#include "duckdb/common/unordered_set.hpp" #include "duckdb/planner/filter/expression_filter.hpp" #include "duckdb/planner/column_binding_map.hpp" @@ -91,6 +92,7 @@ class PlanSignatureTableIndexMap { restore_original_table_index.clear(); restore_original_table_filter_index.clear(); column_ids.clear(); + chunk_column_ids.clear(); projection_ids.clear(); table_indices.clear(); projection_maps.clear(); @@ -318,6 +320,39 @@ class PlanSignatureTableIndexMap { break; } } + if (op.type == LogicalOperatorType::LOGICAL_CHUNK_GET) { + auto &get = op.Cast(); + switch (TYPE) { + case ConversionType::TO_CANONICAL: { + D_ASSERT(chunk_column_ids.empty()); + chunk_column_ids = get.GetColumnIds(); + + unordered_set selected_columns; + for (auto column_id : chunk_column_ids) { + if (!selected_columns.insert(column_id).second) { + return false; + } + } + + // Expose all collection columns so pruning does not change the scan signature. + vector all_column_ids; + all_column_ids.reserve(get.chunk_types.size()); + for (idx_t column_id = 0; column_id < get.chunk_types.size(); column_id++) { + all_column_ids.push_back(column_id); + } + get.SetColumnIds(std::move(all_column_ids)); + + auto &column_index_map = table_index_map.at(table_indices[0]); + for (idx_t column_idx = 0; column_idx < chunk_column_ids.size(); column_idx++) { + column_index_map.Insert(ProjectionIndex(column_idx), ProjectionIndex(chunk_column_ids[column_idx])); + } + break; + } + case ConversionType::RESTORE_ORIGINAL: + get.SetColumnIds(std::move(chunk_column_ids)); + break; + } + } return true; } @@ -415,6 +450,7 @@ class PlanSignatureTableIndexMap { //! Utility to temporarily store column ids, projection_ids, table indices, expression info and children vector column_ids; + vector chunk_column_ids; vector projection_ids; vector> expression_info; vector> children; diff --git a/src/duckdb/src/optimizer/filter_combiner.cpp b/src/duckdb/src/optimizer/filter_combiner.cpp index fbc357d93..66892bba7 100644 --- a/src/duckdb/src/optimizer/filter_combiner.cpp +++ b/src/duckdb/src/optimizer/filter_combiner.cpp @@ -19,6 +19,7 @@ #include "duckdb/planner/table_filter.hpp" #include "duckdb/common/operator/add.hpp" #include "duckdb/common/operator/subtract.hpp" +#include "duckdb/common/string_util.hpp" #include "duckdb/common/types/interval.hpp" #include "duckdb/optimizer/column_lifetime_analyzer.hpp" #include "duckdb/planner/expression_iterator.hpp" @@ -458,13 +459,40 @@ FilterPushdownResult FilterCombiner::TryPushdownPrefixFilter(TableFilterSet &tab return FilterPushdownResult::NO_PUSHDOWN; } +static bool GetCaseInsensitivePrefixBounds(const string &prefix, string &min_prefix, string &max_prefix) { + min_prefix.reserve(prefix.size()); + max_prefix.reserve(prefix.size()); + for (auto c : prefix) { + auto byte = static_cast(c); + if (byte & 0x80) { + return false; + } + auto lower_byte = StringUtil::ASCII_TO_LOWER_MAP[byte]; + min_prefix.push_back(UnsafeNumericCast(StringUtil::ASCII_TO_UPPER_MAP[byte])); + switch (lower_byte) { + case 'i': + max_prefix += "\xC4\xB0"; // U+0130 LATIN CAPITAL LETTER I WITH DOT ABOVE + break; + case 'k': + max_prefix += "\xE2\x84\xAA"; // U+212A KELVIN SIGN + break; + default: + max_prefix.push_back(UnsafeNumericCast(lower_byte)); + break; + } + } + return !min_prefix.empty() && Utf8Proc::FindNextLegalUTF8(max_prefix); +} + FilterPushdownResult FilterCombiner::TryPushdownLikeFilter(TableFilterSet &table_filters, const vector &column_ids, Expression &expr) { if (expr.GetExpressionClass() != ExpressionClass::BOUND_FUNCTION) { return FilterPushdownResult::NO_PUSHDOWN; } auto &func = expr.Cast(); - if (func.Function().GetName() != "~~") { + auto &function_name = func.Function().GetName(); + const bool case_insensitive = function_name == "~~*"; + if (function_name != "~~" && !case_insensitive) { return FilterPushdownResult::NO_PUSHDOWN; } if (func.GetChildren()[0]->GetExpressionClass() != ExpressionClass::BOUND_COLUMN_REF || @@ -500,13 +528,27 @@ FilterPushdownResult FilterCombiner::TryPushdownLikeFilter(TableFilterSet &table } prefix += c; } - if (equality) { + if (equality && !case_insensitive) { //! If the LIKE has no special characters we can turn it into an equality and push that down auto equal_filter = CreateComparisonExpression(*func.GetChildren()[0], ExpressionType::COMPARE_EQUAL, Value(prefix)); table_filters.PushFilter(proj_index, make_uniq(std::move(equal_filter))); return FilterPushdownResult::PUSHED_DOWN_FULLY; } + if (case_insensitive) { + string min_prefix; + string max_prefix; + if (!GetCaseInsensitivePrefixBounds(prefix, min_prefix, max_prefix)) { + return FilterPushdownResult::NO_PUSHDOWN; + } + auto lower_bound = CreateComparisonExpression( + *func.GetChildren()[0], ExpressionType::COMPARE_GREATERTHANOREQUALTO, Value(std::move(min_prefix))); + table_filters.PushFilter(proj_index, make_uniq(std::move(lower_bound))); + auto upper_bound = CreateComparisonExpression(*func.GetChildren()[0], ExpressionType::COMPARE_LESSTHAN, + Value(std::move(max_prefix))); + table_filters.PushFilter(proj_index, make_uniq(std::move(upper_bound))); + return FilterPushdownResult::PUSHED_DOWN_PARTIALLY; + } //! We have a prefix - we can push down the prefix using a bound (x >= PREFIX AND x < next_prefix) // Note that we still need to execute the LIKE filter diff --git a/src/duckdb/src/optimizer/remove_unused_columns.cpp b/src/duckdb/src/optimizer/remove_unused_columns.cpp index 31064c3a4..0b28dccec 100644 --- a/src/duckdb/src/optimizer/remove_unused_columns.cpp +++ b/src/duckdb/src/optimizer/remove_unused_columns.cpp @@ -15,6 +15,7 @@ #include "duckdb/planner/expression_iterator.hpp" #include "duckdb/planner/filter/expression_filter.hpp" #include "duckdb/planner/operator/logical_aggregate.hpp" +#include "duckdb/planner/operator/logical_column_data_get.hpp" #include "duckdb/planner/operator/logical_comparison_join.hpp" #include "duckdb/planner/operator/logical_distinct.hpp" #include "duckdb/planner/operator/logical_filter.hpp" @@ -325,6 +326,12 @@ void RemoveUnusedColumns::VisitOperator(unique_ptr &op_ref) { } return; } + case LogicalOperatorType::LOGICAL_CHUNK_GET: { + LogicalOperatorVisitor::VisitOperatorExpressions(op); + auto &get = op.Cast(); + RemoveColumnsFromLogicalColumnDataGet(get); + return; + } case LogicalOperatorType::LOGICAL_DISTINCT: { auto &distinct = op.Cast(); if (distinct.distinct_type == DistinctType::DISTINCT_ON) { @@ -737,6 +744,16 @@ void RemoveUnusedColumns::CheckPushdownExtract(LogicalOperator &op) { } } +void RemoveUnusedColumns::RemoveColumnsFromLogicalColumnDataGet(LogicalColumnDataGet &get) { + if (everything_referenced) { + return; + } + + auto column_ids = get.GetColumnIds(); + ClearUnusedExpressions(column_ids, get.table_index); + get.SetColumnIds(std::move(column_ids)); +} + void RemoveUnusedColumns::RemoveColumnsFromLogicalGet(LogicalGet &get, unique_ptr &op_ref) { if (everything_referenced) { return; diff --git a/src/duckdb/src/parallel/executor.cpp b/src/duckdb/src/parallel/executor.cpp index c61d954a8..7acd340e3 100644 --- a/src/duckdb/src/parallel/executor.cpp +++ b/src/duckdb/src/parallel/executor.cpp @@ -571,20 +571,9 @@ bool Executor::ResultCollectorIsBlocked() { if (!HasStreamingResultCollector()) { return false; } - if (completed_pipelines + 1 != total_pipelines) { - // The result collector is always in the last pipeline - return false; - } - if (to_be_rescheduled_tasks.empty()) { - return false; - } for (auto &kv : to_be_rescheduled_tasks) { auto &task = kv.second; if (task->TaskBlockedOnResult()) { - // At least one of the blocked tasks is connected to a result collector - // This task could be the only task that could unblock the other non-result-collector tasks - // To prevent a scenario where we halt indefinitely, we return here so it can be unblocked by a call to - // Fetch return true; } } diff --git a/src/duckdb/src/parallel/meta_pipeline.cpp b/src/duckdb/src/parallel/meta_pipeline.cpp index d1d19bc7f..ab4fade1c 100644 --- a/src/duckdb/src/parallel/meta_pipeline.cpp +++ b/src/duckdb/src/parallel/meta_pipeline.cpp @@ -183,15 +183,26 @@ void MetaPipeline::AddRecursiveDependencies(const vector> & const auto thread_count = TaskScheduler::GetScheduler(executor.context).NumberOfThreads(); for (; it != child_meta_pipelines.end(); it++) { for (auto &pipeline : it->get()->pipelines) { - if (dataflow_mode == DataflowDependencyMode::SKIP && pipeline->HasDataflowDependencies()) { - continue; - } if (dependency_mode == RecursiveDependencyMode::RESPECT_PARALLELISM && !PipelineExceedsThreadCount(*pipeline, thread_count)) { continue; } auto &pipeline_deps = pipeline_dependencies[*pipeline]; for (auto &new_dependency : new_dependencies) { + if (dataflow_mode == DataflowDependencyMode::SKIP_CONFLICTING) { + bool conflicts_with_dataflow = false; + for (auto &dataflow_dependency : pipeline->GetDataflowDependencies()) { + auto dependency = dataflow_dependency.lock(); + D_ASSERT(dependency); + if (RefersToSameObject(*dependency, *new_dependency)) { + conflicts_with_dataflow = true; + break; + } + } + if (conflicts_with_dataflow) { + continue; + } + } if (dependency_mode == RecursiveDependencyMode::RESPECT_PARALLELISM && !PipelineExceedsThreadCount(*new_dependency, thread_count)) { continue; diff --git a/src/duckdb/src/parallel/pipeline.cpp b/src/duckdb/src/parallel/pipeline.cpp index 4a9b6e23f..d7e1c0421 100644 --- a/src/duckdb/src/parallel/pipeline.cpp +++ b/src/duckdb/src/parallel/pipeline.cpp @@ -5,6 +5,7 @@ #include "duckdb/common/tree_renderer/text_tree_renderer.hpp" #include "duckdb/execution/executor.hpp" #include "duckdb/execution/operator/aggregate/physical_ungrouped_aggregate.hpp" +#include "duckdb/execution/operator/helper/physical_result_collector.hpp" #include "duckdb/execution/operator/scan/physical_table_scan.hpp" #include "duckdb/execution/operator/set/physical_recursive_cte.hpp" #include "duckdb/main/client_context.hpp" @@ -22,12 +23,15 @@ static shared_ptr ToSharedSourceState(unique_ptr event_p) : ExecutorTask(pipeline_p.executor, std::move(event_p)), pipeline(pipeline_p) { + auto sink = pipeline.GetSink(); + if (sink && sink->RequiredPartitionInfo().AnyRequired()) { + // Account for every task before lazy executor construction can advance the batch minimum. + reserved_batch_index = pipeline.RegisterNewBatchIndex(); + } } bool PipelineTask::TaskBlockedOnResult() const { - // If this returns true, it means the pipeline this task belongs to has a cached chunk - // that was the result of the Sink method returning BLOCKED - return pipeline_executor->RemainingSinkChunk(); + return pipeline.IsStreamingResultPipeline() && pipeline_executor->RemainingSinkChunk(); } const PipelineExecutor &PipelineTask::GetPipelineExecutor() const { @@ -36,7 +40,7 @@ const PipelineExecutor &PipelineTask::GetPipelineExecutor() const { TaskExecutionResult PipelineTask::ExecuteTask(TaskExecutionMode mode) { if (!pipeline_executor) { - pipeline_executor = make_uniq(pipeline.GetClientContext(), pipeline); + pipeline_executor = make_uniq(pipeline.GetClientContext(), pipeline, reserved_batch_index); } pipeline_executor->SetTaskForInterrupts(shared_from_this()); @@ -164,7 +168,7 @@ bool Pipeline::ScheduleParallel(shared_ptr &event) { // Handle partition requirements specific to scheduling auto partition_info = sink->RequiredPartitionInfo(); if (partition_info.batch_index) { - if (!source->SupportsPartitioning(OperatorPartitionInfo::BatchIndex())) { + if (!source->SupportsPartitioning(partition_info)) { throw InternalException( "Attempting to schedule a pipeline where the sink requires batch index but source does not support it"); } @@ -218,14 +222,30 @@ void Pipeline::SetExternalInput() { external_input_event_state = ExternalInputEventState::EXTERNAL_INPUT_UNSET; } -bool Pipeline::CanUseExternalInput() const { +bool Pipeline::IsStreamingResultPipeline() const { + if (external_streaming_result_producer) { + return true; + } + return sink && sink->type == PhysicalOperatorType::RESULT_COLLECTOR && + sink->Cast().IsStreaming(); +} + +bool Pipeline::CanUseExternalInput(const OperatorPartitionInfo &source_partition_info) const { if (!sink || !sink->ParallelSink() || sink->SinkOrderDependent()) { return false; } if (sink->GetExternalInputSupport() != PipelineExternalInputSupport::SUPPORTED) { return false; } - if (sink->RequiredPartitionInfo().AnyRequired()) { + auto required_partition_info = sink->RequiredPartitionInfo(); + if (required_partition_info.RequiresPartitionColumns()) { + return false; + } + if (required_partition_info.RequiresBatchIndex() && base_batch_index != 0) { + // Later ordered sink pipelines rely on pipeline dependencies to sequence their batch ranges. + return false; + } + if (required_partition_info.RequiresBatchIndex() && !source_partition_info.RequiresBatchIndex()) { return false; } for (auto &op_ref : operators) { @@ -363,10 +383,11 @@ void Pipeline::ResetForReschedule(bool reset_sink) { throw InternalException("Source of pipeline does not have IsSource set"); } auto source_state = GetSourceState(); + auto partition_info = sink ? sink->RequiredPartitionInfo() : OperatorPartitionInfo::NoPartitionInfo(); if (allow_reuse && source_state && source_state->SupportsReuse()) { source_state->Reset(client); } else { - SetSourceState(ToSharedSourceState(source->GetGlobalSourceState(client))); + SetSourceState(ToSharedSourceState(source->GetGlobalSourceState(client, partition_info))); } initialized = true; } @@ -377,7 +398,8 @@ void Pipeline::ResetSource(bool force) { } auto source_state = GetSourceState(); if (force || !source_state) { - SetSourceState(ToSharedSourceState(source->GetGlobalSourceState(GetClientContext()))); + auto partition_info = sink ? sink->RequiredPartitionInfo() : OperatorPartitionInfo::NoPartitionInfo(); + SetSourceState(ToSharedSourceState(source->GetGlobalSourceState(GetClientContext(), partition_info))); } } diff --git a/src/duckdb/src/parallel/pipeline_broadcast_exchange.cpp b/src/duckdb/src/parallel/pipeline_broadcast_exchange.cpp index fb823019d..ea98aed70 100644 --- a/src/duckdb/src/parallel/pipeline_broadcast_exchange.cpp +++ b/src/duckdb/src/parallel/pipeline_broadcast_exchange.cpp @@ -2,6 +2,7 @@ #include "duckdb/common/types/column/column_data_collection.hpp" #include "duckdb/common/types/selection_vector.hpp" +#include "duckdb/common/map.hpp" #include "duckdb/logging/log_type.hpp" #include "duckdb/logging/logger.hpp" #include "duckdb/parallel/pipeline.hpp" @@ -110,6 +111,7 @@ struct PipelineBroadcastExchange::BroadcastSpool { struct ChunkEntry { idx_t row_offset; idx_t row_count; + idx_t batch_index; }; BroadcastSpool(ClientContext &context, const vector &types, idx_t base_position_p) @@ -118,12 +120,12 @@ struct PipelineBroadcastExchange::BroadcastSpool { collection.InitializeAppend(append_state); } - void Append(DataChunk &chunk) { + void Append(DataChunk &chunk, idx_t batch_index) { annotated_lock_guard guard(lock); D_ASSERT(chunk.size() > 0); auto row_offset = collection.Count(); collection.Append(append_state, chunk); - chunks.push_back(ChunkEntry {row_offset, chunk.size()}); + chunks.push_back(ChunkEntry {row_offset, chunk.size(), batch_index}); next_position++; append_generation++; } @@ -177,6 +179,14 @@ struct PipelineBroadcastExchange::BroadcastSpool { return position >= base_position && position < next_position; } + idx_t BatchIndex(idx_t position) { + annotated_lock_guard guard(lock); + if (position < base_position || position >= next_position) { + throw InternalException("Attempted to inspect retired pipeline broadcast spool chunk"); + } + return chunks[position - base_position].batch_index; + } + void InitializeReader(BroadcastSpoolReader &reader) { annotated_lock_guard guard(lock); ResetReaderLocked(reader); @@ -222,6 +232,9 @@ PipelineBroadcastExchange::BroadcastSpoolReader::BroadcastSpoolReader(BroadcastS struct PipelineBroadcastExchange::AppendReservation { shared_ptr shared_spool; idx_t position = 0; + idx_t producer_batch_index = DConstants::INVALID_INDEX; + idx_t exchange_batch_index = DConstants::INVALID_INDEX; + bool pending = false; }; struct PipelineBroadcastExchange::SpoolReadReservation { @@ -236,6 +249,18 @@ struct PipelineBroadcastExchange::SpoolReadReservation { struct PipelineBroadcastExchange::BufferedChunk { shared_ptr chunk; + idx_t batch_index; +}; + +class PipelineBroadcastExchangeScanState { + friend class PipelineBroadcastExchange; + +public: + PipelineBroadcastExchangeScanState() = default; + +private: + shared_ptr current_chunk; + shared_ptr spool_reader; }; struct PipelineBroadcastExchange::BufferState { @@ -254,10 +279,17 @@ struct PipelineBroadcastExchange::BufferState { void Reset() { chunks.clear(); + pending_batches.clear(); + active_batches.clear(); shared_spool.reset(); base_position = 0; next_position = 0; buffered_chunks = 0; + pending_chunks = 0; + min_batch_index = 0; + producer_pipeline_index = DConstants::INVALID_INDEX; + producer_pipeline_offset = 0; + max_local_batch_index = 0; } shared_ptr Copy(DataChunk &chunk) { @@ -276,6 +308,41 @@ struct PipelineBroadcastExchange::BufferState { return buffered_chunks; } + idx_t PendingCount() const { + return pending_chunks; + } + + idx_t MinBatchIndex() const { + return min_batch_index; + } + + bool UpdateMinBatchIndex(idx_t new_min_batch_index) { + if (new_min_batch_index <= min_batch_index) { + return false; + } + min_batch_index = new_min_batch_index; + return true; + } + + bool HasReadyBatch() const { + return !pending_batches.empty() && pending_batches.begin()->first <= min_batch_index && + (active_batches.empty() || pending_batches.begin()->first <= *active_batches.begin()); + } + + void RegisterActiveBatch(idx_t batch_index) { + active_batches.insert(batch_index); + } + + void UnregisterActiveBatch(idx_t batch_index) { + auto entry = active_batches.find(batch_index); + D_ASSERT(entry != active_batches.end()); + active_batches.erase(entry); + } + + bool HasEarlierActiveBatch(idx_t batch_index) const { + return !active_batches.empty() && *active_batches.begin() < batch_index; + } + bool HasSharedSpool() const { return shared_spool != nullptr; } @@ -284,14 +351,76 @@ struct PipelineBroadcastExchange::BufferState { return chunks.empty() && !shared_spool; } - void ReserveAppend(AppendReservation &reservation) const { + idx_t ExchangeBatchIndex(idx_t producer_batch_index) { + if (producer_batch_index == DConstants::INVALID_INDEX) { + return DConstants::INVALID_INDEX; + } + auto pipeline_index = producer_batch_index / PipelineBuildState::BATCH_INCREMENT; + auto local_batch_index = producer_batch_index % PipelineBuildState::BATCH_INCREMENT; + if (local_batch_index == 0) { + throw InternalException("Pipeline broadcast exchange received an uninitialized batch index"); + } + if (producer_pipeline_index == DConstants::INVALID_INDEX) { + producer_pipeline_index = pipeline_index; + } else if (pipeline_index != producer_pipeline_index) { + if (pipeline_index < producer_pipeline_index) { + throw InternalException("Pipeline broadcast producer index decreased from %llu to %llu", + producer_pipeline_index, pipeline_index); + } + producer_pipeline_offset += max_local_batch_index; + producer_pipeline_index = pipeline_index; + max_local_batch_index = 0; + } + if (local_batch_index < max_local_batch_index) { + throw InternalException("Pipeline broadcast exchange emitted batch %llu after batch %llu", + local_batch_index, max_local_batch_index); + } + max_local_batch_index = MaxValue(max_local_batch_index, local_batch_index); + auto result = producer_pipeline_offset + local_batch_index - 1; + if (result >= PipelineBuildState::BATCH_INCREMENT - 1) { + throw InternalException("Pipeline broadcast exchange exceeded the batch index range"); + } + return result; + } + + void ReserveAppend(AppendReservation &reservation, idx_t producer_batch_index) { reservation.shared_spool = shared_spool; reservation.position = next_position; + reservation.producer_batch_index = producer_batch_index; + reservation.exchange_batch_index = ExchangeBatchIndex(producer_batch_index); + reservation.pending = false; + } + + shared_ptr ReservePendingAppend(AppendReservation &reservation) { + D_ASSERT(HasReadyBatch()); + auto &pending = pending_batches.begin()->second.front(); + reservation.shared_spool = shared_spool; + reservation.position = next_position; + reservation.producer_batch_index = pending.batch_index; + reservation.exchange_batch_index = ExchangeBatchIndex(pending.batch_index); + reservation.pending = true; + return pending.chunk; + } + + void Stage(shared_ptr copy, idx_t batch_index) { + pending_batches[batch_index].push_back({std::move(copy), batch_index}); + pending_chunks++; } void CompleteAppend(const AppendReservation &reservation, shared_ptr copy) { + if (reservation.pending) { + D_ASSERT(HasReadyBatch()); + auto pending_entry = pending_batches.begin(); + D_ASSERT(pending_entry->first == reservation.producer_batch_index); + D_ASSERT(pending_entry->second.front().chunk == copy); + pending_entry->second.pop_front(); + pending_chunks--; + if (pending_entry->second.empty()) { + pending_batches.erase(pending_entry); + } + } if (!reservation.shared_spool) { - chunks.push_back({std::move(copy)}); + chunks.push_back({std::move(copy), reservation.exchange_batch_index}); } buffered_chunks++; D_ASSERT(next_position == reservation.position); @@ -299,7 +428,7 @@ struct PipelineBroadcastExchange::BufferState { } void ReserveRead(idx_t position, shared_ptr &reader, shared_ptr &next_chunk, - SpoolReadReservation &spool_read) const { + optional_idx &batch_index, SpoolReadReservation &spool_read) const { D_ASSERT(position < next_position); if (shared_spool) { if (!shared_spool->HasPosition(position)) { @@ -311,19 +440,26 @@ struct PipelineBroadcastExchange::BufferState { spool_read.spool = shared_spool; spool_read.reader = reader; spool_read.position = position; + auto stored_batch_index = shared_spool->BatchIndex(position); + if (stored_batch_index != DConstants::INVALID_INDEX) { + batch_index = stored_batch_index; + } return; } D_ASSERT(position >= base_position); auto chunk_idx = position - base_position; D_ASSERT(chunk_idx < chunks.size()); next_chunk = chunks[chunk_idx].chunk; + if (chunks[chunk_idx].batch_index != DConstants::INVALID_INDEX) { + batch_index = chunks[chunk_idx].batch_index; + } } void CreateSharedSpool() { D_ASSERT(!shared_spool); shared_spool = make_shared_ptr(context, types, base_position); for (auto &chunk : chunks) { - shared_spool->Append(*chunk.chunk); + shared_spool->Append(*chunk.chunk, chunk.batch_index); } chunks.clear(); } @@ -344,8 +480,10 @@ struct PipelineBroadcastExchange::BufferState { void Release() { chunks.clear(); + pending_batches.clear(); shared_spool.reset(); buffered_chunks = 0; + pending_chunks = 0; base_position = next_position; } @@ -353,10 +491,17 @@ struct PipelineBroadcastExchange::BufferState { const vector &types; shared_ptr chunk_pool; deque chunks; + map> pending_batches; + multiset active_batches; shared_ptr shared_spool; idx_t base_position = 0; idx_t next_position = 0; idx_t buffered_chunks = 0; + idx_t pending_chunks = 0; + idx_t min_batch_index = 0; + idx_t producer_pipeline_index = DConstants::INVALID_INDEX; + idx_t producer_pipeline_offset = 0; + idx_t max_local_batch_index = 0; }; PipelineBroadcastExchange::ConsumerState::ConsumerState() = default; @@ -368,7 +513,9 @@ PipelineBroadcastExchange::ConsumerState::operator=(ConsumerState &&other) noexc PipelineBroadcastExchange::~PipelineBroadcastExchange() = default; PipelineBroadcastExchangeLocalState::PipelineBroadcastExchangeLocalState(ClientContext &context, - const PipelineBroadcastExchange &exchange) { + const PipelineBroadcastExchange &exchange, + idx_t producer_base_batch_index_p) + : producer_base_batch_index(producer_base_batch_index_p), supports_batch_index(exchange.SupportsBatchIndex()) { vector> direct_pipeline_refs; { annotated_lock_guard guard(exchange.lock); @@ -396,14 +543,34 @@ void PipelineBroadcastExchange::SetLogOperator(const PhysicalOperator &op) { PipelineBroadcastExchangeLocalState::~PipelineBroadcastExchangeLocalState() = default; PipelineBroadcastExchange::PipelineBroadcastExchange(ClientContext &context, vector types_p, - PipelineBroadcastExchangeCompletionMode completion_mode_p) + PipelineBroadcastExchangeCompletionMode completion_mode_p, + OrderPreservationType source_order_p, bool use_batch_index_p) : context(context), types(std::move(types_p)), completion_mode(completion_mode_p), + order_mode(use_batch_index_p ? PipelineBroadcastExchangeOrderMode::BATCH_INDEX + : source_order_p == OrderPreservationType::NO_ORDER ? PipelineBroadcastExchangeOrderMode::UNORDERED + : PipelineBroadcastExchangeOrderMode::SEQUENTIAL), + source_order(source_order_p), max_threads(NumericCast(TaskScheduler::GetScheduler(context).NumberOfThreads())) { + D_ASSERT(!use_batch_index_p || source_order_p != OrderPreservationType::NO_ORDER); buffer = make_uniq(context, types, max_threads); } -unique_ptr PipelineBroadcastExchange::GetLocalState(ClientContext &context) const { - return make_uniq(context, *this); +unique_ptr +PipelineBroadcastExchange::GetLocalState(ClientContext &context, idx_t producer_base_batch_index) const { + return make_uniq(context, *this, producer_base_batch_index); +} + +shared_ptr PipelineBroadcastExchange::GetScanState() const { + return make_shared_ptr(); +} + +void PipelineBroadcastExchange::SetProducerPipelines(const vector> &pipelines) { + D_ASSERT(!pipelines.empty()); + annotated_lock_guard guard(lock); + producer_pipelines.clear(); + for (auto &pipeline : pipelines) { + producer_pipelines.push_back(*pipeline); + } } idx_t PipelineBroadcastExchange::RegisterConsumer() { @@ -426,10 +593,16 @@ void PipelineBroadcastExchange::SelectMaterializedConsumer(idx_t consumer_idx) { } bool PipelineBroadcastExchange::TryRegisterDirectConsumer(Pipeline &pipeline, idx_t consumer_idx) { - if (!pipeline.CanUseExternalInput()) { + auto source_partition_info = + SupportsBatchIndex() ? OperatorPartitionInfo::BatchIndex() : OperatorPartitionInfo::NoPartitionInfo(); + if (!pipeline.CanUseExternalInput(source_partition_info)) { return false; } annotated_lock_guard guard(lock); + auto required_partition_info = pipeline.GetSink()->RequiredPartitionInfo(); + if (required_partition_info.RequiresBatchIndex() && producer_pipelines.size() != 1) { + return false; + } D_ASSERT(consumer_idx < consumers.size()); auto &consumer = consumers[consumer_idx]; if (consumer.mode == PipelineBroadcastExchangeConsumerMode::DIRECT) { @@ -439,6 +612,11 @@ bool PipelineBroadcastExchange::TryRegisterDirectConsumer(Pipeline &pipeline, id consumer.mode = PipelineBroadcastExchangeConsumerMode::DIRECT; DeactivateConsumerLocked(consumer, buffer->BasePosition()); direct_pipelines.push_back(pipeline); + if (pipeline.IsStreamingResultPipeline()) { + for (auto &producer_pipeline : producer_pipelines) { + producer_pipeline.get().SetExternalStreamingResultProducer(); + } + } return true; } @@ -455,6 +633,7 @@ void PipelineBroadcastExchange::ResetConsumerRegistrations() { buffer->EndExecution(); ResetExchangeStateLocked(); direct_pipelines.clear(); + producer_pipelines.clear(); blocked_readers.clear(); blocked_writers.clear(); blocked_appenders.clear(); @@ -496,9 +675,8 @@ void PipelineBroadcastExchange::ResetExchangeStateLocked() { void PipelineBroadcastExchange::ResetConsumerReadStateLocked(ConsumerState &consumer, idx_t position) { consumer.position = position; - consumer.read_state = ConsumerReadState::IDLE; - consumer.read_position = position; - consumer.shared_reader.reset(); + consumer.exhausted = false; + consumer.in_flight_reads.clear(); } void PipelineBroadcastExchange::ResetConsumerRegistrationLocked(ConsumerState &consumer) { @@ -530,18 +708,27 @@ void PipelineBroadcastExchange::DeactivateConsumerLocked(ConsumerState &consumer } SinkResultType PipelineBroadcastExchange::Push(DataChunk &chunk, PipelineBroadcastExchangeLocalState &lstate, + const SourcePartitionInfo &partition_info, const InterruptState &interrupt_state) { if (lstate.HasDirectConsumers() && (lstate.direct_push_state == PipelineBroadcastExchangeDirectPushState::NOT_STARTED || lstate.direct_push_state == PipelineBroadcastExchangeDirectPushState::RESUMING)) { - auto direct_result = lstate.Push(chunk, interrupt_state); + auto direct_result = lstate.Push(chunk, partition_info, interrupt_state); if (direct_result == SinkResultType::BLOCKED) { return SinkResultType::BLOCKED; } } if (lstate.mode == PipelineBroadcastExchangeLocalMode::BUFFERED) { - auto append_result = Append(chunk, interrupt_state); + idx_t batch_index = DConstants::INVALID_INDEX; + idx_t min_batch_index = DConstants::INVALID_INDEX; + if (SupportsBatchIndex()) { + D_ASSERT(partition_info.batch_index.IsValid()); + D_ASSERT(partition_info.min_batch_index.IsValid()); + batch_index = partition_info.batch_index.GetIndex(); + min_batch_index = partition_info.min_batch_index.GetIndex(); + } + auto append_result = Append(chunk, batch_index, min_batch_index, interrupt_state); if (append_result == BufferedPushState::BLOCKED) { return SinkResultType::BLOCKED; } @@ -552,7 +739,8 @@ SinkResultType PipelineBroadcastExchange::Push(DataChunk &chunk, PipelineBroadca SinkResultType PipelineBroadcastExchange::CompletePush(DataChunk &chunk, PipelineBroadcastExchangeLocalState &lstate, BufferedPushState buffered_state) { - if (buffered_state != BufferedPushState::APPENDED && buffered_state != BufferedPushState::CANCELLED) { + if (buffered_state != BufferedPushState::APPENDED && buffered_state != BufferedPushState::STAGED && + buffered_state != BufferedPushState::CANCELLED) { RecordProducedRows(chunk.size()); } const auto direct_consumers_finished = lstate.DirectConsumersFinished(); @@ -565,22 +753,47 @@ SinkResultType PipelineBroadcastExchange::CompletePush(DataChunk &chunk, Pipelin return SinkResultType::NEED_MORE_INPUT; } +SinkNextBatchType PipelineBroadcastExchange::NextBatch(PipelineBroadcastExchangeLocalState &lstate, + const SourcePartitionInfo &partition_info, + const InterruptState &interrupt_state) { + if (!SupportsBatchIndex()) { + return SinkNextBatchType::READY; + } + D_ASSERT(partition_info.min_batch_index.IsValid()); + auto result = FlushReadyBatches(partition_info.min_batch_index.GetIndex(), interrupt_state); + if (result == BufferedPushState::BLOCKED) { + return SinkNextBatchType::BLOCKED; + } + return lstate.NextBatch(partition_info, interrupt_state); +} + SinkCombineResultType PipelineBroadcastExchange::FinishLocal(PipelineBroadcastExchangeLocalState &lstate, + const SourcePartitionInfo &partition_info, const InterruptState &interrupt_state) { - return lstate.Finish(interrupt_state); + if (SupportsBatchIndex()) { + D_ASSERT(partition_info.min_batch_index.IsValid()); + auto result = FlushReadyBatches(partition_info.min_batch_index.GetIndex(), interrupt_state); + if (result == BufferedPushState::BLOCKED) { + return SinkCombineResultType::BLOCKED; + } + } + return lstate.Finish(partition_info, interrupt_state); } -SinkResultType PipelineBroadcastExchangeLocalState::Push(DataChunk &chunk, const InterruptState &interrupt_state) { +SinkResultType PipelineBroadcastExchangeLocalState::Push(DataChunk &chunk, const SourcePartitionInfo &partition_info, + const InterruptState &interrupt_state) { if (direct_push_state != PipelineBroadcastExchangeDirectPushState::RESUMING) { direct_idx = 0; } + auto source_partition_data = GetSourcePartitionData(partition_info); + auto source_min_batch_index = GetSourceMinBatchIndex(partition_info); for (; direct_idx < direct_executors.size(); direct_idx++) { auto &executor = *direct_executors[direct_idx]; executor.SetInterruptState(interrupt_state); if (executor.IsFinishedProcessing()) { continue; } - auto result = executor.PushExternal(chunk); + auto result = executor.PushExternal(chunk, source_partition_data, source_min_batch_index); if (result == PipelineExecuteResult::INTERRUPTED) { direct_push_state = PipelineBroadcastExchangeDirectPushState::RESUMING; return SinkResultType::BLOCKED; @@ -597,13 +810,46 @@ SinkResultType PipelineBroadcastExchangeLocalState::Push(DataChunk &chunk, const return SinkResultType::NEED_MORE_INPUT; } -SinkCombineResultType PipelineBroadcastExchangeLocalState::Finish(const InterruptState &interrupt_state) { +SinkNextBatchType PipelineBroadcastExchangeLocalState::NextBatch(const SourcePartitionInfo &partition_info, + const InterruptState &interrupt_state) { + auto source_min_batch_index = GetSourceMinBatchIndex(partition_info); + if (!partition_info.batch_index.IsValid()) { + throw InternalException("Batch-indexed pipeline broadcast exchange received no batch index"); + } + auto batch_index = partition_info.batch_index.GetIndex(); + auto max_batch_index = producer_base_batch_index + PipelineBuildState::BATCH_INCREMENT - 1; + const bool finalize = batch_index == max_batch_index; + OperatorPartitionData source_partition_data(0); + if (!finalize) { + source_partition_data = GetSourcePartitionData(partition_info); + } + + for (; direct_next_batch_idx < direct_executors.size(); direct_next_batch_idx++) { + auto &executor = *direct_executors[direct_next_batch_idx]; + executor.SetInterruptState(interrupt_state); + PipelineExecuteResult result; + if (finalize || executor.IsFinishedProcessing()) { + result = executor.FinishBatchExternal(source_min_batch_index); + } else { + result = executor.NextBatchExternal(source_partition_data, source_min_batch_index); + } + if (result == PipelineExecuteResult::INTERRUPTED) { + return SinkNextBatchType::BLOCKED; + } + } + direct_next_batch_idx = 0; + return SinkNextBatchType::READY; +} + +SinkCombineResultType PipelineBroadcastExchangeLocalState::Finish(const SourcePartitionInfo &partition_info, + const InterruptState &interrupt_state) { + auto source_min_batch_index = GetSourceMinBatchIndex(partition_info); for (; direct_finalize_idx < direct_executors.size(); direct_finalize_idx++) { auto &executor = *direct_executors[direct_finalize_idx]; executor.SetInterruptState(interrupt_state); auto result = PipelineExecuteResult::NOT_FINISHED; while (result == PipelineExecuteResult::NOT_FINISHED) { - result = executor.FinishExternal(); + result = executor.FinishExternal(source_min_batch_index); } if (result == PipelineExecuteResult::INTERRUPTED) { return SinkCombineResultType::BLOCKED; @@ -612,6 +858,51 @@ SinkCombineResultType PipelineBroadcastExchangeLocalState::Finish(const Interrup return SinkCombineResultType::FINISHED; } +OperatorPartitionData +PipelineBroadcastExchangeLocalState::GetSourcePartitionData(const SourcePartitionInfo &partition_info) const { + OperatorPartitionData result(0); + if (!supports_batch_index) { + return result; + } + if (!partition_info.batch_index.IsValid()) { + throw InternalException("Batch-indexed pipeline broadcast exchange received no batch index"); + } + auto batch_index = partition_info.batch_index.GetIndex(); + if (batch_index <= producer_base_batch_index) { + throw InternalException("Pipeline broadcast exchange received invalid producer batch index %llu", batch_index); + } + result.batch_index = batch_index - producer_base_batch_index - 1; + if (result.batch_index >= PipelineBuildState::BATCH_INCREMENT - 2) { + throw InternalException("Pipeline broadcast exchange received producer batch index outside its pipeline"); + } + return result; +} + +optional_idx +PipelineBroadcastExchangeLocalState::GetSourceMinBatchIndex(const SourcePartitionInfo &partition_info) const { + if (!supports_batch_index) { + return optional_idx(); + } + if (!partition_info.min_batch_index.IsValid()) { + throw InternalException("Batch-indexed pipeline broadcast exchange received no minimum batch index"); + } + auto producer_min_batch_index = partition_info.min_batch_index.GetIndex(); + if (producer_min_batch_index < producer_base_batch_index) { + throw InternalException("Pipeline broadcast exchange received invalid producer minimum batch index %llu", + producer_min_batch_index); + } + if (partition_info.batch_index.IsValid() && producer_min_batch_index > partition_info.batch_index.GetIndex()) { + throw InternalException("Pipeline broadcast exchange received invalid producer minimum batch index %llu", + producer_min_batch_index); + } + auto result = producer_min_batch_index - producer_base_batch_index; + if (result >= PipelineBuildState::BATCH_INCREMENT) { + throw InternalException("Pipeline broadcast exchange received producer minimum batch index outside its " + "pipeline"); + } + return optional_idx(result); +} + bool PipelineBroadcastExchangeLocalState::HasDirectConsumers() const { return !direct_executors.empty(); } @@ -653,20 +944,40 @@ PipelineBroadcastExchange::PrepareAppendLocked(const InterruptState &interrupt_s } PipelineBroadcastExchange::AppendAdmission -PipelineBroadcastExchange::ReserveAppendLocked(const InterruptState &interrupt_state, AppendReservation &reservation, - vector &log_entries) { +PipelineBroadcastExchange::ReserveAppendLocked(idx_t batch_index, const InterruptState &interrupt_state, + AppendReservation &reservation, vector &log_entries) { auto admission = PrepareAppendLocked(interrupt_state, log_entries); if (admission != AppendAdmission::READY) { return admission; } append_reservation_state = AppendReservationState::RESERVED; - buffer->ReserveAppend(reservation); + buffer->ReserveAppend(reservation, batch_index); + return AppendAdmission::READY; +} + +PipelineBroadcastExchange::AppendAdmission +PipelineBroadcastExchange::PrepareStageLocked(idx_t batch_index, const InterruptState &interrupt_state, + vector &log_entries) { + if (producer_state != ProducerState::ACTIVE) { + return AppendAdmission::CANCELLED; + } + if (active_consumers == 0) { + return AppendAdmission::UNCONSUMED; + } + if (batch_index > buffer->MinBatchIndex() && buffer->PendingCount() >= PIPELINE_BROADCAST_HIGH_WATERMARK_CHUNKS) { + if (watermark_state == WatermarkState::BELOW_HIGH_WATERMARK) { + watermark_state = WatermarkState::ABOVE_HIGH_WATERMARK; + log_entries.push_back({ExchangeLogEvent::HIGH_WATERMARK_BLOCKED, active_consumers, buffer->PendingCount()}); + } + blocked_writers.push_back(interrupt_state); + return AppendAdmission::BLOCKED; + } return AppendAdmission::READY; } PipelineBroadcastExchange::BufferedPushState PipelineBroadcastExchange::CompleteAppendLocked( - const AppendReservation &reservation, shared_ptr copy, idx_t row_count, vector &readers, - vector &appenders, vector &log_entries) { + const AppendReservation &reservation, shared_ptr copy, idx_t row_count, bool record_produced_rows, + vector &readers, vector &appenders, vector &log_entries) { D_ASSERT(append_reservation_state == AppendReservationState::RESERVED); append_reservation_state = AppendReservationState::IDLE; if (producer_state == ProducerState::CANCELLED) { @@ -683,7 +994,9 @@ PipelineBroadcastExchange::BufferedPushState PipelineBroadcastExchange::Complete } buffer->CompleteAppend(reservation, std::move(copy)); WakeReadersLocked(readers); - RecordProducedRows(row_count); + if (record_produced_rows) { + RecordProducedRows(row_count); + } WakeAppendersLocked(appenders); return BufferedPushState::APPENDED; } @@ -701,13 +1014,124 @@ void PipelineBroadcastExchange::AbortAppendReservation(vector &r WakeAppendersLocked(appenders); } -PipelineBroadcastExchange::BufferedPushState PipelineBroadcastExchange::Append(DataChunk &chunk, +PipelineBroadcastExchange::BufferedPushState +PipelineBroadcastExchange::FlushReadyBatches(idx_t min_batch_index, const InterruptState &interrupt_state) { + D_ASSERT(SupportsBatchIndex()); + bool appended = false; + while (true) { + vector log_entries; + vector writers; + AppendAdmission admission; + AppendReservation reservation; + shared_ptr copy; + bool has_ready_batch; + try { + annotated_lock_guard guard(lock); + if (buffer->UpdateMinBatchIndex(min_batch_index)) { + WakeWritersLocked(writers, log_entries, WriterWakeMode::FORCE); + } + has_ready_batch = buffer->HasReadyBatch(); + if (has_ready_batch) { + admission = PrepareAppendLocked(interrupt_state, log_entries); + if (admission == AppendAdmission::READY) { + append_reservation_state = AppendReservationState::RESERVED; + copy = buffer->ReservePendingAppend(reservation); + } + } + } catch (...) { + Cancel(); + throw; + } + CallbackAll(writers); + LogTransitions(log_entries); + if (!has_ready_batch) { + return appended ? BufferedPushState::APPENDED : BufferedPushState::NOT_REQUIRED; + } + if (admission == AppendAdmission::BLOCKED) { + return BufferedPushState::BLOCKED; + } + if (admission == AppendAdmission::UNCONSUMED) { + return BufferedPushState::UNCONSUMED; + } + if (admission == AppendAdmission::CANCELLED) { + return BufferedPushState::CANCELLED; + } + + try { + if (reservation.shared_spool) { + reservation.shared_spool->Append(*copy, reservation.exchange_batch_index); + } + } catch (...) { + vector readers; + vector writers; + vector appenders; + { + annotated_lock_guard guard(lock); + AbortAppendReservation(readers, writers, appenders); + } + CallbackAll(readers); + CallbackAll(writers); + CallbackAll(appenders); + buffer->EndExecution(); + throw; + } + + vector readers; + writers.clear(); + vector appenders; + log_entries.clear(); + BufferedPushState result; + { + annotated_lock_guard guard(lock); + result = CompleteAppendLocked(reservation, std::move(copy), 0, false, readers, appenders, log_entries); + WakeWritersLocked(writers, log_entries, WriterWakeMode::FORCE); + } + CallbackAll(readers); + CallbackAll(writers); + CallbackAll(appenders); + LogTransitions(log_entries); + if (result != BufferedPushState::APPENDED) { + return result; + } + appended = true; + } +} + +PipelineBroadcastExchange::BufferedPushState PipelineBroadcastExchange::Append(DataChunk &chunk, idx_t batch_index, + idx_t min_batch_index, const InterruptState &interrupt_state) { + if (SupportsBatchIndex()) { + annotated_lock_guard guard(lock); + buffer->RegisterActiveBatch(batch_index); + } + auto unregister_active_batch = [&](idx_t *) { + if (!SupportsBatchIndex()) { + return; + } + annotated_lock_guard guard(lock); + buffer->UnregisterActiveBatch(batch_index); + }; + unique_ptr active_batch_guard( + SupportsBatchIndex() ? &batch_index : nullptr, unregister_active_batch); + + if (SupportsBatchIndex()) { + auto flush_result = FlushReadyBatches(min_batch_index, interrupt_state); + if (flush_result == BufferedPushState::BLOCKED || flush_result == BufferedPushState::UNCONSUMED || + flush_result == BufferedPushState::CANCELLED) { + return flush_result; + } + } + vector log_entries; AppendAdmission admission; try { annotated_lock_guard guard(lock); - admission = PrepareAppendLocked(interrupt_state, log_entries); + if (SupportsBatchIndex()) { + buffer->UpdateMinBatchIndex(min_batch_index); + } + admission = SupportsBatchIndex() && batch_index > buffer->MinBatchIndex() + ? PrepareStageLocked(batch_index, interrupt_state, log_entries) + : PrepareAppendLocked(interrupt_state, log_entries); } catch (...) { Cancel(); throw; @@ -732,9 +1156,23 @@ PipelineBroadcastExchange::BufferedPushState PipelineBroadcastExchange::Append(D } AppendReservation reservation; log_entries.clear(); + bool staged = false; try { annotated_lock_guard guard(lock); - admission = ReserveAppendLocked(interrupt_state, reservation, log_entries); + if (SupportsBatchIndex()) { + buffer->UpdateMinBatchIndex(min_batch_index); + } + if (SupportsBatchIndex() && (batch_index > buffer->MinBatchIndex() || + buffer->HasEarlierActiveBatch(batch_index) || buffer->HasReadyBatch())) { + admission = PrepareStageLocked(batch_index, interrupt_state, log_entries); + if (admission == AppendAdmission::READY) { + buffer->Stage(copy, batch_index); + RecordProducedRows(chunk.size()); + staged = true; + } + } else { + admission = ReserveAppendLocked(batch_index, interrupt_state, reservation, log_entries); + } } catch (...) { Cancel(); throw; @@ -749,10 +1187,13 @@ PipelineBroadcastExchange::BufferedPushState PipelineBroadcastExchange::Append(D if (admission == AppendAdmission::CANCELLED) { return BufferedPushState::CANCELLED; } + if (staged) { + return BufferedPushState::STAGED; + } try { if (reservation.shared_spool) { - reservation.shared_spool->Append(*copy); + reservation.shared_spool->Append(*copy, reservation.exchange_batch_index); } } catch (...) { vector readers; @@ -775,7 +1216,8 @@ PipelineBroadcastExchange::BufferedPushState PipelineBroadcastExchange::Append(D BufferedPushState result; { annotated_lock_guard guard(lock); - result = CompleteAppendLocked(reservation, std::move(copy), chunk.size(), readers, appenders, log_entries); + result = + CompleteAppendLocked(reservation, std::move(copy), chunk.size(), true, readers, appenders, log_entries); } CallbackAll(readers); CallbackAll(appenders); @@ -793,6 +1235,10 @@ void PipelineBroadcastExchange::Finish() { vector log_entries; { annotated_lock_guard guard(lock); + if (active_consumers > 0 && buffer->PendingCount() > 0) { + throw InternalException("Finishing ordered pipeline broadcast exchange with %llu pending chunks", + buffer->PendingCount()); + } if (producer_state == ProducerState::ACTIVE) { producer_state = ProducerState::FINISHED; } @@ -836,17 +1282,19 @@ void PipelineBroadcastExchange::Cancel() { } SourceResultType PipelineBroadcastExchange::Scan(idx_t consumer_idx, DataChunk &chunk, - shared_ptr ¤t_chunk, - const InterruptState &interrupt_state) { + PipelineBroadcastExchangeScanState &scan_state, + optional_idx &batch_index, const InterruptState &interrupt_state) { vector writers; vector readers; vector log_entries; shared_ptr next_chunk; SpoolReadReservation spool_read; SourceResultType result; + batch_index = optional_idx(); { annotated_lock_guard guard(lock); - result = ReserveScanLocked(consumer_idx, interrupt_state, next_chunk, spool_read, writers, log_entries); + result = ReserveScanLocked(consumer_idx, interrupt_state, scan_state, next_chunk, batch_index, spool_read, + writers, log_entries); } if (spool_read.IsSet()) { @@ -859,7 +1307,7 @@ SourceResultType PipelineBroadcastExchange::Scan(idx_t consumer_idx, DataChunk & { annotated_lock_guard guard(lock); auto &consumer = consumers[consumer_idx]; - consumer.read_state = ConsumerReadState::IDLE; + consumer.in_flight_reads.erase(spool_read.position); producer_state = ProducerState::CANCELLED; DeactivateAllConsumersLocked(); TryReleaseBufferedStorageLocked(); @@ -884,33 +1332,28 @@ SourceResultType PipelineBroadcastExchange::Scan(idx_t consumer_idx, DataChunk & CallbackAll(writers); LogTransitions(log_entries); if (next_chunk) { - current_chunk = std::move(next_chunk); - chunk.Reference(*current_chunk); + scan_state.current_chunk = std::move(next_chunk); + chunk.Reference(*scan_state.current_chunk); } return result; } -SourceResultType PipelineBroadcastExchange::ReserveScanLocked(idx_t consumer_idx, const InterruptState &interrupt_state, - shared_ptr &next_chunk, - SpoolReadReservation &spool_read, - vector &writers, - vector &log_entries) { +SourceResultType PipelineBroadcastExchange::ReserveScanLocked( + idx_t consumer_idx, const InterruptState &interrupt_state, PipelineBroadcastExchangeScanState &scan_state, + shared_ptr &next_chunk, optional_idx &batch_index, SpoolReadReservation &spool_read, + vector &writers, vector &log_entries) { D_ASSERT(consumer_idx < consumers.size()); auto &consumer = consumers[consumer_idx]; if (consumer.lifecycle != ConsumerLifecycle::ACTIVE || producer_state == ProducerState::CANCELLED) { return SourceResultType::FINISHED; } - if (consumer.read_state == ConsumerReadState::READING) { - blocked_readers.push_back(interrupt_state); - return SourceResultType::BLOCKED; - } if (consumer.position < buffer->NextPosition()) { - buffer->ReserveRead(consumer.position, consumer.shared_reader, next_chunk, spool_read); + auto position = consumer.position++; + buffer->ReserveRead(position, scan_state.spool_reader, next_chunk, batch_index, spool_read); if (spool_read.IsSet()) { - consumer.read_state = ConsumerReadState::READING; - consumer.read_position = consumer.position; + auto inserted = consumer.in_flight_reads.insert(position); + D_ASSERT(inserted.second); } else { - consumer.position++; consumer.rows_read += next_chunk->size(); RetireChunksLocked(); WakeWritersLocked(writers, log_entries); @@ -918,9 +1361,12 @@ SourceResultType PipelineBroadcastExchange::ReserveScanLocked(idx_t consumer_idx return SourceResultType::HAVE_MORE_OUTPUT; } if (producer_state == ProducerState::FINISHED) { - consumer.lifecycle = ConsumerLifecycle::INACTIVE; - D_ASSERT(active_consumers > 0); - active_consumers--; + consumer.exhausted = true; + if (consumer.in_flight_reads.empty()) { + consumer.lifecycle = ConsumerLifecycle::INACTIVE; + D_ASSERT(active_consumers > 0); + active_consumers--; + } RetireChunksLocked(); WakeWritersLocked(writers, log_entries, WriterWakeMode::FORCE); return SourceResultType::FINISHED; @@ -936,20 +1382,22 @@ void PipelineBroadcastExchange::CompleteSpoolReadLocked(idx_t consumer_idx, cons vector &log_entries) { D_ASSERT(consumer_idx < consumers.size()); auto &consumer = consumers[consumer_idx]; - D_ASSERT(consumer.read_state == ConsumerReadState::READING); - D_ASSERT(consumer.read_position == spool_read.position); - consumer.read_state = ConsumerReadState::IDLE; + auto entry = consumer.in_flight_reads.find(spool_read.position); + D_ASSERT(entry != consumer.in_flight_reads.end()); + consumer.in_flight_reads.erase(entry); if (consumer.lifecycle != ConsumerLifecycle::ACTIVE || producer_state == ProducerState::CANCELLED) { chunk.Reset(); - consumer.shared_reader.reset(); RetireChunksLocked(); WakeWritersLocked(writers, log_entries, WriterWakeMode::FORCE); WakeReadersLocked(readers); return; } - D_ASSERT(consumer.position == spool_read.position); - consumer.position++; consumer.rows_read += chunk.size(); + if (consumer.exhausted && consumer.in_flight_reads.empty()) { + consumer.lifecycle = ConsumerLifecycle::INACTIVE; + D_ASSERT(active_consumers > 0); + active_consumers--; + } RetireChunksLocked(); WakeWritersLocked(writers, log_entries); WakeReadersLocked(readers); @@ -968,11 +1416,11 @@ void PipelineBroadcastExchange::UnregisterConsumer(idx_t consumer_idx) { if (consumer.lifecycle != ConsumerLifecycle::ACTIVE) { return; } + if (consumer.exhausted) { + return; + } consumer.lifecycle = ConsumerLifecycle::INACTIVE; consumer.position = buffer->NextPosition(); - if (consumer.read_state != ConsumerReadState::READING) { - consumer.shared_reader.reset(); - } D_ASSERT(active_consumers > 0); active_consumers--; RetireChunksLocked(); @@ -1037,7 +1485,7 @@ ProgressData PipelineBroadcastExchange::SinkProgress(const ProgressData &source_ } idx_t PipelineBroadcastExchange::MaxThreads() const { - return MaxValue(max_threads, 1); + return order_mode == PipelineBroadcastExchangeOrderMode::SEQUENTIAL ? 1 : MaxValue(max_threads, 1); } idx_t PipelineBroadcastExchange::RegisteredConsumerCount() const { @@ -1101,39 +1549,22 @@ void PipelineBroadcastExchange::CreateSharedSpoolLocked(vector } void PipelineBroadcastExchange::RetireChunksLocked() { - if (buffer->HasSharedSpool()) { - idx_t min_position = buffer->NextPosition(); - bool found_reader = false; - for (auto &consumer : consumers) { - if (consumer.read_state == ConsumerReadState::READING) { - found_reader = true; - min_position = MinValue(min_position, consumer.read_position); - } - if (consumer.lifecycle == ConsumerLifecycle::ACTIVE) { - found_reader = true; - min_position = MinValue(min_position, consumer.position); - } - } - if (!found_reader) { - min_position = buffer->NextPosition(); - } - buffer->RetireBefore(min_position); - TryReleaseBufferedStorageLocked(); - return; - } if (buffer->Empty()) { return; } idx_t min_position = buffer->NextPosition(); - bool found_active = false; + bool found_reader = false; for (auto &consumer : consumers) { - if (consumer.lifecycle != ConsumerLifecycle::ACTIVE) { - continue; + if (!consumer.in_flight_reads.empty()) { + found_reader = true; + min_position = MinValue(min_position, *consumer.in_flight_reads.begin()); + } + if (consumer.lifecycle == ConsumerLifecycle::ACTIVE) { + found_reader = true; + min_position = MinValue(min_position, consumer.position); } - found_active = true; - min_position = MinValue(min_position, consumer.position); } - if (!found_active) { + if (!found_reader) { min_position = buffer->NextPosition(); } buffer->RetireBefore(min_position); @@ -1145,10 +1576,9 @@ void PipelineBroadcastExchange::TryReleaseBufferedStorageLocked() { return; } for (auto &consumer : consumers) { - if (consumer.read_state == ConsumerReadState::READING) { + if (!consumer.in_flight_reads.empty()) { return; } - consumer.shared_reader.reset(); } buffer->Release(); } @@ -1160,9 +1590,6 @@ void PipelineBroadcastExchange::DeactivateAllConsumersLocked() { } consumer.lifecycle = ConsumerLifecycle::INACTIVE; consumer.position = buffer->NextPosition(); - if (consumer.read_state != ConsumerReadState::READING) { - consumer.shared_reader.reset(); - } } active_consumers = 0; } diff --git a/src/duckdb/src/parallel/pipeline_executor.cpp b/src/duckdb/src/parallel/pipeline_executor.cpp index de3f7a4e9..c440e71e9 100644 --- a/src/duckdb/src/parallel/pipeline_executor.cpp +++ b/src/duckdb/src/parallel/pipeline_executor.cpp @@ -13,17 +13,23 @@ namespace duckdb { -PipelineExecutor::PipelineExecutor(ClientContext &context_p, Pipeline &pipeline_p) +PipelineExecutor::PipelineExecutor(ClientContext &context_p, Pipeline &pipeline_p, optional_idx reserved_batch_index) : pipeline(pipeline_p), thread(context_p), context(context_p, thread, &pipeline_p) { if (pipeline.sink) { local_sink_state = pipeline.sink->GetLocalSinkState(context); required_partition_info = pipeline.sink->RequiredPartitionInfo(); if (required_partition_info.AnyRequired()) { - D_ASSERT(pipeline.source->SupportsPartitioning(OperatorPartitionInfo::BatchIndex())); + D_ASSERT(pipeline.source->SupportsPartitioning(required_partition_info)); auto &partition_info = local_sink_state->partition_info; D_ASSERT(!partition_info.batch_index.IsValid()); // batch index is not set yet - initialize before fetching anything - partition_info.batch_index = pipeline.RegisterNewBatchIndex(); + if (pipeline.IsExternalInput()) { + D_ASSERT(!reserved_batch_index.IsValid()); + partition_info.batch_index = pipeline.GetBaseBatchIndex(); + } else { + partition_info.batch_index = + reserved_batch_index.IsValid() ? reserved_batch_index.GetIndex() : pipeline.RegisterNewBatchIndex(); + } partition_info.min_batch_index = partition_info.batch_index; } } @@ -65,6 +71,7 @@ void PipelineExecutor::Reset() { done_flushing = false; remaining_sink_chunk = false; next_batch_blocked = false; + external_batch_initialized = false; finished_processing_idx = -1; source_profiling_finalized = false; source_finish_notification_state = SourceFinishNotificationState::PENDING; @@ -83,10 +90,11 @@ void PipelineExecutor::Reset() { required_partition_info = pipeline.sink->RequiredPartitionInfo(); local_sink_state->partition_info = SourcePartitionInfo(); if (required_partition_info.AnyRequired()) { - D_ASSERT(pipeline.source->SupportsPartitioning(OperatorPartitionInfo::BatchIndex())); + D_ASSERT(pipeline.source->SupportsPartitioning(required_partition_info)); auto &partition_info = local_sink_state->partition_info; D_ASSERT(!partition_info.batch_index.IsValid()); - partition_info.batch_index = pipeline.RegisterNewBatchIndex(); + partition_info.batch_index = + pipeline.IsExternalInput() ? pipeline.GetBaseBatchIndex() : pipeline.RegisterNewBatchIndex(); partition_info.min_batch_index = partition_info.batch_index; } } @@ -214,20 +222,55 @@ SinkNextBatchType PipelineExecutor::NextBatch(DataChunk &source_chunk, const boo D_ASSERT(local_source_state); D_ASSERT(global_source_state); // if we retrieved data - initialize the next batch index - auto partition_data = pipeline.source->GetPartitionData(context, source_chunk, *global_source_state, - *local_source_state, required_partition_info); - auto batch_index = partition_data.batch_index; - // we start with the base_batch_index as a valid starting value. Make sure that next batch is called below - next_data = std::move(partition_data); - next_data.batch_index = pipeline.base_batch_index + batch_index + 1; - if (next_data.batch_index >= max_batch_index) { - throw InternalException("Pipeline batch index - invalid batch index %llu returned by source operator", - batch_index); - } + auto source_data = pipeline.source->GetPartitionData(context, source_chunk, *global_source_state, + *local_source_state, required_partition_info); + next_data = ToPipelinePartitionData(source_data); } else if (have_more_output) { next_data.batch_index = partition_info.batch_index.GetIndex(); } - if (next_data.batch_index == partition_info.batch_index.GetIndex()) { + return NextBatch(std::move(next_data)); +} + +OperatorPartitionData PipelineExecutor::ToPipelinePartitionData(const OperatorPartitionData &source_data) const { + auto max_batch_index = pipeline.base_batch_index + PipelineBuildState::BATCH_INCREMENT - 1; + OperatorPartitionData result(pipeline.base_batch_index + source_data.batch_index + 1); + result.partition_data = source_data.partition_data; + if (result.batch_index >= max_batch_index) { + throw InternalException("Pipeline batch index - invalid batch index %llu returned by source operator", + source_data.batch_index); + } + return result; +} + +SinkNextBatchType PipelineExecutor::NextBatch(OperatorPartitionData next_data, bool force, + optional_idx external_min_batch_index) { + auto &partition_info = local_sink_state->partition_info; + optional_idx mapped_external_min_batch_index; + if (pipeline.IsExternalInput()) { + if (!external_min_batch_index.IsValid()) { + throw InternalException("External pipeline did not provide a minimum batch index"); + } + auto min_batch_offset = external_min_batch_index.GetIndex(); + if (min_batch_offset >= PipelineBuildState::BATCH_INCREMENT) { + throw InternalException("External pipeline minimum batch index is outside its pipeline"); + } + auto min_batch_index = pipeline.GetBaseBatchIndex() + min_batch_offset; + if (min_batch_index > next_data.batch_index) { + throw InternalException("External pipeline minimum batch index %llu exceeds current batch index %llu", + min_batch_index, next_data.batch_index); + } + if (min_batch_index < partition_info.min_batch_index.GetIndex()) { + throw InternalException("External pipeline minimum batch index decreased from %llu to %llu", + partition_info.min_batch_index.GetIndex(), min_batch_index); + } + mapped_external_min_batch_index = min_batch_index; + } else if (external_min_batch_index.IsValid()) { + throw InternalException("Source-driven pipeline received an external minimum batch index"); + } + if (!force && next_data.batch_index == partition_info.batch_index.GetIndex()) { + if (mapped_external_min_batch_index.IsValid()) { + partition_info.min_batch_index = mapped_external_min_batch_index; + } // no changes, return return SinkNextBatchType::READY; } @@ -263,7 +306,11 @@ SinkNextBatchType PipelineExecutor::NextBatch(DataChunk &source_chunk, const boo return SinkNextBatchType::BLOCKED; } - partition_info.min_batch_index = pipeline.UpdateBatchIndex(current_batch, next_data.batch_index); + if (mapped_external_min_batch_index.IsValid()) { + partition_info.min_batch_index = mapped_external_min_batch_index; + } else { + partition_info.min_batch_index = pipeline.UpdateBatchIndex(current_batch, next_data.batch_index); + } return SinkNextBatchType::READY; } @@ -293,7 +340,7 @@ PipelineExecuteResult PipelineExecutor::Execute(idx_t max_chunks) { // the operators have to be called with the same input chunk to produce the rest of the output D_ASSERT(source_chunk.size() > 0); result = ExecutePushInternal(source_chunk, chunk_budget); - } else if (exhausted_pipeline && !next_batch_blocked && !done_flushing) { + } else if (exhausted_pipeline && (!next_batch_blocked || done_flushing)) { // The pipeline was exhausted, try flushing all operators return FlushAndFinalize(chunk_budget); } else if (!exhausted_pipeline || next_batch_blocked) { @@ -311,7 +358,8 @@ PipelineExecuteResult PipelineExecutor::Execute(idx_t max_chunks) { } } - if (required_partition_info.AnyRequired()) { + if (required_partition_info.AnyRequired() && + (source_result != SourceResultType::FINISHED || source_chunk.size() > 0)) { auto next_batch_result = NextBatch(source_chunk, source_result == SourceResultType::HAVE_MORE_OUTPUT); next_batch_blocked = next_batch_result == SinkNextBatchType::BLOCKED; if (next_batch_blocked) { @@ -355,13 +403,21 @@ PipelineExecuteResult PipelineExecutor::Execute() { return Execute(NumericLimits::Maximum()); } -PipelineExecuteResult PipelineExecutor::PushExternal(DataChunk &input) { +PipelineExecuteResult PipelineExecutor::PushExternal(DataChunk &input, + const OperatorPartitionData &source_partition_data, + optional_idx source_min_batch_index) { D_ASSERT(pipeline.sink); D_ASSERT(pipeline.IsExternalInput()); if (IsFinished()) { return PipelineExecuteResult::FINISHED; } - if (!remaining_sink_chunk) { + if (required_partition_info.AnyRequired() && !remaining_sink_chunk) { + auto next_batch_result = NextBatchExternal(source_partition_data, source_min_batch_index); + if (next_batch_result != PipelineExecuteResult::NOT_FINISHED) { + return next_batch_result; + } + } + if (!remaining_sink_chunk && !next_batch_blocked) { context.thread.profiler.StartOperator(pipeline.source.get()); context.thread.profiler.EndOperator(&input); } @@ -370,12 +426,36 @@ PipelineExecuteResult PipelineExecutor::PushExternal(DataChunk &input) { return PushInputChunk(input, chunk_budget, PipelineInputChunkMode::PUSH_INPUT); } -PipelineExecuteResult PipelineExecutor::FinishExternal() { +PipelineExecuteResult PipelineExecutor::NextBatchExternal(const OperatorPartitionData &source_partition_data, + optional_idx source_min_batch_index) { + D_ASSERT(pipeline.sink); + D_ASSERT(pipeline.IsExternalInput()); + if (IsFinished()) { + return PipelineExecuteResult::FINISHED; + } + if (!required_partition_info.AnyRequired()) { + return PipelineExecuteResult::NOT_FINISHED; + } + auto next_batch_result = + NextBatch(ToPipelinePartitionData(source_partition_data), !external_batch_initialized, source_min_batch_index); + next_batch_blocked = next_batch_result == SinkNextBatchType::BLOCKED; + if (next_batch_blocked) { + return PipelineExecuteResult::INTERRUPTED; + } + external_batch_initialized = true; + return PipelineExecuteResult::NOT_FINISHED; +} + +PipelineExecuteResult PipelineExecutor::FinishExternal(optional_idx source_min_batch_index) { D_ASSERT(pipeline.sink); D_ASSERT(pipeline.IsExternalInput()); if (finalized) { return PipelineExecuteResult::FINISHED; } + auto finish_batch_result = FinishBatchExternal(source_min_batch_index); + if (finish_batch_result == PipelineExecuteResult::INTERRUPTED) { + return finish_batch_result; + } ExecutionBudget chunk_budget(NumericLimits::Maximum()); exhausted_source = true; @@ -383,6 +463,25 @@ PipelineExecuteResult PipelineExecutor::FinishExternal() { return FlushAndFinalize(chunk_budget); } +PipelineExecuteResult PipelineExecutor::FinishBatchExternal(optional_idx source_min_batch_index) { + D_ASSERT(pipeline.sink); + D_ASSERT(pipeline.IsExternalInput()); + if (finalized) { + return PipelineExecuteResult::FINISHED; + } + if (required_partition_info.AnyRequired()) { + auto max_batch_index = pipeline.base_batch_index + PipelineBuildState::BATCH_INCREMENT - 1; + auto next_batch_result = + NextBatch(OperatorPartitionData(max_batch_index), !external_batch_initialized, source_min_batch_index); + next_batch_blocked = next_batch_result == SinkNextBatchType::BLOCKED; + if (next_batch_blocked) { + return PipelineExecuteResult::INTERRUPTED; + } + external_batch_initialized = true; + } + return PipelineExecuteResult::NOT_FINISHED; +} + bool PipelineExecutor::IsFinishedProcessing() const { return IsFinished(); } @@ -521,6 +620,14 @@ PipelineExecuteResult PipelineExecutor::FlushAndFinalize(ExecutionBudget &chunk_ } done_flushing = true; } + if (required_partition_info.AnyRequired() && !pipeline.IsExternalInput()) { + DataChunk empty_chunk; + auto next_batch_result = NextBatch(empty_chunk, false); + next_batch_blocked = next_batch_result == SinkNextBatchType::BLOCKED; + if (next_batch_blocked) { + return PipelineExecuteResult::INTERRUPTED; + } + } return PushFinalize(); } diff --git a/src/duckdb/src/planner/binder/statement/bind_export.cpp b/src/duckdb/src/planner/binder/statement/bind_export.cpp index 87f332f8d..60a5e1c2d 100644 --- a/src/duckdb/src/planner/binder/statement/bind_export.cpp +++ b/src/duckdb/src/planner/binder/statement/bind_export.cpp @@ -167,6 +167,10 @@ BoundStatement Binder::Bind(ExportStatement &stmt) { catalog_entry_vector_t tables; auto schemas = Catalog::GetSchemas(context, catalog); for (auto &schema : schemas) { + auto &schema_entry = schema.get(); + if (schema_entry.ParentCatalog().IsTemporaryCatalog()) { + continue; + } schema.get().Scan(context, CatalogType::TABLE_ENTRY, [&](CatalogEntry &entry) { if (entry.type == CatalogType::TABLE_ENTRY) { tables.push_back(entry.Cast()); diff --git a/src/duckdb/src/planner/operator/logical_column_data_get.cpp b/src/duckdb/src/planner/operator/logical_column_data_get.cpp index cb20528bb..a96f04d8b 100644 --- a/src/duckdb/src/planner/operator/logical_column_data_get.cpp +++ b/src/duckdb/src/planner/operator/logical_column_data_get.cpp @@ -1,16 +1,29 @@ #include "duckdb/planner/operator/logical_column_data_get.hpp" +#include "duckdb/common/exception.hpp" +#include "duckdb/common/serializer/deserializer.hpp" +#include "duckdb/common/serializer/serializer.hpp" #include "duckdb/common/types/data_chunk.hpp" #include "duckdb/main/config.hpp" namespace duckdb { +static vector GenerateColumnDataColumnIds(idx_t column_count) { + vector column_ids; + column_ids.reserve(column_count); + for (idx_t i = 0; i < column_count; i++) { + column_ids.push_back(i); + } + return column_ids; +} + LogicalColumnDataGet::LogicalColumnDataGet(TableIndex table_index, vector types, unique_ptr collection_p) : LogicalOperator(LogicalOperatorType::LOGICAL_CHUNK_GET), table_index(table_index), collection(std::move(collection_p)) { D_ASSERT(!types.empty()); chunk_types = std::move(types); + SetColumnIds(GenerateColumnDataColumnIds(chunk_types.size())); } LogicalColumnDataGet::LogicalColumnDataGet(TableIndex table_index, vector types, @@ -18,6 +31,7 @@ LogicalColumnDataGet::LogicalColumnDataGet(TableIndex table_index, vector types, @@ -26,10 +40,26 @@ LogicalColumnDataGet::LogicalColumnDataGet(TableIndex table_index, vector column_ids_p) { + for (auto column_id : column_ids_p) { + if (column_id >= chunk_types.size()) { + throw SerializationException("LogicalColumnDataGet column id %llu is out of bounds for %llu chunk types", + column_id, chunk_types.size()); + } + } + column_ids = std::move(column_ids_p); + ResolveTypes(); +} + +const vector &LogicalColumnDataGet::GetColumnIds() const { + return column_ids; } vector LogicalColumnDataGet::GetColumnBindings() { - return GenerateColumnBindings(table_index, chunk_types.size()); + return GenerateColumnBindings(table_index, column_ids.size()); } vector LogicalColumnDataGet::GetTableIndex() const { @@ -45,4 +75,28 @@ string LogicalColumnDataGet::GetName() const { return LogicalOperator::GetName(); } +void LogicalColumnDataGet::Serialize(Serializer &serializer) const { + LogicalOperator::Serialize(serializer); + serializer.WritePropertyWithDefault(200, "table_index", table_index); + serializer.WritePropertyWithDefault>(201, "chunk_types", chunk_types); + serializer.WritePropertyWithDefault>(202, "collection", collection); + serializer.WritePropertyWithDefault>(203, "column_ids", column_ids, + GenerateColumnDataColumnIds(chunk_types.size())); +} + +unique_ptr LogicalColumnDataGet::Deserialize(Deserializer &deserializer) { + auto table_index = deserializer.ReadPropertyWithDefault(200, "table_index"); + auto chunk_types = deserializer.ReadPropertyWithDefault>(201, "chunk_types"); + auto collection = + deserializer.ReadPropertyWithDefault>(202, "collection"); + auto result = duckdb::unique_ptr( + new LogicalColumnDataGet(table_index, std::move(chunk_types), std::move(collection))); + if (deserializer.CanDeserializeProperty(203, "column_ids")) { + vector column_ids; + deserializer.ReadProperty(203, "column_ids", column_ids); + result->SetColumnIds(std::move(column_ids)); + } + return std::move(result); +} + } // namespace duckdb diff --git a/src/duckdb/src/storage/serialization/serialize_logical_operator.cpp b/src/duckdb/src/storage/serialization/serialize_logical_operator.cpp index a18f6e794..fde83d3bc 100644 --- a/src/duckdb/src/storage/serialization/serialize_logical_operator.cpp +++ b/src/duckdb/src/storage/serialization/serialize_logical_operator.cpp @@ -347,21 +347,6 @@ unique_ptr LogicalCTERef::Deserialize(Deserializer &deserialize return std::move(result); } -void LogicalColumnDataGet::Serialize(Serializer &serializer) const { - LogicalOperator::Serialize(serializer); - serializer.WritePropertyWithDefault(200, "table_index", table_index); - serializer.WritePropertyWithDefault>(201, "chunk_types", chunk_types); - serializer.WritePropertyWithDefault>(202, "collection", collection); -} - -unique_ptr LogicalColumnDataGet::Deserialize(Deserializer &deserializer) { - auto table_index = deserializer.ReadPropertyWithDefault(200, "table_index"); - auto chunk_types = deserializer.ReadPropertyWithDefault>(201, "chunk_types"); - auto collection = deserializer.ReadPropertyWithDefault>(202, "collection"); - auto result = duckdb::unique_ptr(new LogicalColumnDataGet(table_index, std::move(chunk_types), std::move(collection))); - return std::move(result); -} - void LogicalComparisonJoin::Serialize(Serializer &serializer) const { LogicalOperator::Serialize(serializer); serializer.WriteProperty(200, "join_type", join_type); diff --git a/src/duckdb/src/storage/wal_replay.cpp b/src/duckdb/src/storage/wal_replay.cpp index 40085f12d..00f321c43 100644 --- a/src/duckdb/src/storage/wal_replay.cpp +++ b/src/duckdb/src/storage/wal_replay.cpp @@ -62,7 +62,8 @@ class ReplayState { struct ReplayIndexInfo { ReplayIndexInfo(TableIndexList &index_list, unique_ptr index, const Identifier &table_schema, const Identifier &table_name) - : index_list(index_list), index(std::move(index)), table_schema(table_schema), table_name(table_name) {}; + : index_list(index_list), index(std::move(index)), table_schema(table_schema), table_name(table_name) { + } reference index_list; unique_ptr index; @@ -796,6 +797,16 @@ void WriteAheadLogDeserializer::ReplayIndexData(IndexStorageInfo &info) { list.ReadElement(data_ptr, data_info.allocation_sizes[j]); + // For read-only mode, retain the transient block handle and release the pin held by the buffer handle. The + // buffer can then be evicted to temporary storage until the index is bound. + if (db.IsReadOnly()) { + if (!data_info.transient_block_handles) { + data_info.transient_block_handles = make_shared_ptr>>(); + } + data_info.transient_block_handles->push_back(std::move(block_handle)); + continue; + } + // Convert the buffer handle to a persistent block and store the block id. if (!deserialize_only) { auto block_id = block_manager->GetFreeBlockIdForCheckpoint();