From 4ad28275a6b0c94feb759ffcc22413b6c801855a Mon Sep 17 00:00:00 2001 From: "xiaolei.zl" Date: Fri, 10 Jul 2026 18:24:27 +0800 Subject: [PATCH 1/8] feat: bulk build terminal CSV COPY inputs try parallel bulk loading optimize for single --- .../function/import/csv_read_function.h | 29 +- .../neug/compiler/function/read_function.h | 10 +- .../execute/ops/batch/batch_insert_edge.h | 21 + .../execute/ops/batch/batch_insert_vertex.h | 21 + .../execute/ops/batch/batch_update_utils.h | 10 + include/neug/storages/csr/mutable_csr.h | 276 ++++ include/neug/storages/graph/edge_table.h | 17 + include/neug/storages/graph/graph_interface.h | 11 + include/neug/storages/graph/property_graph.h | 10 + include/neug/storages/graph/vertex_table.h | 34 + .../storages/loader/chunk_pipeline_utils.h | 277 ++++ include/neug/storages/loader/loader_utils.h | 80 +- .../neug/utils/io/read/csv/csv_read_config.h | 1 + include/neug/utils/io/read/csv/csv_reader.h | 5 + .../execute/ops/batch/batch_insert_edge.cc | 155 ++- .../execute/ops/batch/batch_insert_vertex.cc | 150 +- .../execute/ops/batch/batch_update_utils.cc | 134 ++ src/execution/execute/plan_parser.cc | 4 + src/storages/graph/edge_table.cc | 1192 +++++++++++++++- src/storages/graph/graph_interface.cc | 32 + src/storages/graph/property_graph.cc | 45 + src/storages/graph/vertex_table.cc | 149 +- src/storages/loader/loader_utils.cc | 883 +++++++++++- src/utils/io/read/common/options.cc | 1 + src/utils/io/read/csv/csv_reader.cc | 29 + tests/storage/test_copy_temp.cc | 83 ++ tests/storage/test_edge_table.cc | 1226 ++++++++++++++++- tests/storage/test_mutable_csr.cc | 69 + tests/storage/test_vertex_table.cc | 24 + tests/unittest/utils.h | 38 + tests/utils/test_reader.cc | 334 +++++ 31 files changed, 5268 insertions(+), 82 deletions(-) create mode 100644 include/neug/storages/loader/chunk_pipeline_utils.h diff --git a/include/neug/compiler/function/import/csv_read_function.h b/include/neug/compiler/function/import/csv_read_function.h index 9a55756f2..1c9be2f9d 100644 --- a/include/neug/compiler/function/import/csv_read_function.h +++ b/include/neug/compiler/function/import/csv_read_function.h @@ -36,6 +36,7 @@ struct CSVReadFunction { std::vector{common::DataTypeId::kVarchar}; auto readFunction = std::make_unique(name, typeIDs); readFunction->execFunc = execFunc; + readFunction->sourceFunc = sourceFunc; readFunction->sniffFunc = sniffFunc; function_set functionSet; functionSet.push_back(std::move(readFunction)); @@ -132,6 +133,32 @@ struct CSVReadFunction { return ctx; } + static std::shared_ptr sourceFunc( + std::shared_ptr state) { + if (!state) { + THROW_INVALID_ARGUMENT_EXCEPTION("State is null"); + } + // sourceFunc is speculative: a caller may create a source only to decide + // that the destination is ineligible, then execute the legacy path. Keep + // the original state pristine for that fallback. + auto source_state = std::make_shared(*state); + validateAndConvertExecOptions(source_state); + const auto& vfs = neug::main::MetadataRegistry::getVFS(); + const auto& fs = vfs->Provide(source_state->schema.file); + std::vector resolved_paths; + for (const auto& path : source_state->schema.file.paths) { + const auto& resolved = fs->glob(path); + resolved_paths.insert(resolved_paths.end(), resolved.begin(), + resolved.end()); + } + source_state->schema.file.paths = std::move(resolved_paths); + auto options_builder = + std::make_unique(source_state); + auto reader = std::make_unique( + source_state, std::move(options_builder)); + return reader->createChunkSource(); + } + static std::shared_ptr sniffFunc( const reader::FileSchema& schema) { auto state = std::make_shared(); @@ -184,4 +211,4 @@ struct CSVReadFunction { } }; } // namespace function -} // namespace neug \ No newline at end of file +} // namespace neug diff --git a/include/neug/compiler/function/read_function.h b/include/neug/compiler/function/read_function.h index cfc8c74e1..e0253ffc8 100644 --- a/include/neug/compiler/function/read_function.h +++ b/include/neug/compiler/function/read_function.h @@ -27,6 +27,7 @@ #include "neug/utils/io/reader.h" namespace neug { +class IDataChunkSource; namespace function { // The exec function invoked by data source operators to load data from external @@ -34,6 +35,12 @@ namespace function { using read_exec_func_t = std::function state)>; +/// Creates a repeatable source for an opt-in bulk ingestion fast path. The +/// function must leave its input state unchanged because callers may fall back +/// to execFunc when the source or destination is not eligible for bulk build. +using read_source_func_t = std::function( + std::shared_ptr state)>; + // The function used to sniff/infer file column names and their types from // external data sources. using read_sniff_func_t = std::function( @@ -41,10 +48,11 @@ using read_sniff_func_t = std::function( struct ReadFunction : public TableFunction { read_exec_func_t execFunc = nullptr; + read_source_func_t sourceFunc = nullptr; read_sniff_func_t sniffFunc = nullptr; ReadFunction(std::string name, std::vector inputTypes) : TableFunction{std::move(name), std::move(inputTypes)} {} }; } // namespace function -} // namespace neug \ No newline at end of file +} // namespace neug diff --git a/include/neug/execution/execute/ops/batch/batch_insert_edge.h b/include/neug/execution/execute/ops/batch/batch_insert_edge.h index 12d7e9118..f4219ae84 100644 --- a/include/neug/execution/execute/ops/batch/batch_insert_edge.h +++ b/include/neug/execution/execute/ops/batch/batch_insert_edge.h @@ -40,6 +40,27 @@ class BatchInsertEdgeOprBuilder : public IOperatorBuilder { } }; +/// Fuses only a terminal, empty-sink COPY FROM plan. The implementation +/// reverts to the normal reader/Context path unless runtime bulk eligibility +/// is established. +class BatchInsertEdgeFromSourceOprBuilder : public IOperatorBuilder { + public: + BatchInsertEdgeFromSourceOprBuilder() = default; + ~BatchInsertEdgeFromSourceOprBuilder() = default; + + neug::result Build(const Schema& schema, + const ContextMeta& ctx_meta, + const physical::PhysicalPlan& plan, + int op_idx) override; + + std::vector GetOpKinds() + const override { + return {physical::PhysicalOpr_Operator::OpKindCase::kSource, + physical::PhysicalOpr_Operator::OpKindCase::kLoadEdge, + physical::PhysicalOpr_Operator::OpKindCase::kSink}; + } +}; + } // namespace ops } // namespace execution } // namespace neug diff --git a/include/neug/execution/execute/ops/batch/batch_insert_vertex.h b/include/neug/execution/execute/ops/batch/batch_insert_vertex.h index 77946a43f..798ad59fd 100644 --- a/include/neug/execution/execute/ops/batch/batch_insert_vertex.h +++ b/include/neug/execution/execute/ops/batch/batch_insert_vertex.h @@ -39,6 +39,27 @@ class BatchInsertVertexOprBuilder : public IOperatorBuilder { } }; +/// Fuses only a terminal, empty-sink COPY FROM plan. The implementation +/// reverts to the normal reader/Context path unless runtime bulk eligibility +/// is established. +class BatchInsertVertexFromSourceOprBuilder : public IOperatorBuilder { + public: + BatchInsertVertexFromSourceOprBuilder() = default; + ~BatchInsertVertexFromSourceOprBuilder() = default; + + neug::result Build(const Schema& schema, + const ContextMeta& ctx_meta, + const physical::PhysicalPlan& plan, + int op_idx) override; + + std::vector GetOpKinds() + const override { + return {physical::PhysicalOpr_Operator::OpKindCase::kSource, + physical::PhysicalOpr_Operator::OpKindCase::kLoadVertex, + physical::PhysicalOpr_Operator::OpKindCase::kSink}; + } +}; + } // namespace ops } // namespace execution } // namespace neug diff --git a/include/neug/execution/execute/ops/batch/batch_update_utils.h b/include/neug/execution/execute/ops/batch/batch_update_utils.h index 70ba91dec..af53624f0 100644 --- a/include/neug/execution/execute/ops/batch/batch_update_utils.h +++ b/include/neug/execution/execute/ops/batch/batch_update_utils.h @@ -32,6 +32,7 @@ class RepeatedPtrField; namespace neug { class IDataChunkSupplier; +class IDataChunkSource; class Schema; class StorageReadInterface; namespace execution { @@ -65,6 +66,15 @@ std::shared_ptr create_data_chunk_supplier( const Context& ctx, const std::vector>& prop_mappings); +std::shared_ptr create_data_chunk_source( + std::shared_ptr source, + const std::vector>& prop_mappings); + +/// Selects staged bulk builders only for sources large enough to amortize their +/// two bounded parsing passes. Set NEUG_COPY_BULK_BUILD=true/false to force a +/// decision while benchmarking or rolling back. +bool should_use_copy_bulk_build(const IDataChunkSource& source); + std::vector match_files_with_pattern(const std::string& file_path); std::vector> create_csv_chunk_suppliers( diff --git a/include/neug/storages/csr/mutable_csr.h b/include/neug/storages/csr/mutable_csr.h index 61a685502..892456216 100644 --- a/include/neug/storages/csr/mutable_csr.h +++ b/include/neug/storages/csr/mutable_csr.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -31,6 +32,7 @@ #include #include +#include "neug/config.h" #include "neug/storages/allocators.h" #include "neug/storages/container/i_container.h" #include "neug/storages/csr/csr_base.h" @@ -50,6 +52,12 @@ static_assert( sizeof(std::atomic) == sizeof(int), "atomic must have the same size as int on supported platforms"); +template +class MutableCsrBulkBuildAccess; + +template +class SingleMutableCsrBulkBuildAccess; + template class MutableCsr : public TypedCsrBase { public: @@ -235,6 +243,8 @@ class MutableCsr : public TypedCsrBase { } private: + friend class MutableCsrBulkBuildAccess; + std::unique_ptr locks_; std::shared_ptr adj_list_buffer_; std::shared_ptr degree_list_; @@ -382,6 +392,8 @@ class SingleMutableCsr : public TypedCsrBase { } private: + friend class SingleMutableCsrBulkBuildAccess; + std::shared_ptr nbr_list_; std::atomic edge_num_{0}; CsrPrefetchPolicy prefetch_policy_; @@ -480,4 +492,268 @@ class EmptyCsr : public TypedCsrBase { } }; +/// Internal append-only access used while a fresh CSR is being bulk built. +/// The CSR is not published until both passes complete, so this bypasses the +/// incremental-growth path used by transactional updates. +template +class MutableCsrBulkBuildAccess { + public: + using csr_t = MutableCsr; + using nbr_t = typename csr_t::nbr_t; + static constexpr bool kSupportsDisjointConcurrentFill = true; + static constexpr bool kStoresEdges = true; + static constexpr bool kNeedsDegreeCount = true; + static constexpr bool kChecksSingleUniqueness = false; + static constexpr bool kNeedsConcurrentGrouping = true; + static constexpr bool kTracksInputEdgeCount = false; + + explicit MutableCsrBulkBuildAccess(csr_t& csr) : csr_(csr) {} + + void PrepareBuild(vid_t vertex_count) { + CHECK(csr_.adj_list_buffer_ != nullptr); + CHECK(csr_.degree_list_ != nullptr); + CHECK(csr_.cap_list_ != nullptr); + CHECK(csr_.nbr_list_ != nullptr); + csr_.adj_list_buffer_->Resize(static_cast(vertex_count) * + sizeof(nbr_t*)); + csr_.degree_list_->Resize(static_cast(vertex_count) * sizeof(int)); + csr_.cap_list_->Resize(static_cast(vertex_count) * sizeof(int)); + csr_.locks_ = std::make_unique(vertex_count); + refresh_metadata_ptrs(); + CHECK(adj_lists_ != nullptr || vertex_count == 0); + CHECK(degrees_ != nullptr || vertex_count == 0); + CHECK(capacities_ != nullptr || vertex_count == 0); + for (vid_t i = 0; i < vertex_count; ++i) { + adj_lists_[i] = nullptr; + degrees_[i].store(0, std::memory_order_relaxed); + capacities_[i] = 0; + } + csr_.edge_num_.store(0, std::memory_order_relaxed); + csr_.unsorted_since_ = 0; + } + + void CountSerial(vid_t src, int count = 1) { + CHECK_LT(src, vertex_capacity_); + CHECK_GT(count, 0); + auto degree = degrees_[src].load(std::memory_order_relaxed); + CHECK_LE(degree, std::numeric_limits::max() - count); + degrees_[src].store(degree + count, std::memory_order_relaxed); + } + + void CountConcurrent(vid_t src, int count = 1) { + CHECK_LT(src, vertex_capacity_); + CHECK_GT(count, 0); + const auto degree = + degrees_[src].fetch_add(count, std::memory_order_relaxed); + CHECK_LE(degree, std::numeric_limits::max() - count); + } + + void AllocateFromCounts() { + refresh_metadata_ptrs(); + size_t total_capacity = 0; + for (size_t i = 0; i < vertex_capacity_; ++i) { + const auto degree = degrees_[i].load(std::memory_order_relaxed); + const auto reserved_capacity = ReservedCapacity(degree); + // During the fill pass cap_list_ stores the exact expected degree. This + // lets range reservations validate both passes without allocating a + // second O(V) metadata array. Finish() converts it to runtime capacity. + capacities_[i] = degree; + CHECK_LE(static_cast(reserved_capacity), + std::numeric_limits::max() - total_capacity); + total_capacity += static_cast(reserved_capacity); + } + CHECK_LE(total_capacity, + std::numeric_limits::max() / sizeof(nbr_t)); + csr_.nbr_list_->Resize(total_capacity * sizeof(nbr_t)); + nbrs_ = reinterpret_cast(csr_.nbr_list_->GetData()); + CHECK(nbrs_ != nullptr || total_capacity == 0); + + size_t offset = 0; + for (size_t i = 0; i < vertex_capacity_; ++i) { + const auto reserved_capacity = ReservedCapacity(capacities_[i]); + adj_lists_[i] = reserved_capacity == 0 ? nullptr : nbrs_ + offset; + offset += static_cast(reserved_capacity); + degrees_[i].store(0, std::memory_order_relaxed); + } + } + + int ExpectedDegree(vid_t src) const { + CHECK_LT(src, vertex_capacity_); + return capacities_[src]; + } + + int ReservedCapacityForVertex(vid_t src) const { + CHECK_LT(src, vertex_capacity_); + return ReservedCapacity(capacities_[src]); + } + + int ReserveSerial(vid_t src, int count) { + CHECK_LT(src, vertex_capacity_); + CHECK_GT(count, 0); + const auto slot = degrees_[src].load(std::memory_order_relaxed); + CHECK_LE(slot, capacities_[src] - count); + degrees_[src].store(slot + count, std::memory_order_relaxed); + return slot; + } + + int ReserveConcurrent(vid_t src, int count) { + CHECK_LT(src, vertex_capacity_); + CHECK_GT(count, 0); + CHECK_LE(count, capacities_[src]); + const auto slot = degrees_[src].fetch_add(count, std::memory_order_relaxed); + CHECK_LE(slot, capacities_[src] - count); + return slot; + } + + void PutAt(vid_t src, int slot, vid_t dst, const EDATA_T& data, + timestamp_t ts) { + CHECK_LT(src, vertex_capacity_); + CHECK_GE(slot, 0); + CHECK_LT(slot, capacities_[src]); + auto& nbr = adj_lists_[src][slot]; + nbr.neighbor = dst; + nbr.data = data; + nbr.timestamp.store(ts, std::memory_order_relaxed); + } + + void PutSerial(vid_t src, vid_t dst, const EDATA_T& data, timestamp_t ts) { + PutAt(src, ReserveSerial(src, 1), dst, data, ts); + } + + void Finish() { + uint64_t edge_num = 0; + for (size_t i = 0; i < vertex_capacity_; ++i) { + const auto degree = degrees_[i].load(std::memory_order_relaxed); + CHECK_EQ(degree, capacities_[i]) + << "Bulk edge count/fill mismatch for vertex " << i; + edge_num += static_cast(degree); + capacities_[i] = ReservedCapacity(degree); + } + csr_.edge_num_.store(edge_num, std::memory_order_relaxed); + csr_.refresh_prefetch_policy(); + } + + private: + static int ReservedCapacity(int degree) { + CHECK_GE(degree, 0); + if (degree == 0) { + return 0; + } + const auto reserved = + std::ceil(degree * NeugDBConfig::DEFAULT_RESERVE_RATIO); + CHECK_LE(reserved, static_cast(std::numeric_limits::max())); + return static_cast(reserved); + } + + void refresh_metadata_ptrs() { + vertex_capacity_ = csr_.vertex_capacity(); + adj_lists_ = reinterpret_cast( + csr_.adj_list_buffer_ == nullptr ? nullptr + : csr_.adj_list_buffer_->GetData()); + degrees_ = reinterpret_cast*>( + csr_.degree_list_ == nullptr ? nullptr : csr_.degree_list_->GetData()); + capacities_ = reinterpret_cast( + csr_.cap_list_ == nullptr ? nullptr : csr_.cap_list_->GetData()); + } + + csr_t& csr_; + size_t vertex_capacity_ = 0; + nbr_t** adj_lists_ = nullptr; + std::atomic* degrees_ = nullptr; + int* capacities_ = nullptr; + nbr_t* nbrs_ = nullptr; +}; + +/// Append-only writer for single-edge CSR layouts. A preceding degree pass can +/// use the timestamp field as a build-private seen marker. If every endpoint is +/// unique, workers can then fill fixed vertex slots concurrently without +/// locks. Duplicate endpoints require the ordered serial fill path to preserve +/// the pre-existing last-write-wins behavior. +template +class SingleMutableCsrBulkBuildAccess { + public: + using csr_t = SingleMutableCsr; + using nbr_t = typename csr_t::nbr_t; + static constexpr bool kSupportsDisjointConcurrentFill = true; + static constexpr bool kStoresEdges = true; + static constexpr bool kNeedsDegreeCount = false; + static constexpr bool kChecksSingleUniqueness = true; + static constexpr bool kNeedsConcurrentGrouping = false; + static constexpr bool kTracksInputEdgeCount = true; + + explicit SingleMutableCsrBulkBuildAccess(csr_t& csr) : csr_(csr) {} + + void PrepareBuild(vid_t vertex_count) { + CHECK(csr_.nbr_list_ != nullptr); + csr_.nbr_list_->Resize(static_cast(vertex_count) * sizeof(nbr_t)); + refresh_ptrs(); + CHECK(nbrs_ != nullptr || vertex_count == 0); + for (vid_t i = 0; i < vertex_count; ++i) { + nbrs_[i].timestamp.store(INVALID_TIMESTAMP, std::memory_order_relaxed); + } + csr_.edge_num_.store(0, std::memory_order_relaxed); + edge_count_ = 0; + input_edge_count_precomputed_ = false; + } + + void AllocateFromCounts() {} + + bool CheckUniqueSerial(vid_t src) { + CHECK_LT(src, vertex_capacity_); + auto& timestamp = nbrs_[src].timestamp; + const auto previous = timestamp.load(std::memory_order_relaxed); + timestamp.store(0, std::memory_order_relaxed); + return previous != INVALID_TIMESTAMP; + } + + bool CheckUniqueConcurrent(vid_t src) { + CHECK_LT(src, vertex_capacity_); + const auto previous = + nbrs_[src].timestamp.exchange(0, std::memory_order_relaxed); + return previous != INVALID_TIMESTAMP; + } + + void SetInputEdgeCount(uint64_t count) { + edge_count_ = count; + input_edge_count_precomputed_ = true; + } + + void PutSerial(vid_t src, vid_t dst, const EDATA_T& data, timestamp_t ts) { + PutConcurrent(src, dst, data, ts); + } + + void PutConcurrent(vid_t src, vid_t dst, const EDATA_T& data, + timestamp_t ts) { + CHECK_LT(src, vertex_capacity_); + auto& nbr = nbrs_[src]; + nbr.neighbor = dst; + nbr.data = data; + nbr.timestamp.store(ts, std::memory_order_relaxed); + } + + void RecordFilledEdges(size_t count) { + if (!input_edge_count_precomputed_) { + edge_count_ += static_cast(count); + } + } + + void Finish() { + csr_.edge_num_.store(edge_count_, std::memory_order_relaxed); + csr_.refresh_prefetch_policy(); + } + + private: + void refresh_ptrs() { + vertex_capacity_ = csr_.vertex_capacity(); + nbrs_ = reinterpret_cast( + csr_.nbr_list_ == nullptr ? nullptr : csr_.nbr_list_->GetData()); + } + + csr_t& csr_; + uint64_t edge_count_ = 0; + bool input_edge_count_precomputed_ = false; + size_t vertex_capacity_ = 0; + nbr_t* nbrs_ = nullptr; +}; + } // namespace neug diff --git a/include/neug/storages/graph/edge_table.h b/include/neug/storages/graph/edge_table.h index 56ac97a6a..652d5cfc9 100644 --- a/include/neug/storages/graph/edge_table.h +++ b/include/neug/storages/graph/edge_table.h @@ -41,6 +41,7 @@ class CheckpointManifest; class PropertyGraph; class IDataChunkSupplier; +class IDataChunkSource; class EdgeTable { public: @@ -135,6 +136,22 @@ class EdgeTable { const IndexerType& dst_indexer, std::shared_ptr supplier); + bool CanBatchBuild() const { + if (!meta_ || !meta_->is_bundled() || EdgeNum() != 0) { + return false; + } + // The staged writer handles mutable CSR layouts and the no-adjacency + // layout. Immutable CSR keeps its established incremental path. + return (meta_->oe_strategy == EdgeStrategy::kNone || meta_->oe_mutable) && + (meta_->ie_strategy == EdgeStrategy::kNone || meta_->ie_mutable); + } + + /// Builds a fresh bundled edge table from a repeatable source and swaps it + /// into place only after both passes succeed. + void BatchBuildEdges(const IndexerType& src_indexer, + const IndexerType& dst_indexer, + std::shared_ptr source); + // Add edges in batch to the edge table. void BatchAddEdges(const std::vector& src_lid_list, const std::vector& dst_lid_list, diff --git a/include/neug/storages/graph/graph_interface.h b/include/neug/storages/graph/graph_interface.h index d604239a1..5c83ac9b6 100644 --- a/include/neug/storages/graph/graph_interface.h +++ b/include/neug/storages/graph/graph_interface.h @@ -649,8 +649,19 @@ class StorageAPUpdateInterface : public StorageUpdateInterface { Status BatchAddVertices( label_t v_label_id, std::shared_ptr supplier) override; + /// Narrow AP-only bulk-build entry points. They intentionally do not + /// widen StorageInsertInterface: transactional stores keep their existing + /// COPY semantics and the execution fast path explicitly opts into AP. + bool CanBatchBuildVertices(label_t v_label_id) const; + Status BatchBuildVertices(label_t v_label_id, + std::shared_ptr source); Status BatchAddEdges(label_t src_label, label_t dst_label, label_t edge_label, std::shared_ptr supplier) override; + bool CanBatchBuildEdges(label_t src_label, label_t dst_label, + label_t edge_label) const; + Status BatchBuildEdges(label_t src_label, label_t dst_label, + label_t edge_label, + std::shared_ptr source); Status BatchDeleteVertices(label_t v_label_id, const std::vector& vids) override; Status BatchDeleteEdges( diff --git a/include/neug/storages/graph/property_graph.h b/include/neug/storages/graph/property_graph.h index 4b12e6403..ad8135ba6 100644 --- a/include/neug/storages/graph/property_graph.h +++ b/include/neug/storages/graph/property_graph.h @@ -293,9 +293,19 @@ class PropertyGraph { Status BatchAddVertices(label_t v_label_id, std::shared_ptr supplier); + bool CanBatchBuildVertices(label_t v_label_id) const; + Status BatchBuildVertices(label_t v_label_id, + std::shared_ptr source); + Status BatchAddEdges(label_t src_label, label_t dst_label, label_t edge_label, std::shared_ptr supplier); + bool CanBatchBuildEdges(label_t src_label, label_t dst_label, + label_t edge_label) const; + Status BatchBuildEdges(label_t src_label, label_t dst_label, + label_t edge_label, + std::shared_ptr source); + Status BatchDeleteVertices(label_t v_label_id, const std::vector& vids); diff --git a/include/neug/storages/graph/vertex_table.h b/include/neug/storages/graph/vertex_table.h index 8ca36284a..e5600b9bf 100644 --- a/include/neug/storages/graph/vertex_table.h +++ b/include/neug/storages/graph/vertex_table.h @@ -259,6 +259,13 @@ class VertexTable { void insert_vertices(std::shared_ptr suppliers); + bool CanBatchBuild() const { return Size() == 0; } + + /// Builds an empty table from a repeatable source. CSV sources use a cheap + /// raw row count to reserve before their single typed parse. The existing + /// table is untouched until the staged build succeeds. + void BatchBuildVertices(std::shared_ptr source); + const VertexTimestamp& get_vertex_timestamp() const { return *v_ts_; } const Table& get_table() const { return *table_; } @@ -287,6 +294,33 @@ class VertexTable { return vids; } + void insert_vertices_preallocated( + std::shared_ptr supplier) { + while (auto chunk = supplier->GetNextChunk()) { + auto& columns = chunk->columns; + const auto& property_names = vertex_schema_->property_names; + CHECK_EQ(columns.size(), property_names.size() + 1) + << "Number of columns in the chunk (" << columns.size() + << ") does not match the number of properties (" + << property_names.size() + 1 << ")."; + auto pk_index = std::get<2>(vertex_schema_->primary_keys[0]); + + std::vector> property_columns; + property_columns.reserve(columns.size() - 1); + for (size_t i = 0; i < columns.size(); ++i) { + if (static_cast(i) != pk_index) { + property_columns.push_back(columns[i]); + } + } + + auto vids = insert_primary_keys(columns[pk_index]); + for (size_t i = 0; i < property_columns.size(); ++i) { + set_properties_from_context_column(table_->get_column_by_id(i), + property_columns[i], vids); + } + } + } + std::shared_ptr ckp_; std::unique_ptr indexer_; std::unique_ptr table_; diff --git a/include/neug/storages/loader/chunk_pipeline_utils.h b/include/neug/storages/loader/chunk_pipeline_utils.h new file mode 100644 index 000000000..fc80d129d --- /dev/null +++ b/include/neug/storages/loader/chunk_pipeline_utils.h @@ -0,0 +1,277 @@ +/** Copyright 2020 Alibaba Group Holding Limited. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "neug/storages/loader/loader_utils.h" + +namespace neug { + +struct ChunkPipelineOptions { + int32_t consumer_count = 1; + size_t queue_capacity = 2; +}; + +struct ChunkPipelineAllocation { + bool parallel_enabled = false; + int32_t producer_count = 1; + int32_t consumer_count = 1; + size_t queue_capacity = 2; +}; + +inline int32_t hardware_worker_count() { + auto workers = static_cast(std::thread::hardware_concurrency()); + return workers <= 0 ? 1 : workers; +} + +inline ChunkPipelineAllocation resolve_chunk_pipeline_allocation( + int64_t source_bytes, bool parallel_enabled, bool preserve_order, + int32_t hardware_workers = 0) { + constexpr int64_t kMinParallelBytes = 256LL * 1024 * 1024; + constexpr int64_t kMinPartitionBytes = 64LL * 1024 * 1024; + constexpr size_t kMaxQueuedChunks = 64; + + ChunkPipelineAllocation result; + const auto workers = + hardware_workers > 0 ? hardware_workers : hardware_worker_count(); + if (!parallel_enabled || preserve_order || workers <= 1 || + source_bytes < kMinParallelBytes) { + return result; + } + + // Compute ceil(source_bytes / kMinPartitionBytes) without overflowing when + // source_bytes is close to INT64_MAX. + const auto useful_partitions = std::max( + 1, source_bytes / kMinPartitionBytes + + (source_bytes % kMinPartitionBytes == 0 ? 0 : 1)); + const auto balanced_producers = (workers + 1) / 2; + result.producer_count = static_cast(std::min( + balanced_producers, std::min(useful_partitions, workers - 1))); + result.producer_count = std::max(1, result.producer_count); + result.consumer_count = std::max(1, workers - result.producer_count); + result.parallel_enabled = true; + result.queue_capacity = std::clamp( + static_cast(result.producer_count) * 2, 2, kMaxQueuedChunks); + return result; +} + +namespace chunk_pipeline_detail { + +template +class BoundedQueue { + public: + explicit BoundedQueue(size_t capacity) + : capacity_(std::max(1, capacity)) {} + + bool Push(T value) { + std::unique_lock lock(mutex_); + not_full_.wait(lock, [&] { return closed_ || values_.size() < capacity_; }); + if (closed_) { + return false; + } + values_.push_back(std::move(value)); + not_empty_.notify_one(); + return true; + } + + bool Pop(T& value) { + std::unique_lock lock(mutex_); + not_empty_.wait(lock, [&] { return closed_ || !values_.empty(); }); + if (values_.empty()) { + return false; + } + value = std::move(values_.front()); + values_.pop_front(); + not_full_.notify_one(); + return true; + } + + void Close() { + { + std::lock_guard lock(mutex_); + closed_ = true; + } + not_empty_.notify_all(); + not_full_.notify_all(); + } + + private: + size_t capacity_; + std::mutex mutex_; + std::condition_variable not_empty_; + std::condition_variable not_full_; + std::deque values_; + bool closed_ = false; +}; + +} // namespace chunk_pipeline_detail + +/// Reads chunks from one non-thread-safe supplier and delivers them to a +/// bounded pool of consumers. It preserves the first exception, closes the +/// queue on cancellation, and joins every thread before rethrowing. +template +inline void consume_chunk_pipeline_impl(IDataChunkSupplier& supplier, + ChunkPipelineOptions options, + Consume&& consume) { + const auto consumer_count = std::max(1, options.consumer_count); + if (consumer_count == 1) { + while (auto chunk = supplier.GetNextChunk()) { + consume(0, chunk); + } + return; + } + + chunk_pipeline_detail::BoundedQueue> queue( + options.queue_capacity); + std::atomic cancelled{false}; + std::mutex error_mutex; + std::exception_ptr first_error; + auto capture_error = [&](std::exception_ptr error) { + bool expected = false; + if (cancelled.compare_exchange_strong(expected, true, + std::memory_order_acq_rel)) { + { + std::lock_guard lock(error_mutex); + first_error = std::move(error); + } + queue.Close(); + // A producer may be blocked inside GetNextChunk() when a consumer fails. + // Closing the local queue only wakes this pipeline's threads; cancelling + // the supplier is what gives it a chance to stop background work and + // unblock the producer before join(). + supplier.Cancel(); + } + }; + + std::thread producer([&] { + try { + while (!cancelled.load(std::memory_order_acquire)) { + auto chunk = supplier.GetNextChunk(); + if (!chunk || !queue.Push(std::move(chunk))) { + break; + } + } + } catch (...) { capture_error(std::current_exception()); } + queue.Close(); + }); + + std::vector consumers; + consumers.reserve(static_cast(consumer_count)); + for (int32_t i = 0; i < consumer_count; ++i) { + consumers.emplace_back([&, i] { + try { + std::shared_ptr chunk; + while (!cancelled.load(std::memory_order_acquire) && queue.Pop(chunk)) { + consume(i, chunk); + } + } catch (...) { capture_error(std::current_exception()); } + }); + } + + producer.join(); + for (auto& consumer : consumers) { + consumer.join(); + } + if (first_error) { + std::rethrow_exception(first_error); + } +} + +inline void consume_chunk_pipeline( + IDataChunkSupplier& supplier, ChunkPipelineOptions options, + const std::function&)>& consume) { + consume_chunk_pipeline_impl( + supplier, options, + [&](int32_t /*consumer*/, const std::shared_ptr& chunk) { + consume(chunk); + }); +} + +inline void consume_chunk_pipeline_indexed( + IDataChunkSupplier& supplier, ChunkPipelineOptions options, + const std::function&)>& + consume) { + consume_chunk_pipeline_impl(supplier, options, consume); +} + +/// Concurrently pulls chunks from a supplier that owns its producer queue. +/// This avoids adding another producer thread and bounded queue between a +/// partitioned source and its consumers. +template +inline void consume_concurrent_supplier_indexed(IDataChunkSupplier& supplier, + int32_t consumer_count, + Consume&& consume) { + CHECK(supplier.SupportsConcurrentGetNext()); + consumer_count = std::max(1, consumer_count); + if (consumer_count == 1) { + while (auto chunk = supplier.GetNextChunk()) { + consume(0, chunk); + } + return; + } + + std::atomic cancelled{false}; + std::mutex error_mutex; + std::exception_ptr first_error; + auto capture_error = [&](std::exception_ptr error) { + bool expected = false; + if (cancelled.compare_exchange_strong(expected, true, + std::memory_order_acq_rel)) { + { + std::lock_guard lock(error_mutex); + first_error = std::move(error); + } + supplier.Cancel(); + } + }; + + std::vector consumers; + consumers.reserve(static_cast(consumer_count)); + for (int32_t i = 0; i < consumer_count; ++i) { + consumers.emplace_back([&, i] { + try { + while (!cancelled.load(std::memory_order_acquire)) { + auto chunk = supplier.GetNextChunk(); + if (!chunk) { + break; + } + consume(i, chunk); + } + } catch (...) { capture_error(std::current_exception()); } + }); + } + for (auto& consumer : consumers) { + consumer.join(); + } + if (first_error) { + std::rethrow_exception(first_error); + } +} + +} // namespace neug diff --git a/include/neug/storages/loader/loader_utils.h b/include/neug/storages/loader/loader_utils.h index 32b33aab9..436fef4c1 100644 --- a/include/neug/storages/loader/loader_utils.h +++ b/include/neug/storages/loader/loader_utils.h @@ -73,25 +73,99 @@ class IDataChunkSupplier { virtual ~IDataChunkSupplier() = default; virtual std::shared_ptr GetNextChunk() = 0; virtual int64_t RowNum() const = 0; + + /// Whether GetNextChunk() may be called concurrently by several consumers. + virtual bool SupportsConcurrentGetNext() const { return false; } + + /// Stops any background producers and wakes blocked GetNextChunk() calls. + virtual void Cancel() {} +}; + +struct ChunkSourceOptions { + bool parallel_enabled = false; + int32_t producer_count = 1; + size_t queue_capacity = 2; + bool preserve_order = true; + + /// Zero-based columns in the source's logical output to materialize. An + /// empty list means all columns. Sources must preserve the requested order. + std::vector projected_columns; +}; + +/// A repeatable source of data chunks. +/// +/// Keeping source setup separate from the supplier cursor lets storage consume +/// the input directly without making execution::Context stateful or lazy. +class IDataChunkSource { + public: + virtual ~IDataChunkSource() = default; + + /// Opens a new supplier positioned at the beginning of the source. + virtual std::shared_ptr Open() const = 0; + + /// Opens a supplier with execution-specific concurrency settings. Sources + /// that do not implement parsing-time projection receive a generic + /// projection wrapper around Open(). + virtual std::shared_ptr Open( + const ChunkSourceOptions& options) const; + + /// Whether Open() can be called more than once with identical contents. + virtual bool rewindable() const = 0; + + /// Returns the source size when cheaply known, otherwise -1. + virtual int64_t EstimatedBytes() const { return -1; } + + /// Whether the source options permit parallel parsing and consumption. + virtual bool ParallelEnabled() const { return true; } +}; + +inline constexpr int64_t kUnknownRowNum = -1; + +enum class CsvRowCountMode { + kCountOnOpen, + kUnknown, }; /// csv-parser based supplier. Reads CSV in chunks and yields ValueColumns. class CSVChunkSupplier : public IDataChunkSupplier { public: - CSVChunkSupplier(const std::string& file_path, CsvReadConfig config); + CSVChunkSupplier( + const std::string& file_path, CsvReadConfig config, + CsvRowCountMode row_count_mode = CsvRowCountMode::kCountOnOpen); ~CSVChunkSupplier() override; std::shared_ptr GetNextChunk() override; - int64_t RowNum() const override { return row_num_; } + int64_t RowNum() const override; private: - int64_t row_num_ = 0; std::string file_path_; std::unique_ptr runtime_; }; +struct CsvPartitionPlanCache; + +/// Reopens the public CSV parser for each pass. Parallel suppliers share a +/// cached record-aligned partition plan so repeated passes do not repeat the +/// raw row-count scan. +class CSVChunkSource final : public IDataChunkSource { + public: + CSVChunkSource(std::vector file_paths, CsvReadConfig config); + + std::shared_ptr Open() const override; + std::shared_ptr Open( + const ChunkSourceOptions& options) const override; + bool rewindable() const override { return true; } + int64_t EstimatedBytes() const override; + bool ParallelEnabled() const override { return config_.use_threads; } + + private: + std::vector file_paths_; + CsvReadConfig config_; + std::shared_ptr partition_plan_cache_; +}; + using CSVStreamChunkSupplier = CSVChunkSupplier; using CSVTableChunkSupplier = CSVChunkSupplier; diff --git a/include/neug/utils/io/read/csv/csv_read_config.h b/include/neug/utils/io/read/csv/csv_read_config.h index cf8aba834..93f3448f1 100644 --- a/include/neug/utils/io/read/csv/csv_read_config.h +++ b/include/neug/utils/io/read/csv/csv_read_config.h @@ -31,6 +31,7 @@ struct CsvReadConfig { bool double_quote = true; bool escaping = false; char escape_char = '\\'; + bool use_threads = true; int64_t skip_rows = 0; int64_t chunk_size = 4096; diff --git a/include/neug/utils/io/read/csv/csv_reader.h b/include/neug/utils/io/read/csv/csv_reader.h index bf4f66d74..bbbeec875 100644 --- a/include/neug/utils/io/read/csv/csv_reader.h +++ b/include/neug/utils/io/read/csv/csv_reader.h @@ -27,6 +27,7 @@ namespace neug { class IDataChunkSupplier; +class IDataChunkSource; namespace execution { class Context; @@ -43,6 +44,10 @@ class CsvReader { void read(std::shared_ptr localState, execution::Context& ctx); + /// Creates a repeatable CSV source for direct COPY FROM bulk loading. + /// Returns nullptr when the read needs a row filter and must materialize. + std::shared_ptr createChunkSource(); + result> inferSchema(); private: diff --git a/src/execution/execute/ops/batch/batch_insert_edge.cc b/src/execution/execute/ops/batch/batch_insert_edge.cc index f4fa3c25b..a652f2c0f 100644 --- a/src/execution/execute/ops/batch/batch_insert_edge.cc +++ b/src/execution/execute/ops/batch/batch_insert_edge.cc @@ -14,8 +14,11 @@ */ #include "neug/execution/execute/ops/batch/batch_insert_edge.h" +#include "neug/compiler/function/read_function.h" +#include "neug/compiler/main/metadata_registry.h" #include "neug/execution/common/context.h" #include "neug/execution/execute/ops/batch/batch_update_utils.h" +#include "neug/execution/execute/ops/batch/data_source.h" #include "neug/storages/graph/graph_interface.h" #include "neug/utils/exception/exception.h" #include "neug/utils/result.h" @@ -35,14 +38,14 @@ namespace ops { namespace { -bool resolve_vertex_label_id(const Schema& schema, const common::NameOrId& ni, +bool resolve_vertex_label_id(const Schema& schema, const ::common::NameOrId& ni, label_t& out) { switch (ni.item_case()) { - case common::NameOrId::kId: { + case ::common::NameOrId::kId: { out = ni.id(); return true; } - case common::NameOrId::kName: { + case ::common::NameOrId::kName: { if (!schema.is_vertex_label_valid(ni.name())) { LOG(ERROR) << "Unknown vertex type: " << ni.DebugString(); return false; @@ -62,10 +65,10 @@ bool resolve_edge_triplet(const Schema& schema, label_t& edge_label, label_t& src_type, label_t& dst_type) { switch (edge_type.type_name().item_case()) { - case common::NameOrId::kId: + case ::common::NameOrId::kId: edge_label = edge_type.type_name().id(); break; - case common::NameOrId::kName: { + case ::common::NameOrId::kName: { const auto& name = edge_type.type_name().name(); if (!schema.is_edge_label_valid(name)) { LOG(ERROR) << "Unknown edge type: " @@ -88,6 +91,22 @@ bool resolve_edge_triplet(const Schema& schema, return true; } +std::vector> build_total_edge_mappings( + const std::vector>& source_mappings, + const std::vector>& destination_mappings, + const std::vector>& property_mappings) { + std::vector> mappings; + mappings.reserve(source_mappings.size() + destination_mappings.size() + + property_mappings.size()); + mappings.insert(mappings.end(), source_mappings.begin(), + source_mappings.end()); + mappings.insert(mappings.end(), destination_mappings.begin(), + destination_mappings.end()); + mappings.insert(mappings.end(), property_mappings.begin(), + property_mappings.end()); + return mappings; +} + } // namespace class BatchInsertEdgeOpr : public IOperator { @@ -115,6 +134,37 @@ class BatchInsertEdgeOpr : public IOperator { src_vertex_bindings_, dst_vertex_bindings_; }; +class BatchInsertEdgeFromSourceOpr : public IOperator { + public: + BatchInsertEdgeFromSourceOpr( + std::shared_ptr shared_state, + function::ReadFunction* read_function, physical::EdgeType edge_type, + std::vector> property_mappings, + std::vector> source_mappings, + std::vector> destination_mappings) + : shared_state_(std::move(shared_state)), + read_function_(read_function), + edge_type_(std::move(edge_type)), + property_mappings_(std::move(property_mappings)), + source_mappings_(std::move(source_mappings)), + destination_mappings_(std::move(destination_mappings)) {} + + std::string get_operator_name() const override { + return "BatchInsertEdgeFromSourceOpr"; + } + + neug::result Eval(IStorageInterface& graph, const ParamsMap& params, + Context&& ctx, OprTimer* timer) override; + + private: + std::shared_ptr shared_state_; + function::ReadFunction* read_function_; + physical::EdgeType edge_type_; + std::vector> property_mappings_; + std::vector> source_mappings_; + std::vector> destination_mappings_; +}; + neug::result BatchInsertEdgeOpr::Eval( IStorageInterface& graph_interface, const ParamsMap& params, Context&& ctx, OprTimer* timer) { @@ -131,18 +181,8 @@ neug::result BatchInsertEdgeOpr::Eval( "BatchInsertEdge"); } - std::vector> total_mappings; - total_mappings.reserve(src_vertex_bindings_.size() + - dst_vertex_bindings_.size() + prop_mappings_.size()); - for (const auto& mapping : src_vertex_bindings_) { - total_mappings.emplace_back(mapping); - } - for (const auto& mapping : dst_vertex_bindings_) { - total_mappings.emplace_back(mapping); - } - for (const auto& mapping : prop_mappings_) { - total_mappings.emplace_back(mapping); - } + auto total_mappings = build_total_edge_mappings( + src_vertex_bindings_, dst_vertex_bindings_, prop_mappings_); auto supplier = create_data_chunk_supplier(ctx, total_mappings); RETURN_STATUS_ERROR_IF_NOT_OK( @@ -150,6 +190,48 @@ neug::result BatchInsertEdgeOpr::Eval( return neug::result(std::move(ctx)); } +neug::result BatchInsertEdgeFromSourceOpr::Eval( + IStorageInterface& graph_interface, const ParamsMap& params, Context&& ctx, + OprTimer* timer) { + (void) params; + (void) ctx; + (void) timer; + CHECK(read_function_ != nullptr); + auto& graph = dynamic_cast(graph_interface); + label_t edge_label_id = 0; + label_t src_label_id = 0; + label_t dst_label_id = 0; + if (!resolve_edge_triplet(graph.schema(), edge_type_, edge_label_id, + src_label_id, dst_label_id)) { + RETURN_STATUS_ERROR(StatusCode::ERR_INVALID_ARGUMENT, + "Failed to resolve edge type for " + "BatchInsertEdgeFromSource"); + } + + auto mappings = build_total_edge_mappings( + source_mappings_, destination_mappings_, property_mappings_); + auto* ap_graph = dynamic_cast(&graph_interface); + if (ap_graph && + ap_graph->CanBatchBuildEdges(src_label_id, dst_label_id, edge_label_id) && + read_function_->sourceFunc) { + auto raw_source = read_function_->sourceFunc(shared_state_); + auto source = create_data_chunk_source(std::move(raw_source), mappings); + if (source && should_use_copy_bulk_build(*source)) { + RETURN_STATUS_ERROR_IF_NOT_OK(ap_graph->BatchBuildEdges( + src_label_id, dst_label_id, edge_label_id, std::move(source))); + return Context{}; + } + } + + auto materialized = read_function_->execFunc(shared_state_); + auto supplier = create_data_chunk_supplier(materialized, mappings); + RETURN_STATUS_ERROR_IF_NOT_OK(graph.BatchAddEdges( + src_label_id, dst_label_id, edge_label_id, std::move(supplier))); + // Match the empty terminal sink consumed by the fused plan. + materialized.tag_ids.clear(); + return neug::result(std::move(materialized)); +} + neug::result BatchInsertEdgeOprBuilder::Build( const Schema& schema, const ContextMeta& ctx_meta, const physical::PhysicalPlan& plan, int op_idx) { @@ -177,6 +259,45 @@ neug::result BatchInsertEdgeOprBuilder::Build( ret_meta); } +neug::result BatchInsertEdgeFromSourceOprBuilder::Build( + const Schema& schema, const ContextMeta& ctx_meta, + const physical::PhysicalPlan& plan, int op_idx) { + (void) schema; + ContextMeta result_meta = ctx_meta; + if (op_idx + 3 != plan.plan_size() || + plan.plan(op_idx + 2).opr().sink().tags_size() != 0) { + return std::make_pair(nullptr, result_meta); + } + const auto& source_pb = plan.plan(op_idx).opr().source(); + const auto& edge_pb = plan.plan(op_idx + 1).opr().load_edge(); + if (!edge_pb.has_edge_type()) { + THROW_INTERNAL_EXCEPTION( + "BatchInsertEdgeFromSourceOprBuilder: edge type is not set"); + } + + ReadStateBuilder state_builder; + auto state = state_builder.build(source_pb); + auto catalog = neug::main::MetadataRegistry::getCatalog(); + auto function = catalog->getFunctionWithSignature(source_pb.extension_name()); + auto read_function = function->ptrCast(); + + std::vector> property_mappings; + std::vector> source_mappings; + std::vector> destination_mappings; + parse_property_mappings(edge_pb.property_mappings(), property_mappings); + parse_property_mappings(edge_pb.source_vertex_binding(), source_mappings); + parse_property_mappings(edge_pb.destination_vertex_binding(), + destination_mappings); + physical::EdgeType edge_type; + edge_type.CopyFrom(edge_pb.edge_type()); + return std::make_pair( + std::make_unique( + std::move(state), read_function, std::move(edge_type), + std::move(property_mappings), std::move(source_mappings), + std::move(destination_mappings)), + result_meta); +} + } // namespace ops } // namespace execution } // namespace neug diff --git a/src/execution/execute/ops/batch/batch_insert_vertex.cc b/src/execution/execute/ops/batch/batch_insert_vertex.cc index b4454f3a8..bc2284714 100644 --- a/src/execution/execute/ops/batch/batch_insert_vertex.cc +++ b/src/execution/execute/ops/batch/batch_insert_vertex.cc @@ -14,8 +14,11 @@ */ #include "neug/execution/execute/ops/batch/batch_insert_vertex.h" +#include "neug/compiler/function/read_function.h" +#include "neug/compiler/main/metadata_registry.h" #include "neug/execution/common/context.h" #include "neug/execution/execute/ops/batch/batch_update_utils.h" +#include "neug/execution/execute/ops/batch/data_source.h" #include "neug/storages/graph/graph_interface.h" #include "neug/utils/exception/exception.h" #include "neug/utils/result.h" @@ -31,10 +34,34 @@ class OprTimer; namespace ops { +namespace { + +bool resolve_vertex_label_id(const Schema& schema, + const ::common::NameOrId& type, + label_t& label_id) { + switch (type.item_case()) { + case ::common::NameOrId::kId: + label_id = type.id(); + return true; + case ::common::NameOrId::kName: + if (!schema.is_vertex_label_valid(type.name())) { + LOG(ERROR) << "Unknown vertex type: " << type.DebugString(); + return false; + } + label_id = schema.get_vertex_label_id(type.name()); + return true; + default: + LOG(ERROR) << "Invalid vertex type: " << type.DebugString(); + return false; + } +} + +} // namespace + class BatchInsertVertexOpr : public IOperator { public: BatchInsertVertexOpr( - common::NameOrId vertex_type, + ::common::NameOrId vertex_type, std::vector> prop_mappings) : vertex_type_(std::move(vertex_type)), prop_mappings_(std::move(prop_mappings)) {} @@ -47,10 +74,35 @@ class BatchInsertVertexOpr : public IOperator { Context&& ctx, OprTimer* timer) override; private: - common::NameOrId vertex_type_; + ::common::NameOrId vertex_type_; std::vector> prop_mappings_; }; +class BatchInsertVertexFromSourceOpr : public IOperator { + public: + BatchInsertVertexFromSourceOpr( + std::shared_ptr shared_state, + function::ReadFunction* read_function, ::common::NameOrId vertex_type, + std::vector> property_mappings) + : shared_state_(std::move(shared_state)), + read_function_(read_function), + vertex_type_(std::move(vertex_type)), + property_mappings_(std::move(property_mappings)) {} + + std::string get_operator_name() const override { + return "BatchInsertVertexFromSourceOpr"; + } + + neug::result Eval(IStorageInterface& graph, const ParamsMap& params, + Context&& ctx, OprTimer* timer) override; + + private: + std::shared_ptr shared_state_; + function::ReadFunction* read_function_; + ::common::NameOrId vertex_type_; + std::vector> property_mappings_; +}; + neug::result BatchInsertVertexOpr::Eval( IStorageInterface& graph_interface, const ParamsMap& params, Context&& ctx, OprTimer* timer) { @@ -58,24 +110,9 @@ neug::result BatchInsertVertexOpr::Eval( (void) timer; auto& graph = dynamic_cast(graph_interface); label_t vertex_label_id = 0; - switch (vertex_type_.item_case()) { - case common::NameOrId::kId: - vertex_label_id = vertex_type_.id(); - break; - case common::NameOrId::kName: { - const auto& name = vertex_type_.name(); - if (!graph.schema().is_vertex_label_valid(name)) { - LOG(ERROR) << "Unknown vertex type: " << vertex_type_.DebugString(); - RETURN_STATUS_ERROR(StatusCode::ERR_INVALID_ARGUMENT, - "Unknown vertex type: " + name); - } - vertex_label_id = graph.schema().get_vertex_label_id(name); - break; - } - default: - THROW_INVALID_ARGUMENT_EXCEPTION( - "BatchInsertVertexOpr: invalid vertex_type: " + - vertex_type_.DebugString()); + if (!resolve_vertex_label_id(graph.schema(), vertex_type_, vertex_label_id)) { + RETURN_STATUS_ERROR(StatusCode::ERR_INVALID_ARGUMENT, + "Failed to resolve vertex type for BatchInsertVertex"); } auto supplier = create_data_chunk_supplier(ctx, prop_mappings_); RETURN_STATUS_ERROR_IF_NOT_OK( @@ -83,6 +120,44 @@ neug::result BatchInsertVertexOpr::Eval( return neug::result(std::move(ctx)); } +neug::result BatchInsertVertexFromSourceOpr::Eval( + IStorageInterface& graph_interface, const ParamsMap& params, Context&& ctx, + OprTimer* timer) { + (void) params; + (void) ctx; + (void) timer; + CHECK(read_function_ != nullptr); + auto& graph = dynamic_cast(graph_interface); + label_t vertex_label_id = 0; + if (!resolve_vertex_label_id(graph.schema(), vertex_type_, vertex_label_id)) { + RETURN_STATUS_ERROR( + StatusCode::ERR_INVALID_ARGUMENT, + "Failed to resolve vertex type for BatchInsertVertexFromSource"); + } + + auto* ap_graph = dynamic_cast(&graph_interface); + if (ap_graph && ap_graph->CanBatchBuildVertices(vertex_label_id) && + read_function_->sourceFunc) { + auto raw_source = read_function_->sourceFunc(shared_state_); + auto source = + create_data_chunk_source(std::move(raw_source), property_mappings_); + if (source && should_use_copy_bulk_build(*source)) { + RETURN_STATUS_ERROR_IF_NOT_OK( + ap_graph->BatchBuildVertices(vertex_label_id, std::move(source))); + // The fused builder only accepts a terminal COPY with an empty sink. + return Context{}; + } + } + + auto materialized = read_function_->execFunc(shared_state_); + auto supplier = create_data_chunk_supplier(materialized, property_mappings_); + RETURN_STATUS_ERROR_IF_NOT_OK( + graph.BatchAddVertices(vertex_label_id, std::move(supplier))); + // Match the empty terminal sink consumed by the fused plan. + materialized.tag_ids.clear(); + return neug::result(std::move(materialized)); +} + neug::result BatchInsertVertexOprBuilder::Build( const Schema& schema, const ContextMeta& ctx_meta, const physical::PhysicalPlan& plan, int op_idx) { @@ -96,13 +171,46 @@ neug::result BatchInsertVertexOprBuilder::Build( std::vector> prop_mappings; parse_property_mappings(opr.property_mappings(), prop_mappings); - common::NameOrId vertex_type; + ::common::NameOrId vertex_type; vertex_type.CopyFrom(opr.vertex_type()); return std::make_pair(std::make_unique( std::move(vertex_type), std::move(prop_mappings)), ret_meta); } +neug::result BatchInsertVertexFromSourceOprBuilder::Build( + const Schema& schema, const ContextMeta& ctx_meta, + const physical::PhysicalPlan& plan, int op_idx) { + (void) schema; + ContextMeta result_meta = ctx_meta; + if (op_idx + 3 != plan.plan_size() || + plan.plan(op_idx + 2).opr().sink().tags_size() != 0) { + return std::make_pair(nullptr, result_meta); + } + const auto& source_pb = plan.plan(op_idx).opr().source(); + const auto& vertex_pb = plan.plan(op_idx + 1).opr().load_vertex(); + if (!vertex_pb.has_vertex_type()) { + THROW_INTERNAL_EXCEPTION( + "BatchInsertVertexFromSourceOprBuilder: vertex type is not set"); + } + + ReadStateBuilder state_builder; + auto state = state_builder.build(source_pb); + auto catalog = neug::main::MetadataRegistry::getCatalog(); + auto function = catalog->getFunctionWithSignature(source_pb.extension_name()); + auto read_function = function->ptrCast(); + + std::vector> property_mappings; + parse_property_mappings(vertex_pb.property_mappings(), property_mappings); + ::common::NameOrId vertex_type; + vertex_type.CopyFrom(vertex_pb.vertex_type()); + return std::make_pair( + std::make_unique( + std::move(state), read_function, std::move(vertex_type), + std::move(property_mappings)), + result_meta); +} + } // namespace ops } // namespace execution } // namespace neug diff --git a/src/execution/execute/ops/batch/batch_update_utils.cc b/src/execution/execute/ops/batch/batch_update_utils.cc index 4339f2947..fc765f408 100644 --- a/src/execution/execute/ops/batch/batch_update_utils.cc +++ b/src/execution/execute/ops/batch/batch_update_utils.cc @@ -22,7 +22,10 @@ #include #include +#include +#include #include +#include #include #include #include @@ -316,6 +319,108 @@ class MultiChunkSupplier : public IDataChunkSupplier { size_t index_; }; +class ProjectingChunkSupplier final : public IDataChunkSupplier { + public: + ProjectingChunkSupplier(std::shared_ptr input, + std::vector aliases) + : input_(std::move(input)), aliases_(std::move(aliases)) {} + + std::shared_ptr GetNextChunk() override { + auto input = input_->GetNextChunk(); + if (!input) { + return nullptr; + } + auto output = std::make_shared(); + for (size_t index = 0; index < aliases_.size(); ++index) { + auto column = input->get(aliases_[index]); + if (!column) { + THROW_INTERNAL_EXCEPTION("Column not found for tag id: " + + std::to_string(aliases_[index])); + } + output->set(static_cast(index), std::move(column)); + } + return output; + } + + int64_t RowNum() const override { return input_->RowNum(); } + + bool SupportsConcurrentGetNext() const override { + return input_->SupportsConcurrentGetNext(); + } + + void Cancel() override { input_->Cancel(); } + + private: + std::shared_ptr input_; + std::vector aliases_; +}; + +class ProjectingChunkSource final : public IDataChunkSource { + public: + ProjectingChunkSource(std::shared_ptr input, + std::vector aliases) + : input_(std::move(input)), aliases_(std::move(aliases)) {} + + std::shared_ptr Open() const override { + auto supplier = input_->Open(); + if (!supplier) { + return nullptr; + } + return std::make_shared(std::move(supplier), + aliases_); + } + + std::shared_ptr Open( + const ChunkSourceOptions& options) const override { + if (!options.projected_columns.empty()) { + ChunkSourceOptions input_options = options; + input_options.projected_columns.clear(); + input_options.projected_columns.reserve(options.projected_columns.size()); + for (const auto output_column : options.projected_columns) { + if (output_column < 0 || + static_cast(output_column) >= aliases_.size()) { + THROW_INVALID_ARGUMENT_EXCEPTION( + "Projected source column is out of range: " + + std::to_string(output_column)); + } + input_options.projected_columns.push_back( + aliases_[static_cast(output_column)]); + } + // IDataChunkSource::Open(options) guarantees that projected columns are + // returned in the requested order. CSV sources push this into typed + // parsing; other repeatable sources receive the generic wrapper. + return input_->Open(input_options); + } + + auto supplier = input_->Open(options); + if (!supplier) { + return nullptr; + } + return std::make_shared(std::move(supplier), + aliases_); + } + + bool rewindable() const override { return input_->rewindable(); } + + int64_t EstimatedBytes() const override { return input_->EstimatedBytes(); } + + bool ParallelEnabled() const override { return input_->ParallelEnabled(); } + + private: + std::shared_ptr input_; + std::vector aliases_; +}; + +std::vector property_mapping_aliases( + const std::vector>& prop_mappings) { + std::vector aliases; + aliases.reserve(prop_mappings.size()); + for (const auto& mapping : prop_mappings) { + aliases.push_back(mapping.first); + } + return aliases; +} + std::shared_ptr create_data_chunk_supplier( const Context& ctx, const std::vector>& prop_mappings) { @@ -338,6 +443,35 @@ std::shared_ptr create_data_chunk_supplier( return std::make_shared(std::move(projected_chunks)); } +std::shared_ptr create_data_chunk_source( + std::shared_ptr source, + const std::vector>& prop_mappings) { + if (!source) { + return nullptr; + } + return std::make_shared( + std::move(source), property_mapping_aliases(prop_mappings)); +} + +bool should_use_copy_bulk_build(const IDataChunkSource& source) { + const char* configured = std::getenv("NEUG_COPY_BULK_BUILD"); + if (configured != nullptr) { + std::string value(configured); + std::transform(value.begin(), value.end(), value.begin(), + [](unsigned char ch) { return std::tolower(ch); }); + if (value == "0" || value == "false" || value == "off" || value == "no") { + return false; + } + if (value == "1" || value == "true" || value == "on" || value == "yes") { + return source.rewindable(); + } + LOG(WARNING) << "Ignore invalid NEUG_COPY_BULK_BUILD=" << configured; + } + + constexpr int64_t kMinBulkBuildBytes = 256LL * 1024 * 1024; + return source.rewindable() && source.EstimatedBytes() >= kMinBulkBuildBytes; +} + std::vector match_files_with_pattern( const std::string& file_path) { std::vector result; diff --git a/src/execution/execute/plan_parser.cc b/src/execution/execute/plan_parser.cc index 5a2afb051..91a143e54 100644 --- a/src/execution/execute/plan_parser.cc +++ b/src/execution/execute/plan_parser.cc @@ -126,6 +126,10 @@ void PlanParser::init() { register_operator_builder(std::make_unique()); + register_operator_builder( + std::make_unique()); + register_operator_builder( + std::make_unique()); register_operator_builder(std::make_unique()); register_operator_builder( std::make_unique()); diff --git a/src/storages/graph/edge_table.cc b/src/storages/graph/edge_table.cc index 98664b616..4045aff43 100644 --- a/src/storages/graph/edge_table.cc +++ b/src/storages/graph/edge_table.cc @@ -20,18 +20,29 @@ #include "neug/storages/module/module_factory.h" #include +#include #include +#include +#include #include +#include #include +#include +#include +#include #include #include +#include #include +#include #include "neug/common/columns/value_columns.h" +#include "neug/common/types/container_types.h" #include "neug/storages/checkpoint_manager.h" #include "neug/storages/csr/csr_view_utils.h" #include "neug/storages/csr/immutable_csr.h" #include "neug/storages/csr/mutable_csr.h" +#include "neug/storages/loader/chunk_pipeline_utils.h" #include "neug/storages/loader/loader_utils.h" #include "neug/storages/module/type_name.h" #include "neug/storages/module_descriptor.h" @@ -40,6 +51,39 @@ namespace neug { +namespace { + +size_t vector_bool_storage_bytes(size_t bit_capacity) { + return (bit_capacity + 7) / 8; +} + +uint64_t process_peak_rss_bytes() { + struct rusage usage {}; + if (getrusage(RUSAGE_SELF, &usage) != 0) { + return 0; + } +#if defined(__APPLE__) + return static_cast(usage.ru_maxrss); +#else + return static_cast(usage.ru_maxrss) * 1024; +#endif +} + +size_t estimate_edge_fallback_buffer_bytes( + const std::vector& src_lid, const std::vector& dst_lid, + const std::vector& valid_flags, + const std::vector>& bundled_data_cols, + const std::vector>& unbundled_data_chunks) { + return src_lid.capacity() * sizeof(vid_t) + + dst_lid.capacity() * sizeof(vid_t) + + vector_bool_storage_bytes(valid_flags.capacity()) + + bundled_data_cols.capacity() * + sizeof(std::shared_ptr) + + unbundled_data_chunks.capacity() * sizeof(std::shared_ptr); +} + +} // namespace + void filterInvalidEdges(std::vector& src_lid, std::vector& dst_lid, std::vector& valid_flags) { @@ -384,6 +428,1041 @@ void batch_add_bundled_edges_impl( } } +template +class EmptyCsrBulkWriter { + public: + static constexpr bool kSupportsDisjointConcurrentFill = true; + static constexpr bool kStoresEdges = false; + static constexpr bool kNeedsDegreeCount = false; + static constexpr bool kChecksSingleUniqueness = false; + static constexpr bool kNeedsConcurrentGrouping = false; + static constexpr bool kTracksInputEdgeCount = false; + + void PrepareBuild(vid_t /*vertex_count*/) {} + void AllocateFromCounts() {} + int ReserveConcurrent(vid_t /*src*/, int /*count*/) { return 0; } + void PutAt(vid_t /*src*/, int /*slot*/, vid_t /*dst*/, + const EDATA_T& /*data*/, timestamp_t /*ts*/) {} + void PutSerial(vid_t /*src*/, vid_t /*dst*/, const EDATA_T& /*data*/, + timestamp_t /*ts*/) {} + void Finish() {} +}; + +template +bool with_csr_bulk_writer(CsrBase* csr, F&& callback) { + if (auto* typed = dynamic_cast*>(csr)) { + MutableCsrBulkBuildAccess writer(*typed); + std::forward(callback)(writer); + return true; + } + if (auto* typed = dynamic_cast*>(csr)) { + SingleMutableCsrBulkBuildAccess writer(*typed); + std::forward(callback)(writer); + return true; + } + if (dynamic_cast*>(csr) != nullptr) { + EmptyCsrBulkWriter writer; + std::forward(callback)(writer); + return true; + } + return false; +} + +template +class BulkEdgeDataReader { + public: + explicit BulkEdgeDataReader(const std::shared_ptr& column) + : column_(column.get()) { + CHECK(column_ != nullptr); + auto values = std::dynamic_pointer_cast>(column); + if (values) { + values_ = &values->data(); + } + } + + EDATA_T Get(size_t row) const { + if (values_ != nullptr) { + return (*values_)[row]; + } + return column_->get_elem(row).template GetValue(); + } + + private: + const IContextColumn* column_; + const vector_t* values_ = nullptr; +}; + +template <> +class BulkEdgeDataReader { + public: + explicit BulkEdgeDataReader( + const std::shared_ptr& /*column*/) {} + + EmptyType Get(size_t /*row*/) const { return EmptyType(); } +}; + +struct BulkEdgeEndpointScratch { + std::vector src_lids; + std::vector dst_lids; +}; + +void index_bulk_edge_endpoints(const std::shared_ptr& chunk, + const IndexerType& src_indexer, + const IndexerType& dst_indexer, + BulkEdgeEndpointScratch& scratch) { + CHECK(chunk != nullptr); + CHECK_GE(chunk->col_num(), 2); + auto src_column = chunk->get(0); + auto dst_column = chunk->get(1); + CHECK(src_column != nullptr); + CHECK(dst_column != nullptr); + CHECK_EQ(src_column->size(), dst_column->size()); + + scratch.src_lids.clear(); + scratch.dst_lids.clear(); + src_indexer.get_index(*src_column, scratch.src_lids); + dst_indexer.get_index(*dst_column, scratch.dst_lids); + CHECK_EQ(scratch.src_lids.size(), scratch.dst_lids.size()); +} + +inline bool is_valid_bulk_edge(vid_t src, vid_t dst) { + return src != std::numeric_limits::max() && + dst != std::numeric_limits::max(); +} + +size_t count_valid_bulk_edges(const BulkEdgeEndpointScratch& endpoints) { + size_t valid_edges = 0; + for (size_t row = 0; row < endpoints.src_lids.size(); ++row) { + valid_edges += + is_valid_bulk_edge(endpoints.src_lids[row], endpoints.dst_lids[row]); + } + return valid_edges; +} + +struct BulkEdgeCountScratch { + BulkEdgeEndpointScratch endpoints; + flat_hash_map out_counts; + flat_hash_map in_counts; + bool out_single_duplicate = false; + bool in_single_duplicate = false; + uint64_t valid_edges = 0; +}; + +struct BulkEdgeCountChunkResult { + size_t valid_edges = 0; + size_t out_single_checks = 0; + size_t in_single_checks = 0; +}; + +struct BulkEdgeCountProfile { + size_t out_updates = 0; + size_t in_updates = 0; +}; + +struct BulkEdgeCountSummary { + uint64_t valid_edges = 0; + bool out_single_duplicate = false; + bool in_single_duplicate = false; +}; + +struct BulkEdgeCountConfig { + bool concurrent = false; + bool check_out_single = false; + bool check_in_single = false; +}; + +constexpr size_t kBulkEdgeInitialGroupReserve = 4096; + +size_t bulk_edge_group_reserve(size_t rows) { + // Eagerly reserving one hash slot per edge wastes substantial memory for a + // supernode. The map can grow past this cap for high-cardinality chunks, and + // the worker-local scratch reuses that larger backing store on later chunks. + return std::min(rows, kBulkEdgeInitialGroupReserve); +} + +template +BulkEdgeCountChunkResult count_bulk_edge_chunk( + const BulkEdgeEndpointScratch& endpoints, OutWriter& out, InWriter& in, + BulkEdgeCountScratch& scratch, const BulkEdgeCountConfig& config) { + CHECK_LE(endpoints.src_lids.size(), + static_cast(std::numeric_limits::max())); + BulkEdgeCountChunkResult result; + + if (config.concurrent) { + if constexpr (OutWriter::kNeedsDegreeCount) { + scratch.out_counts.clear(); + scratch.out_counts.reserve( + bulk_edge_group_reserve(endpoints.src_lids.size())); + } + if constexpr (InWriter::kNeedsDegreeCount) { + scratch.in_counts.clear(); + scratch.in_counts.reserve( + bulk_edge_group_reserve(endpoints.dst_lids.size())); + } + } + for (size_t row = 0; row < endpoints.src_lids.size(); ++row) { + const auto src = endpoints.src_lids[row]; + const auto dst = endpoints.dst_lids[row]; + if (!is_valid_bulk_edge(src, dst)) { + continue; + } + ++result.valid_edges; + if (config.concurrent) { + if constexpr (OutWriter::kNeedsDegreeCount) { + ++scratch.out_counts[src]; + } + if constexpr (InWriter::kNeedsDegreeCount) { + ++scratch.in_counts[dst]; + } + if constexpr (OutWriter::kChecksSingleUniqueness) { + if (config.check_out_single && !scratch.out_single_duplicate) { + ++result.out_single_checks; + scratch.out_single_duplicate = out.CheckUniqueConcurrent(src); + } + } + if constexpr (InWriter::kChecksSingleUniqueness) { + if (config.check_in_single && !scratch.in_single_duplicate) { + ++result.in_single_checks; + scratch.in_single_duplicate = in.CheckUniqueConcurrent(dst); + } + } + } else { + if constexpr (OutWriter::kNeedsDegreeCount) { + out.CountSerial(src); + } + if constexpr (InWriter::kNeedsDegreeCount) { + in.CountSerial(dst); + } + if constexpr (OutWriter::kChecksSingleUniqueness) { + if (config.check_out_single && !scratch.out_single_duplicate) { + ++result.out_single_checks; + scratch.out_single_duplicate = out.CheckUniqueSerial(src); + } + } + if constexpr (InWriter::kChecksSingleUniqueness) { + if (config.check_in_single && !scratch.in_single_duplicate) { + ++result.in_single_checks; + scratch.in_single_duplicate = in.CheckUniqueSerial(dst); + } + } + } + } + if (config.concurrent) { + if constexpr (OutWriter::kNeedsDegreeCount) { + for (const auto& [src, count] : scratch.out_counts) { + CHECK_LE(count, static_cast(std::numeric_limits::max())); + out.CountConcurrent(src, static_cast(count)); + } + } + if constexpr (InWriter::kNeedsDegreeCount) { + for (const auto& [dst, count] : scratch.in_counts) { + CHECK_LE(count, static_cast(std::numeric_limits::max())); + in.CountConcurrent(dst, static_cast(count)); + } + } + } + return result; +} + +BulkEdgeCountProfile profile_bulk_edge_count( + const BulkEdgeCountScratch& scratch, + const BulkEdgeCountChunkResult& chunk_result, bool concurrent_count, + bool counts_out_degree, bool counts_in_degree) { + BulkEdgeCountProfile result; + if (counts_out_degree) { + result.out_updates = + concurrent_count ? scratch.out_counts.size() : chunk_result.valid_edges; + } + if (counts_in_degree) { + result.in_updates = + concurrent_count ? scratch.in_counts.size() : chunk_result.valid_edges; + } + return result; +} + +ChunkSourceOptions make_bulk_edge_source_options( + const ChunkPipelineAllocation& allocation, bool preserve_order) { + ChunkSourceOptions source_options; + source_options.parallel_enabled = allocation.parallel_enabled; + source_options.producer_count = allocation.producer_count; + source_options.queue_capacity = allocation.queue_capacity; + source_options.preserve_order = preserve_order; + return source_options; +} + +struct BulkEdgeCountWorkerStats { + size_t chunks = 0; + size_t edges = 0; + size_t out_updates = 0; + size_t in_updates = 0; + size_t out_single_checks = 0; + size_t in_single_checks = 0; + int64_t work_time_ns = 0; +}; + +using BulkEdgeCountChunk = std::function; + +BulkEdgeCountSummary count_bulk_edges( + const std::shared_ptr& source, + const IndexerType& src_indexer, const IndexerType& dst_indexer, + const ChunkPipelineAllocation& allocation, bool counts_out_degree, + bool counts_in_degree, bool check_out_single, bool check_in_single, + const BulkEdgeCountChunk& count_chunk) { + // Degree accumulation and single-slot uniqueness checks are commutative, so + // this pass never needs input order even when fill may later fall back to the + // ordered last-write-wins path. + auto source_options = make_bulk_edge_source_options(allocation, false); + // The degree pass only needs endpoint OIDs. Push this projection through + // ProjectingChunkSource into CSVChunkSource so edge properties are not typed, + // allocated, and discarded during the first parse. + source_options.projected_columns = {0, 1}; + auto supplier = source->Open(source_options); + CHECK(supplier != nullptr); + + const bool profile_stages = VLOG_IS_ON(1); + std::chrono::steady_clock::time_point row_count_start; + std::chrono::steady_clock::time_point row_count_end; + int64_t row_num = 0; + if (profile_stages) { + row_count_start = std::chrono::steady_clock::now(); + row_num = supplier->RowNum(); + row_count_end = std::chrono::steady_clock::now(); + } + const auto worker_count = allocation.consumer_count; + std::vector scratches( + static_cast(worker_count)); + std::vector worker_stats; + if (profile_stages) { + worker_stats.resize(static_cast(worker_count)); + } + const BulkEdgeCountConfig count_config{ + .concurrent = worker_count > 1, + .check_out_single = check_out_single, + .check_in_single = check_in_single, + }; + auto count = [&](int32_t worker, const std::shared_ptr& chunk) { + CHECK_GE(worker, 0); + CHECK_LT(worker, worker_count); + const auto start = profile_stages ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + auto& scratch = scratches[static_cast(worker)]; + index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, + scratch.endpoints); + const auto chunk_result = count_chunk(scratch, count_config); + scratch.valid_edges += static_cast(chunk_result.valid_edges); + if (profile_stages) { + const auto work_end = std::chrono::steady_clock::now(); + const auto profile = profile_bulk_edge_count( + scratch, chunk_result, count_config.concurrent, counts_out_degree, + counts_in_degree); + auto& stats = worker_stats[static_cast(worker)]; + ++stats.chunks; + stats.edges += chunk_result.valid_edges; + stats.out_updates += profile.out_updates; + stats.in_updates += profile.in_updates; + stats.out_single_checks += chunk_result.out_single_checks; + stats.in_single_checks += chunk_result.in_single_checks; + stats.work_time_ns += + std::chrono::duration_cast(work_end - start) + .count(); + } + }; + VLOG(1) << "Bulk edge count pass: bytes=" << source->EstimatedBytes() + << ", producers=" << allocation.producer_count + << ", consumers=" << allocation.consumer_count + << ", queue_capacity=" << allocation.queue_capacity + << ", check_out_single=" << check_out_single + << ", check_in_single=" << check_in_single + << ", direct_supplier_queue=" + << supplier->SupportsConcurrentGetNext(); + const auto consume_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + if (supplier->SupportsConcurrentGetNext()) { + consume_concurrent_supplier_indexed(*supplier, worker_count, count); + } else if (worker_count == 1) { + while (auto chunk = supplier->GetNextChunk()) { + count(0, chunk); + } + } else { + ChunkPipelineOptions options; + options.consumer_count = worker_count; + options.queue_capacity = allocation.queue_capacity; + consume_chunk_pipeline_indexed(*supplier, options, count); + } + const auto consume_end = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + BulkEdgeCountSummary summary; + for (size_t worker = 0; worker < scratches.size(); ++worker) { + summary.valid_edges += scratches[worker].valid_edges; + summary.out_single_duplicate = + summary.out_single_duplicate || scratches[worker].out_single_duplicate; + summary.in_single_duplicate = + summary.in_single_duplicate || scratches[worker].in_single_duplicate; + } + if (!profile_stages) { + return summary; + } + size_t out_updates = 0; + size_t in_updates = 0; + size_t out_single_checks = 0; + size_t in_single_checks = 0; + int64_t vid_degree_time_ns = 0; + for (const auto& stats : worker_stats) { + out_updates += stats.out_updates; + in_updates += stats.in_updates; + out_single_checks += stats.out_single_checks; + in_single_checks += stats.in_single_checks; + vid_degree_time_ns += stats.work_time_ns; + } + for (size_t worker = 0; worker < worker_stats.size(); ++worker) { + const auto& stats = worker_stats[worker]; + VLOG(2) << "Bulk edge count worker: worker=" << worker + << ", chunks=" << stats.chunks << ", edges=" << stats.edges + << ", out_atomic_updates=" << stats.out_updates + << ", in_atomic_updates=" << stats.in_updates + << ", out_single_checks=" << stats.out_single_checks + << ", in_single_checks=" << stats.in_single_checks + << ", worker_ms=" << stats.work_time_ns / 1000000; + } + VLOG(1) << "Bulk edge count stages: rows=" << row_num << ", row_count_ms=" + << std::chrono::duration_cast( + row_count_end - row_count_start) + .count() + << ", pipeline_wall_ms=" + << std::chrono::duration_cast( + consume_end - consume_start) + .count() + << ", valid_edges=" << summary.valid_edges + << ", out_atomic_updates=" << out_updates + << ", in_atomic_updates=" << in_updates + << ", out_single_checks=" << out_single_checks + << ", in_single_checks=" << in_single_checks + << ", out_single_duplicate=" << summary.out_single_duplicate + << ", in_single_duplicate=" << summary.in_single_duplicate + << ", vid_degree_worker_ms=" << vid_degree_time_ns / 1000000; + return summary; +} + +constexpr uint32_t kInvalidBulkEdgeIndex = std::numeric_limits::max(); + +struct BulkEdgeGroup { + uint32_t count = 0; + uint32_t head = kInvalidBulkEdgeIndex; + uint32_t tail = kInvalidBulkEdgeIndex; +}; + +struct BulkEdgeFillScratch { + BulkEdgeEndpointScratch endpoints; + flat_hash_map out_groups; + flat_hash_map in_groups; + std::vector out_next; + std::vector in_next; +}; + +struct BulkEdgeFillResult { + size_t valid_edges = 0; + size_t out_reservations = 0; + size_t in_reservations = 0; +}; + +void append_bulk_edge_group(flat_hash_map& groups, + std::vector& next, vid_t vertex, + uint32_t edge_index) { + auto [it, inserted] = groups.try_emplace(vertex); + auto& group = it->second; + if (inserted) { + group.head = edge_index; + } else { + CHECK_NE(group.tail, kInvalidBulkEdgeIndex); + next[group.tail] = edge_index; + } + group.tail = edge_index; + CHECK_LT(group.count, std::numeric_limits::max()); + ++group.count; +} + +template +void group_bulk_edge_chunk(const BulkEdgeEndpointScratch& endpoints, + BulkEdgeFillScratch& scratch) { + CHECK_LE(endpoints.src_lids.size(), + static_cast(std::numeric_limits::max())); + if constexpr (OutWriter::kNeedsConcurrentGrouping) { + scratch.out_groups.clear(); + scratch.out_groups.reserve( + bulk_edge_group_reserve(endpoints.src_lids.size())); + scratch.out_next.assign(endpoints.src_lids.size(), kInvalidBulkEdgeIndex); + } + if constexpr (InWriter::kNeedsConcurrentGrouping) { + scratch.in_groups.clear(); + scratch.in_groups.reserve( + bulk_edge_group_reserve(endpoints.dst_lids.size())); + scratch.in_next.assign(endpoints.dst_lids.size(), kInvalidBulkEdgeIndex); + } + for (size_t row = 0; row < endpoints.src_lids.size(); ++row) { + const auto src = endpoints.src_lids[row]; + const auto dst = endpoints.dst_lids[row]; + if (!is_valid_bulk_edge(src, dst)) { + continue; + } + const auto edge_index = static_cast(row); + if constexpr (OutWriter::kNeedsConcurrentGrouping) { + append_bulk_edge_group(scratch.out_groups, scratch.out_next, src, + edge_index); + } + if constexpr (InWriter::kNeedsConcurrentGrouping) { + append_bulk_edge_group(scratch.in_groups, scratch.in_next, dst, + edge_index); + } + } +} + +template +void fill_bulk_edge_chunk_serial(const BulkEdgeEndpointScratch& endpoints, + const BulkEdgeDataReader& data_reader, + OutWriter& out, InWriter& in) { + size_t valid_edges = 0; + for (size_t row = 0; row < endpoints.src_lids.size(); ++row) { + const auto src = endpoints.src_lids[row]; + const auto dst = endpoints.dst_lids[row]; + if (!is_valid_bulk_edge(src, dst)) { + continue; + } + ++valid_edges; + const auto data = data_reader.Get(row); + if constexpr (OutWriter::kStoresEdges) { + out.PutSerial(src, dst, data, 0); + } + if constexpr (InWriter::kStoresEdges) { + in.PutSerial(dst, src, data, 0); + } + } + if constexpr (OutWriter::kTracksInputEdgeCount) { + out.RecordFilledEdges(valid_edges); + } + if constexpr (InWriter::kTracksInputEdgeCount) { + in.RecordFilledEdges(valid_edges); + } +} + +template +void fill_bulk_edge_chunk_concurrent( + const BulkEdgeEndpointScratch& endpoints, + const BulkEdgeDataReader& data_reader, OutWriter& out, + InWriter& in, BulkEdgeFillScratch& scratch) { + group_bulk_edge_chunk(endpoints, scratch); + constexpr bool kDirectOut = + OutWriter::kStoresEdges && !OutWriter::kNeedsConcurrentGrouping; + constexpr bool kDirectIn = + InWriter::kStoresEdges && !InWriter::kNeedsConcurrentGrouping; + if constexpr (kDirectOut || kDirectIn) { + size_t valid_edges = 0; + for (size_t row = 0; row < endpoints.src_lids.size(); ++row) { + const auto src = endpoints.src_lids[row]; + const auto dst = endpoints.dst_lids[row]; + if (!is_valid_bulk_edge(src, dst)) { + continue; + } + ++valid_edges; + const auto data = data_reader.Get(row); + if constexpr (kDirectOut) { + out.PutConcurrent(src, dst, data, 0); + } + if constexpr (kDirectIn) { + in.PutConcurrent(dst, src, data, 0); + } + } + if constexpr (OutWriter::kTracksInputEdgeCount) { + out.RecordFilledEdges(valid_edges); + } + if constexpr (InWriter::kTracksInputEdgeCount) { + in.RecordFilledEdges(valid_edges); + } + } + if constexpr (OutWriter::kNeedsConcurrentGrouping) { + for (const auto& [src, group] : scratch.out_groups) { + CHECK_LE(group.count, + static_cast(std::numeric_limits::max())); + auto slot = out.ReserveConcurrent(src, static_cast(group.count)); + auto edge_index = group.head; + for (uint32_t i = 0; i < group.count; ++i) { + CHECK_NE(edge_index, kInvalidBulkEdgeIndex); + const auto data = data_reader.Get(edge_index); + out.PutAt(src, slot++, endpoints.dst_lids[edge_index], data, 0); + edge_index = scratch.out_next[edge_index]; + } + CHECK_EQ(edge_index, kInvalidBulkEdgeIndex); + } + } + if constexpr (InWriter::kNeedsConcurrentGrouping) { + for (const auto& [dst, group] : scratch.in_groups) { + CHECK_LE(group.count, + static_cast(std::numeric_limits::max())); + auto slot = in.ReserveConcurrent(dst, static_cast(group.count)); + auto edge_index = group.head; + for (uint32_t i = 0; i < group.count; ++i) { + CHECK_NE(edge_index, kInvalidBulkEdgeIndex); + const auto data = data_reader.Get(edge_index); + in.PutAt(dst, slot++, endpoints.src_lids[edge_index], data, 0); + edge_index = scratch.in_next[edge_index]; + } + CHECK_EQ(edge_index, kInvalidBulkEdgeIndex); + } + } +} + +BulkEdgeFillResult profile_bulk_edge_fill( + const BulkEdgeEndpointScratch& endpoints, + const BulkEdgeFillScratch* concurrent_scratch, bool stores_out_direction, + bool stores_in_direction, bool groups_out_direction, + bool groups_in_direction) { + BulkEdgeFillResult result; + result.valid_edges = count_valid_bulk_edges(endpoints); + if (stores_out_direction) { + result.out_reservations = + concurrent_scratch == nullptr + ? result.valid_edges + : (groups_out_direction ? concurrent_scratch->out_groups.size() + : 0); + } + if (stores_in_direction) { + result.in_reservations = + concurrent_scratch == nullptr + ? result.valid_edges + : (groups_in_direction ? concurrent_scratch->in_groups.size() : 0); + } + return result; +} + +template +using BulkEdgeSerialFillChunk = std::function&)>; + +template +using BulkEdgeConcurrentFillChunk = std::function&, + BulkEdgeFillScratch&)>; + +template +struct BulkEdgeBuildOps { + bool stores_out_direction = false; + bool stores_in_direction = false; + bool counts_out_degree = false; + bool counts_in_degree = false; + bool checks_out_single = false; + bool checks_in_single = false; + bool groups_out_direction = false; + bool groups_in_direction = false; + bool supports_disjoint_concurrent_fill = false; + std::function prepare_build; + BulkEdgeCountChunk count_chunk; + std::function set_input_edge_count; + std::function allocate_from_counts; + BulkEdgeSerialFillChunk fill_serial_chunk; + BulkEdgeConcurrentFillChunk fill_concurrent_chunk; + std::function finish; +}; + +template +BulkEdgeBuildOps make_bulk_edge_build_ops(OutWriter& out, + InWriter& in) { + BulkEdgeBuildOps ops; + constexpr bool kStoresAnyDirection = + OutWriter::kStoresEdges || InWriter::kStoresEdges; + constexpr bool kNeedsAnyDegreeCount = + OutWriter::kNeedsDegreeCount || InWriter::kNeedsDegreeCount; + constexpr bool kSupportsCountPass = kNeedsAnyDegreeCount || + OutWriter::kChecksSingleUniqueness || + InWriter::kChecksSingleUniqueness; + ops.stores_out_direction = OutWriter::kStoresEdges; + ops.stores_in_direction = InWriter::kStoresEdges; + ops.counts_out_degree = OutWriter::kNeedsDegreeCount; + ops.counts_in_degree = InWriter::kNeedsDegreeCount; + ops.checks_out_single = OutWriter::kChecksSingleUniqueness; + ops.checks_in_single = InWriter::kChecksSingleUniqueness; + ops.groups_out_direction = OutWriter::kNeedsConcurrentGrouping; + ops.groups_in_direction = InWriter::kNeedsConcurrentGrouping; + ops.prepare_build = [&out, &in](vid_t out_vertices, vid_t in_vertices) { + out.PrepareBuild(out_vertices); + in.PrepareBuild(in_vertices); + }; + ops.allocate_from_counts = [&out, &in]() { + out.AllocateFromCounts(); + in.AllocateFromCounts(); + }; + if constexpr (kSupportsCountPass) { + ops.count_chunk = [&out, &in](BulkEdgeCountScratch& scratch, + const BulkEdgeCountConfig& config) { + return count_bulk_edge_chunk(scratch.endpoints, out, in, scratch, config); + }; + } + if constexpr (OutWriter::kTracksInputEdgeCount && + InWriter::kTracksInputEdgeCount) { + ops.set_input_edge_count = [&out, &in](uint64_t count) { + out.SetInputEdgeCount(count); + in.SetInputEdgeCount(count); + }; + } else if constexpr (OutWriter::kTracksInputEdgeCount) { + ops.set_input_edge_count = [&out](uint64_t count) { + out.SetInputEdgeCount(count); + }; + } else if constexpr (InWriter::kTracksInputEdgeCount) { + ops.set_input_edge_count = [&in](uint64_t count) { + in.SetInputEdgeCount(count); + }; + } + if constexpr (kStoresAnyDirection) { + ops.supports_disjoint_concurrent_fill = + OutWriter::kSupportsDisjointConcurrentFill && + InWriter::kSupportsDisjointConcurrentFill; + ops.fill_serial_chunk = + [&out, &in](const BulkEdgeEndpointScratch& endpoints, + const BulkEdgeDataReader& data_reader) { + fill_bulk_edge_chunk_serial(endpoints, data_reader, out, in); + }; + if constexpr (OutWriter::kSupportsDisjointConcurrentFill && + InWriter::kSupportsDisjointConcurrentFill) { + ops.fill_concurrent_chunk = + [&out, &in](const BulkEdgeEndpointScratch& endpoints, + const BulkEdgeDataReader& data_reader, + BulkEdgeFillScratch& scratch) { + fill_bulk_edge_chunk_concurrent(endpoints, data_reader, out, in, + scratch); + }; + } + } + ops.finish = [&out, &in]() { + out.Finish(); + in.Finish(); + }; + return ops; +} + +struct BulkEdgeFillWorkerStats { + size_t chunks = 0; + size_t edges = 0; + size_t out_reservations = 0; + size_t in_reservations = 0; + int64_t endpoint_time_ns = 0; + int64_t fill_time_ns = 0; +}; + +template +void fill_bulk_edges(const std::shared_ptr& source, + const IndexerType& src_indexer, + const IndexerType& dst_indexer, + const ChunkPipelineAllocation& allocation, + bool preserve_order, bool allow_concurrent_fill, + const BulkEdgeBuildOps& ops) { + CHECK(!allow_concurrent_fill || ops.supports_disjoint_concurrent_fill); + CHECK(!allow_concurrent_fill || !preserve_order); + const auto source_options = + make_bulk_edge_source_options(allocation, preserve_order); + auto supplier = source->Open(source_options); + CHECK(supplier != nullptr); + + const bool profile_stages = VLOG_IS_ON(1); + const auto pipeline_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + const auto worker_count = allocation.consumer_count; + std::vector worker_stats; + if (profile_stages) { + worker_stats.resize(static_cast(worker_count)); + } + bool concurrent_fill = false; + + if (allow_concurrent_fill) { + if (worker_count > 1) { + CHECK(static_cast(ops.fill_concurrent_chunk)); + concurrent_fill = true; + std::vector scratches( + static_cast(worker_count)); + auto fill = [&](int32_t worker, const std::shared_ptr& chunk) { + CHECK_GE(worker, 0); + CHECK_LT(worker, worker_count); + auto& scratch = scratches[static_cast(worker)]; + const auto endpoint_start = + profile_stages ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, + scratch.endpoints); + if (profile_stages) { + worker_stats[static_cast(worker)].endpoint_time_ns += + std::chrono::duration_cast( + std::chrono::steady_clock::now() - endpoint_start) + .count(); + } + const auto fill_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; + BulkEdgeDataReader data_reader(data_column); + ops.fill_concurrent_chunk(scratch.endpoints, data_reader, scratch); + if (profile_stages) { + const auto fill_end = std::chrono::steady_clock::now(); + const auto result = profile_bulk_edge_fill( + scratch.endpoints, &scratch, ops.stores_out_direction, + ops.stores_in_direction, ops.groups_out_direction, + ops.groups_in_direction); + auto& stats = worker_stats[static_cast(worker)]; + ++stats.chunks; + stats.edges += result.valid_edges; + stats.out_reservations += result.out_reservations; + stats.in_reservations += result.in_reservations; + stats.fill_time_ns += + std::chrono::duration_cast(fill_end - + fill_start) + .count(); + } + }; + + if (supplier->SupportsConcurrentGetNext()) { + consume_concurrent_supplier_indexed(*supplier, worker_count, fill); + } else { + ChunkPipelineOptions options; + options.consumer_count = worker_count; + options.queue_capacity = allocation.queue_capacity; + consume_chunk_pipeline_indexed(*supplier, options, fill); + } + } + } + + if (!concurrent_fill) { + BulkEdgeEndpointScratch endpoints; + while (auto chunk = supplier->GetNextChunk()) { + const auto endpoint_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, endpoints); + if (profile_stages) { + worker_stats.front().endpoint_time_ns += + std::chrono::duration_cast( + std::chrono::steady_clock::now() - endpoint_start) + .count(); + } + const auto fill_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; + BulkEdgeDataReader data_reader(data_column); + ops.fill_serial_chunk(endpoints, data_reader); + if (profile_stages) { + const auto fill_end = std::chrono::steady_clock::now(); + const auto result = profile_bulk_edge_fill( + endpoints, nullptr, ops.stores_out_direction, + ops.stores_in_direction, ops.groups_out_direction, + ops.groups_in_direction); + auto& stats = worker_stats.front(); + ++stats.chunks; + stats.edges += result.valid_edges; + stats.out_reservations += result.out_reservations; + stats.in_reservations += result.in_reservations; + stats.fill_time_ns += + std::chrono::duration_cast(fill_end - + fill_start) + .count(); + } + } + } + + const auto pipeline_end = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + if (!profile_stages) { + return; + } + const auto pipeline_wall_ms = + std::chrono::duration_cast(pipeline_end - + pipeline_start) + .count(); + size_t total_edges = 0; + size_t total_chunks = 0; + size_t out_reservations = 0; + size_t in_reservations = 0; + size_t min_worker_edges = std::numeric_limits::max(); + size_t max_worker_edges = 0; + int64_t endpoint_time_ns = 0; + int64_t fill_time_ns = 0; + for (const auto& stats : worker_stats) { + total_edges += stats.edges; + total_chunks += stats.chunks; + out_reservations += stats.out_reservations; + in_reservations += stats.in_reservations; + min_worker_edges = std::min(min_worker_edges, stats.edges); + max_worker_edges = std::max(max_worker_edges, stats.edges); + endpoint_time_ns += stats.endpoint_time_ns; + fill_time_ns += stats.fill_time_ns; + } + for (size_t worker = 0; worker < worker_stats.size(); ++worker) { + const auto& stats = worker_stats[worker]; + VLOG(2) << "Bulk edge fill worker: worker=" << worker + << ", chunks=" << stats.chunks << ", edges=" << stats.edges + << ", out_range_reservations=" << stats.out_reservations + << ", in_range_reservations=" << stats.in_reservations + << ", endpoint_index_ms=" << stats.endpoint_time_ns / 1000000 + << ", fill_ms=" << stats.fill_time_ns / 1000000; + } + VLOG(1) << "Bulk edge fill stages: pipeline_wall_ms=" << pipeline_wall_ms + << ", concurrent_fill=" << concurrent_fill + << ", fill_workers=" << (concurrent_fill ? worker_count : 1) + << ", chunks=" << total_chunks << ", edges=" << total_edges + << ", out_range_reservations=" << out_reservations + << ", in_range_reservations=" << in_reservations + << ", min_worker_edges=" << min_worker_edges + << ", max_worker_edges=" << max_worker_edges + << ", endpoint_index_worker_ms=" << endpoint_time_ns / 1000000 + << ", csr_fill_worker_ms=" << fill_time_ns / 1000000; +} + +vid_t csr_vertex_capacity(const IndexerType& indexer) { + const size_t capacity = indexer.capacity(); + CHECK_LE(capacity, static_cast(std::numeric_limits::max())) + << "CSR vertex capacity exceeds the vertex id range"; + return static_cast(capacity); +} + +template +void build_bundled_edges_with_ops( + const BulkEdgeBuildOps& ops, const IndexerType& src_indexer, + const IndexerType& dst_indexer, + const std::shared_ptr& source) { + // VertexTable owns its reserve policy. Edge storage consumes the resulting + // indexer capacity instead of duplicating PropertyGraph::Dump's policy. + ops.prepare_build(csr_vertex_capacity(src_indexer), + csr_vertex_capacity(dst_indexer)); + if (!ops.stores_out_direction && !ops.stores_in_direction) { + ops.allocate_from_counts(); + ops.finish(); + return; + } + const auto estimated_bytes = source->EstimatedBytes(); + const auto count_allocation = resolve_chunk_pipeline_allocation( + estimated_bytes, source->ParallelEnabled(), false); + const bool needs_degree_count = ops.counts_out_degree || ops.counts_in_degree; + const bool has_single_direction = + ops.checks_out_single || ops.checks_in_single; + // Single-only layouts retain their one-pass ordered path. Mixed + // Single/Mutable layouts already need the degree pass, so that pass also + // verifies whether fixed Single slots are disjoint across all chunks. + const bool check_single_uniqueness = + has_single_direction && needs_degree_count; + const bool run_count_pass = needs_degree_count || check_single_uniqueness; + VLOG(1) << "Bulk edge count allocation: bytes=" << estimated_bytes + << ", producers=" << count_allocation.producer_count + << ", consumers=" << count_allocation.consumer_count + << ", queue_capacity=" << count_allocation.queue_capacity + << ", degree_count_pass=" << needs_degree_count + << ", single_uniqueness_check=" << check_single_uniqueness + << ", count_pass=" << run_count_pass; + BulkEdgeCountSummary count_summary; + if (run_count_pass) { + CHECK(static_cast(ops.count_chunk)); + count_summary = count_bulk_edges( + source, src_indexer, dst_indexer, count_allocation, + ops.counts_out_degree, ops.counts_in_degree, + check_single_uniqueness && ops.checks_out_single, + check_single_uniqueness && ops.checks_in_single, ops.count_chunk); + } + const bool single_uniqueness_verified = + has_single_direction && check_single_uniqueness; + if (single_uniqueness_verified) { + CHECK(static_cast(ops.set_input_edge_count)); + ops.set_input_edge_count(count_summary.valid_edges); + } + const bool single_duplicate = + (ops.checks_out_single && count_summary.out_single_duplicate) || + (ops.checks_in_single && count_summary.in_single_duplicate); + const bool single_slots_disjoint = + !has_single_direction || + (single_uniqueness_verified && !single_duplicate); + const bool allow_concurrent_fill = + ops.supports_disjoint_concurrent_fill && single_slots_disjoint; + const bool preserve_fill_order = + has_single_direction && !single_slots_disjoint; + const auto fill_allocation = + preserve_fill_order + ? resolve_chunk_pipeline_allocation(estimated_bytes, + source->ParallelEnabled(), true) + : count_allocation; + VLOG(1) << "Bulk edge fill allocation: bytes=" << estimated_bytes + << ", producers=" << fill_allocation.producer_count + << ", consumers=" << fill_allocation.consumer_count + << ", queue_capacity=" << fill_allocation.queue_capacity + << ", single_uniqueness_verified=" << single_uniqueness_verified + << ", out_single_duplicate=" << count_summary.out_single_duplicate + << ", in_single_duplicate=" << count_summary.in_single_duplicate + << ", allow_concurrent_fill=" << allow_concurrent_fill + << ", preserve_order=" << preserve_fill_order; + const bool profile_stages = VLOG_IS_ON(1); + const auto allocate_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + ops.allocate_from_counts(); + const auto fill_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + fill_bulk_edges(source, src_indexer, dst_indexer, fill_allocation, + preserve_fill_order, allow_concurrent_fill, ops); + ops.finish(); + if (!profile_stages) { + return; + } + const auto fill_end = std::chrono::steady_clock::now(); + VLOG(1) << "Bulk edge CSR stages: allocate_ms=" + << std::chrono::duration_cast( + fill_start - allocate_start) + .count() + << ", fill_pass_ms=" + << std::chrono::duration_cast(fill_end - + fill_start) + .count(); +} + +template +bool build_bundled_edges_typed( + CsrBase* out_csr, CsrBase* in_csr, const IndexerType& src_indexer, + const IndexerType& dst_indexer, + const std::shared_ptr& source) { + bool built = false; + const bool out_supported = + with_csr_bulk_writer(out_csr, [&](auto& out) { + const bool in_supported = + with_csr_bulk_writer(in_csr, [&](auto& in) { + auto ops = make_bulk_edge_build_ops(out, in); + build_bundled_edges_with_ops(ops, src_indexer, dst_indexer, + source); + }); + built = in_supported; + }); + return out_supported && built; +} + +bool build_bundled_edges(CsrBase* out_csr, CsrBase* in_csr, + const std::shared_ptr& schema, + const IndexerType& src_indexer, + const IndexerType& dst_indexer, + const std::shared_ptr& source) { + const auto property_type = schema->properties.empty() + ? DataTypeId::kEmpty + : schema->properties[0].id(); + switch (property_type) { +#define TYPE_DISPATCHER(enum_val, cpp_type) \ + case DataTypeId::enum_val: \ + return build_bundled_edges_typed(out_csr, in_csr, src_indexer, \ + dst_indexer, source); + FOR_EACH_DATA_TYPE_NO_STRING(TYPE_DISPATCHER) +#undef TYPE_DISPATCHER + case DataTypeId::kEmpty: + return build_bundled_edges_typed(out_csr, in_csr, src_indexer, + dst_indexer, source); + default: + return false; + } +} + void EdgeTable::Init(std::shared_ptr ckp, MemoryLevel level) { CHECK(meta_ != nullptr) << "EdgeTable::Init requires schema"; @@ -766,11 +1845,34 @@ std::pair EdgeTable::AddEdge( void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, const IndexerType& dst_indexer, std::shared_ptr supplier) { - in_csr_->resize(dst_indexer.size()); - out_csr_->resize(src_indexer.size()); + const bool profile_stages = VLOG_IS_ON(1); + const auto total_start = std::chrono::steady_clock::now(); + int64_t row_num_time_ns = 0; + int64_t collect_time_ns = 0; + int64_t filter_time_ns = 0; + int64_t ensure_capacity_time_ns = 0; + int64_t write_time_ns = 0; + size_t chunk_count = 0; + size_t input_rows = 0; + size_t cached_property_columns = 0; + size_t peak_buffer_bytes = 0; + const auto peak_rss_start_bytes = + profile_stages ? process_peak_rss_bytes() : 0; + // Keep fallback COPY paths aligned with the vertex table's actual capacity, + // while leaving completely unloaded edge tables lazy until persistence. + in_csr_->resize(csr_vertex_capacity(dst_indexer)); + out_csr_->resize(csr_vertex_capacity(src_indexer)); std::vector src_lid, dst_lid; // Pre-reserve capacity to reduce vector reallocation on large graphs. + const auto row_num_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; auto total_rows = supplier->RowNum(); + if (profile_stages) { + row_num_time_ns = std::chrono::duration_cast( + std::chrono::steady_clock::now() - row_num_start) + .count(); + } if (total_rows > 0) { src_lid.reserve(total_rows); dst_lid.reserve(total_rows); @@ -779,11 +1881,17 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, // for unbundled: full property DataChunks). std::vector> bundled_data_cols; std::vector> unbundled_data_chunks; + std::vector valid_flags; while (true) { + const auto collect_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; auto chunk = supplier->GetNextChunk(); if (chunk == nullptr) { break; } + ++chunk_count; + input_rows += chunk->row_num(); auto src_col = chunk->get(0); auto dst_col = chunk->get(1); src_indexer.get_index(*src_col, src_lid); @@ -792,6 +1900,7 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, if (meta_->is_bundled()) { // Bundled: only one property column (index 2). bundled_data_cols.push_back(chunk->get(2)); + ++cached_property_columns; } else { // Unbundled: collect remaining columns as a DataChunk. auto prop_chunk = std::make_shared(); @@ -799,14 +1908,40 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, auto c = chunk->get(static_cast(i)); if (c) { prop_chunk->set(static_cast(i - 2), c); + ++cached_property_columns; } } unbundled_data_chunks.push_back(prop_chunk); } } + if (profile_stages) { + collect_time_ns += std::chrono::duration_cast( + std::chrono::steady_clock::now() - collect_start) + .count(); + peak_buffer_bytes = std::max( + peak_buffer_bytes, estimate_edge_fallback_buffer_bytes( + src_lid, dst_lid, valid_flags, + bundled_data_cols, unbundled_data_chunks)); + } } - std::vector valid_flags; + const auto filter_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; filterInvalidEdges(src_lid, dst_lid, valid_flags); + if (profile_stages) { + filter_time_ns = std::chrono::duration_cast( + std::chrono::steady_clock::now() - filter_start) + .count(); + peak_buffer_bytes = std::max( + peak_buffer_bytes, estimate_edge_fallback_buffer_bytes( + src_lid, dst_lid, valid_flags, bundled_data_cols, + unbundled_data_chunks)); + } + const auto valid_rows = src_lid.size(); + const auto invalid_rows = valid_flags.size() - valid_rows; + const auto ensure_capacity_start = + profile_stages ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; size_t new_size = table_idx_.load() + src_lid.size(); if (new_size >= Capacity()) { auto new_cap = new_size; @@ -815,6 +1950,15 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, } EnsureCapacity(new_cap); } + if (profile_stages) { + ensure_capacity_time_ns = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - ensure_capacity_start) + .count(); + } + const auto write_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; if (meta_->is_bundled()) { batch_add_bundled_edges_impl(out_csr_.get(), in_csr_.get(), meta_, src_lid, dst_lid, bundled_data_cols, valid_flags); @@ -826,6 +1970,48 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, src_lid, dst_lid, oe_csr, ie_csr, table_.get(), table_idx_, capacity_, meta_->properties, unbundled_data_chunks, valid_flags); } + if (profile_stages) { + write_time_ns = std::chrono::duration_cast( + std::chrono::steady_clock::now() - write_start) + .count(); + const auto total_time_ns = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - total_start) + .count(); + const auto peak_rss_end_bytes = process_peak_rss_bytes(); + VLOG(1) << "Fallback edge load stages: input_rows=" << input_rows + << ", valid_rows=" << valid_rows + << ", invalid_rows=" << invalid_rows << ", chunks=" << chunk_count + << ", cached_property_columns=" << cached_property_columns + << ", peak_buffer_bytes=" << peak_buffer_bytes + << ", process_peak_rss_start_bytes=" << peak_rss_start_bytes + << ", process_peak_rss_end_bytes=" << peak_rss_end_bytes + << ", total_ms=" << total_time_ns / 1000000 + << ", row_num_ms=" << row_num_time_ns / 1000000 + << ", collect_lookup_ms=" << collect_time_ns / 1000000 + << ", filter_ms=" << filter_time_ns / 1000000 + << ", ensure_capacity_ms=" << ensure_capacity_time_ns / 1000000 + << ", write_ms=" << write_time_ns / 1000000; + } +} + +void EdgeTable::BatchBuildEdges(const IndexerType& src_indexer, + const IndexerType& dst_indexer, + std::shared_ptr source) { + if (!source) { + THROW_INVALID_ARGUMENT_EXCEPTION( + "BatchBuildEdges requires a non-null data source"); + } + CHECK(source->rewindable()); + CHECK(CanBatchBuild()) + << "Bulk edge build requires an empty bundled edge table"; + + EdgeTable staged(meta_); + staged.Init(ckp_, memory_level_); + CHECK(build_bundled_edges(staged.out_csr_.get(), staged.in_csr_.get(), meta_, + src_indexer, dst_indexer, source)) + << "Bulk edge build does not support this CSR layout"; + Swap(staged); } void EdgeTable::BatchAddEdges( diff --git a/src/storages/graph/graph_interface.cc b/src/storages/graph/graph_interface.cc index eba93c230..8c2150924 100644 --- a/src/storages/graph/graph_interface.cc +++ b/src/storages/graph/graph_interface.cc @@ -120,6 +120,20 @@ Status StorageAPUpdateInterface::BatchAddVertices( return graph_.BatchAddVertices(v_label_id, std::move(supplier)); } +bool StorageAPUpdateInterface::CanBatchBuildVertices(label_t v_label_id) const { + return graph_.CanBatchBuildVertices(v_label_id); +} + +Status StorageAPUpdateInterface::BatchBuildVertices( + label_t v_label_id, std::shared_ptr source) { + auto status = graph_.BatchBuildVertices(v_label_id, std::move(source)); + if (status.ok()) { + // The staged table swap replaces raw pointers cached by GraphView. + mut_view_.Rebuild(graph_); + } + return status; +} + Status StorageAPUpdateInterface::BatchAddEdges( label_t src_label, label_t dst_label, label_t edge_label, std::shared_ptr supplier) { @@ -127,6 +141,24 @@ Status StorageAPUpdateInterface::BatchAddEdges( std::move(supplier)); } +bool StorageAPUpdateInterface::CanBatchBuildEdges(label_t src_label, + label_t dst_label, + label_t edge_label) const { + return graph_.CanBatchBuildEdges(src_label, dst_label, edge_label); +} + +Status StorageAPUpdateInterface::BatchBuildEdges( + label_t src_label, label_t dst_label, label_t edge_label, + std::shared_ptr source) { + auto status = graph_.BatchBuildEdges(src_label, dst_label, edge_label, + std::move(source)); + if (status.ok()) { + // The staged CSR swap replaces raw pointers cached by GraphView. + mut_view_.Rebuild(graph_); + } + return status; +} + Status StorageAPUpdateInterface::BatchDeleteVertices( label_t v_label_id, const std::vector& vids) { return graph_.BatchDeleteVertices(v_label_id, vids); diff --git a/src/storages/graph/property_graph.cc b/src/storages/graph/property_graph.cc index 54d9b1631..635974784 100644 --- a/src/storages/graph/property_graph.cc +++ b/src/storages/graph/property_graph.cc @@ -141,6 +141,23 @@ Status PropertyGraph::BatchAddVertices( return neug::Status::OK(); } +bool PropertyGraph::CanBatchBuildVertices(label_t v_label) const { + return vertex_label_check(v_label).ok() && + vertex_tables_[v_label].CanBatchBuild(); +} + +Status PropertyGraph::BatchBuildVertices( + label_t v_label, std::shared_ptr source) { + RETURN_IF_NOT_OK(vertex_label_check(v_label)); + if (!source || !source->rewindable() || !CanBatchBuildVertices(v_label)) { + return Status(StatusCode::ERR_NOT_SUPPORTED, + "Bulk vertex build requires an empty label and a " + "repeatable source."); + } + vertex_tables_[v_label].BatchBuildVertices(std::move(source)); + return neug::Status::OK(); +} + Status PropertyGraph::BatchAddEdges( label_t src_v_label, label_t dst_v_label, label_t e_label, std::shared_ptr supplier) { @@ -153,6 +170,34 @@ Status PropertyGraph::BatchAddEdges( return neug::Status::OK(); } +bool PropertyGraph::CanBatchBuildEdges(label_t src_v_label, label_t dst_v_label, + label_t e_label) const { + if (!edge_triplet_check(src_v_label, dst_v_label, e_label).ok()) { + return false; + } + const auto index = + schema_.generate_edge_label(src_v_label, dst_v_label, e_label); + auto it = edge_tables_.find(index); + return it != edge_tables_.end() && it->second.CanBatchBuild(); +} + +Status PropertyGraph::BatchBuildEdges( + label_t src_v_label, label_t dst_v_label, label_t e_label, + std::shared_ptr source) { + RETURN_IF_NOT_OK(edge_triplet_check(src_v_label, dst_v_label, e_label)); + if (!source || !source->rewindable() || + !CanBatchBuildEdges(src_v_label, dst_v_label, e_label)) { + return Status(StatusCode::ERR_NOT_SUPPORTED, + "Bulk edge build requires an empty bundled edge table and " + "a repeatable source."); + } + auto index = schema_.generate_edge_label(src_v_label, dst_v_label, e_label); + edge_tables_.at(index).BatchBuildEdges( + vertex_tables_.at(src_v_label).get_indexer(), + vertex_tables_.at(dst_v_label).get_indexer(), std::move(source)); + return neug::Status::OK(); +} + Status PropertyGraph::CreateVertexType(const CreateVertexTypeParam& config) { if (schema_.is_vertex_label_valid(config.GetVertexLabel())) { return Status(StatusCode::ERR_SCHEMA_MISMATCH, diff --git a/src/storages/graph/vertex_table.cc b/src/storages/graph/vertex_table.cc index e1150fc22..031bafc65 100644 --- a/src/storages/graph/vertex_table.cc +++ b/src/storages/graph/vertex_table.cc @@ -15,6 +15,8 @@ #include "neug/storages/graph/vertex_table.h" +#include + #include "neug/storages/checkpoint_manifest.h" #include "neug/storages/module/module_broker.h" #include "neug/storages/module/module_factory.h" @@ -47,21 +49,66 @@ void VertexTable::Init(std::shared_ptr ckp, MemoryLevel level) { void VertexTable::insert_vertices( std::shared_ptr supplier) { + const bool profile_stages = VLOG_IS_ON(1); + const auto total_start = std::chrono::steady_clock::now(); + int64_t row_num_time_ns = 0; + int64_t reserve_time_ns = 0; + int64_t get_chunk_time_ns = 0; + int64_t pk_time_ns = 0; + int64_t property_time_ns = 0; + size_t chunk_count = 0; + size_t loaded_rows = 0; + auto reserve_checkpoint_headroom = [this](size_t required_size) { + const size_t headroom = required_size / 4; + if (headroom > std::numeric_limits::max() - required_size) { + THROW_RUNTIME_ERROR("Vertex capacity overflow"); + } + const size_t required_capacity = required_size + headroom; + if (required_capacity > indexer_->capacity()) { + EnsureCapacity(required_capacity); + } + }; + + const auto row_num_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; auto row_nums = supplier->RowNum(); + if (profile_stages) { + row_num_time_ns = std::chrono::duration_cast( + std::chrono::steady_clock::now() - row_num_start) + .count(); + } if (row_nums < 0) { VLOG(1) << "Row number from supplier is unknown, skip pre-reserve."; row_nums = 0; } - size_t new_size = indexer_->size() + row_nums; - if (new_size > indexer_->capacity()) { - size_t cap = indexer_->capacity(); - while (new_size >= cap) { - cap = cap < 4096 ? 4096 : cap + cap / 4; - } - EnsureCapacity(cap); + if (static_cast(row_nums) > std::numeric_limits::max()) { + THROW_RUNTIME_ERROR("Vertex row count exceeds addressable capacity"); + } + const auto row_count = static_cast(row_nums); + if (row_count > std::numeric_limits::max() - indexer_->size()) { + THROW_RUNTIME_ERROR("Vertex row count overflow"); + } + const auto pre_reserve_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + reserve_checkpoint_headroom(indexer_->size() + row_count); + if (profile_stages) { + reserve_time_ns += std::chrono::duration_cast( + std::chrono::steady_clock::now() - pre_reserve_start) + .count(); } while (true) { + const auto get_chunk_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; auto chunk = supplier->GetNextChunk(); + if (profile_stages) { + get_chunk_time_ns += + std::chrono::duration_cast( + std::chrono::steady_clock::now() - get_chunk_start) + .count(); + } if (chunk == nullptr) { break; } @@ -85,24 +132,100 @@ void VertexTable::insert_vertices( // Capacity check for actual batch size. size_t chunk_rows = chunk->row_num(); + if (chunk_rows > std::numeric_limits::max() - indexer_->size()) { + THROW_RUNTIME_ERROR("Vertex row count overflow"); + } size_t new_size = indexer_->size() + chunk_rows; - if (new_size > indexer_->capacity()) { - size_t cap = indexer_->capacity(); - while (new_size >= cap) { - cap = cap < 4096 ? 4096 : cap + cap / 4; - } - EnsureCapacity(cap); + const auto chunk_reserve_start = + profile_stages ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + reserve_checkpoint_headroom(new_size); + if (profile_stages) { + reserve_time_ns += + std::chrono::duration_cast( + std::chrono::steady_clock::now() - chunk_reserve_start) + .count(); } + const auto pk_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; auto vids = insert_primary_keys(pk_col); + if (profile_stages) { + pk_time_ns += std::chrono::duration_cast( + std::chrono::steady_clock::now() - pk_start) + .count(); + } + const auto property_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; for (size_t i = 0; i < prop_cols.size(); ++i) { auto col = table_->get_column_by_id(i); set_properties_from_context_column(col, prop_cols[i], vids); } + if (profile_stages) { + property_time_ns += std::chrono::duration_cast( + std::chrono::steady_clock::now() - property_start) + .count(); + ++chunk_count; + loaded_rows += chunk_rows; + } VLOG(10) << "Inserted " << chunk_rows << " vertices, current vertex num: " << VertexNum(); } + if (profile_stages) { + const auto total_time_ns = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - total_start) + .count(); + VLOG(1) << "Vertex load stages: rows=" << loaded_rows + << ", chunks=" << chunk_count + << ", total_ms=" << total_time_ns / 1000000 + << ", row_num_ms=" << row_num_time_ns / 1000000 + << ", reserve_ms=" << reserve_time_ns / 1000000 + << ", get_chunk_ms=" << get_chunk_time_ns / 1000000 + << ", insert_pk_ms=" << pk_time_ns / 1000000 + << ", set_properties_ms=" << property_time_ns / 1000000; + } +} + +void VertexTable::BatchBuildVertices(std::shared_ptr source) { + if (!source) { + THROW_INVALID_ARGUMENT_EXCEPTION( + "BatchBuildVertices requires a non-null data source"); + } + CHECK(source->rewindable()); + CHECK(CanBatchBuild()) + << "Bulk vertex build requires an empty destination table"; + + auto supplier = source->Open(); + if (!supplier) { + THROW_INTERNAL_EXCEPTION("Data source returned a null supplier"); + } + const auto row_num = supplier->RowNum(); + + VertexTable staged(vertex_schema_); + staged.Init(ckp_, memory_level_); + if (row_num >= 0) { + if (static_cast(row_num) > std::numeric_limits::max()) { + THROW_RUNTIME_ERROR("Vertex row count exceeds addressable capacity"); + } + const auto row_count = static_cast(row_num); + // Keep the same headroom that PropertyGraph::DisassembleTo requires before + // writing a checkpoint. This avoids a second resize after the bulk build. + const size_t checkpoint_headroom = row_count / 4; + if (checkpoint_headroom > std::numeric_limits::max() - row_count) { + THROW_RUNTIME_ERROR("Vertex bulk load capacity overflow"); + } + staged.EnsureCapacity(row_count + checkpoint_headroom); + staged.insert_vertices_preallocated(std::move(supplier)); + } else { + // Unknown-cardinality suppliers still stream once and grow only when an + // actual chunk proves that more capacity is needed. + staged.insert_vertices(std::move(supplier)); + } + Swap(staged); } void VertexTable::Close() { diff --git a/src/storages/loader/loader_utils.cc b/src/storages/loader/loader_utils.cc index 4c2d0e8ba..b8ae84ffe 100644 --- a/src/storages/loader/loader_utils.cc +++ b/src/storages/loader/loader_utils.cc @@ -22,20 +22,28 @@ #include #include +#include +#include #include #include +#include #include +#include +#include #include #include #include #include +#include #include #include #include +#include #include #include #include #include +#include #include #include @@ -43,6 +51,7 @@ #include "neug/common/columns/columns_utils.h" #include "neug/common/columns/value_columns.h" #include "neug/common/types/value.h" +#include "neug/storages/loader/chunk_pipeline_utils.h" #include "neug/utils/datetime_parsers.h" #include "neug/utils/exception/exception.h" #include "neug/utils/property/column.h" @@ -50,6 +59,38 @@ namespace neug { +struct CsvPartitionRange { + size_t start = 0; + size_t end = 0; + int64_t start_row = 0; +}; + +struct CsvRangeTask { + std::string file_path; + CsvPartitionRange range; + int64_t skip_rows = 0; +}; + +struct CsvPartitionPlan { + int64_t row_count = 0; + std::vector tasks; +}; + +struct CsvPartitionPlanCache { + std::shared_ptr GetOrCreate( + const std::vector& file_paths, const CsvReadConfig& config, + int32_t producer_count); + + private: + struct Entry { + std::once_flag once; + std::shared_ptr plan; + }; + + std::mutex mutex; + std::unordered_map> entries; +}; + namespace { constexpr size_t kDefaultCsvChunkRows = 4096; @@ -80,6 +121,7 @@ csv::CSVFormat build_csv_format(const CsvReadConfig& config) { csv_format.variable_columns(csv::VariableColumnPolicy::KEEP); csv_format.no_header(); csv_format.column_names(config.column_names); + csv_format.threading(config.use_threads); return csv_format; } @@ -416,6 +458,11 @@ struct RowCounterState { /// Fast CSV row counter: scans raw bytes to count row boundaries /// without parsing fields. Uses a quote-tracking state machine; /// parallelizes via speculative dual-state-machine scan for large files. +struct CsvScanResult { + int64_t row_count = 0; + std::vector ranges; +}; + class CsvRowCountCounter { public: // rows_to_skip is intentionally NOT a parameter: the counter counts @@ -423,12 +470,18 @@ class CsvRowCountCounter { // separately, and RowNum() is only a pre-allocation hint, so a slight // overcount (by at most skip_rows, typically 1 for header) is safe. CsvRowCountCounter(std::string file_path, bool quoting, char quote_char, - bool double_quote, char delimiter) + bool /*double_quote*/, char delimiter, + bool use_threads = true) : file_path_(std::move(file_path)), quoting_(quoting), quote_char_(quote_char), - double_quote_(double_quote), - delimiter_(delimiter) {} + // csv::CSVReader always uses RFC-style doubled quotes and exposes no + // switch for disabling that behavior. Partition planning must follow + // the parser's actual DFA even when the legacy DOUBLE_QUOTE option is + // false, otherwise a range can start in the middle of a record. + double_quote_(true), + delimiter_(delimiter), + use_threads_(use_threads) {} int64_t count() const { struct stat st; @@ -440,7 +493,8 @@ class CsvRowCountCounter { return 0; constexpr size_t kMinChunkSize = 4 << 20; // 4 MB - unsigned num_threads = std::thread::hardware_concurrency(); + unsigned num_threads = + use_threads_ ? std::thread::hardware_concurrency() : 1; if (num_threads == 0) num_threads = 1; if (file_size < kMinChunkSize * num_threads) { @@ -452,6 +506,49 @@ class CsvRowCountCounter { return count_parallel(file_size, num_threads); } + CsvScanResult ScanWithRanges(int32_t requested_partitions, + int32_t scan_threads) const { + struct stat st; + if (stat(file_path_.c_str(), &st) != 0) { + THROW_IO_EXCEPTION("Failed to get file size: " + file_path_); + } + const auto file_size = static_cast(st.st_size); + CsvScanResult result; + if (file_size == 0) { + return result; + } + const auto target_partitions = static_cast(std::max( + 1, std::min(requested_partitions, + static_cast(file_size)))); + const auto workers = static_cast(std::max( + 1, std::min(scan_threads, static_cast(file_size)))); + if (workers == 1) { + result.row_count = count_single(file_size); + result.ranges.push_back({0, file_size, 0}); + return result; + } + result = scan_parallel(file_size, workers, true); + if (result.ranges.size() <= target_partitions) { + return result; + } + + // The scan may use all hardware workers even when the parser only needs P + // producer ranges. Merge adjacent record-aligned scan ranges so planning + // parallelism and producer parallelism remain independent. + std::vector merged; + merged.reserve(target_partitions); + const auto scanned_ranges = result.ranges.size(); + for (size_t partition = 0; partition < target_partitions; ++partition) { + const auto begin = partition * scanned_ranges / target_partitions; + const auto end = (partition + 1) * scanned_ranges / target_partitions; + CHECK_LT(begin, end); + merged.push_back({result.ranges[begin].start, result.ranges[end - 1].end, + result.ranges[begin].start_row}); + } + result.ranges = std::move(merged); + return result; + } + private: struct ChunkResult { RowCounterState outside; // assumed start outside quotes @@ -504,8 +601,22 @@ class CsvRowCountCounter { if (n == 0) break; for (size_t i = 0; i < n; ++i) { - if (buf[i] == '\n' || buf[i] == '\r') + if (buf[i] == '\n') { return pos + i + 1; + } + if (buf[i] == '\r') { + const auto boundary = pos + i + 1; + if (i + 1 < n && buf[i + 1] == '\n') { + return boundary + 1; + } + // file.read() has already advanced beyond this buffer. When CR is + // its last byte, inspect the next byte so a CRLF pair never gets + // split across two parser ranges. + if (i + 1 == n && boundary < end && file.peek() == '\n') { + return boundary + 1; + } + return boundary; + } } pos += n; } @@ -557,6 +668,45 @@ class CsvRowCountCounter { /// Parallel scan with speculative dual state machines. int64_t count_parallel(size_t file_size, unsigned num_threads) const { + return scan_parallel(file_size, num_threads, false).row_count; + } + + CsvPartitionRange find_record_boundary(size_t start, size_t end, + int64_t start_row) const { + RowCounterState state; + state.init(true, quoting_, quote_char_, double_quote_, delimiter_); + size_t boundary = end; + bool found = false; + size_t position = start; + scan_range(start, end, [&](const char* data, size_t n) { + if (found) { + position += n; + return; + } + for (size_t i = 0; i < n; ++i) { + const auto before = state.count; + const char c = data[i]; + state.step(c); + if (state.count != before) { + boundary = position + i + 1; + found = true; + break; + } + } + position += n; + }); + if (found && boundary < end) { + std::ifstream file(file_path_, std::ios::binary); + file.seekg(static_cast(boundary)); + if (file && file.peek() == '\n') { + ++boundary; + } + } + return {boundary, end, start_row + state.count}; + } + + CsvScanResult scan_parallel(size_t file_size, unsigned num_threads, + bool build_ranges) const { // Compute newline-aligned chunk boundaries. size_t approx_chunk = file_size / num_threads; std::vector bounds; @@ -572,20 +722,36 @@ class CsvRowCountCounter { // Parallel scan: each thread scans its chunk with dual state machines. std::vector results(actual_threads); + std::vector errors(actual_threads); std::vector threads; - for (unsigned i = 0; i < actual_threads; ++i) { - threads.emplace_back([this, &results, i, &bounds]() { - results[i] = scan_chunk(bounds[i], bounds[i + 1]); + threads.reserve(actual_threads > 0 ? actual_threads - 1 : 0); + for (unsigned i = 1; i < actual_threads; ++i) { + threads.emplace_back([this, &results, &errors, i, &bounds]() { + try { + results[i] = scan_chunk(bounds[i], bounds[i + 1]); + } catch (...) { errors[i] = std::current_exception(); } }); } + try { + results[0] = scan_chunk(bounds[0], bounds[1]); + } catch (...) { errors[0] = std::current_exception(); } for (auto& t : threads) t.join(); + for (const auto& error : errors) { + if (error) { + std::rethrow_exception(error); + } + } // Sequential resolution: chain quote state across chunks. int64_t total = 0; bool in_quotes = false; // chunk 0 starts outside quotes const RowCounterState* last_selected = nullptr; + std::vector starts_inside(actual_threads, false); + std::vector rows_before(actual_threads, 0); for (unsigned i = 0; i < actual_threads; ++i) { + starts_inside[i] = in_quotes; + rows_before[i] = total; const RowCounterState& s = in_quotes ? results[i].inside : results[i].outside; total += s.count; @@ -599,7 +765,30 @@ class CsvRowCountCounter { if (last_selected && last_selected->has_content) ++total; - return total; + CsvScanResult output; + output.row_count = total; + if (!build_ranges) { + return output; + } + + std::vector> safe_bounds; + safe_bounds.emplace_back(0, 0); + for (unsigned i = 1; i < actual_threads; ++i) { + CsvPartitionRange safe{bounds[i], file_size, rows_before[i]}; + if (starts_inside[i]) { + safe = find_record_boundary(bounds[i], file_size, rows_before[i]); + } + if (safe.start > safe_bounds.back().first && safe.start < file_size) { + safe_bounds.emplace_back(safe.start, safe.start_row); + } + } + safe_bounds.emplace_back(file_size, total); + output.ranges.reserve(safe_bounds.size() - 1); + for (size_t i = 0; i + 1 < safe_bounds.size(); ++i) { + output.ranges.push_back({safe_bounds[i].first, safe_bounds[i + 1].first, + safe_bounds[i].second}); + } + return output; } std::string file_path_; @@ -607,13 +796,273 @@ class CsvRowCountCounter { char quote_char_; bool double_quote_; char delimiter_; + bool use_threads_; +}; + +class BoundedFileStreamBuf final : public std::streambuf { + public: + BoundedFileStreamBuf(const std::string& file_path, size_t start, size_t end) + : file_(file_path, std::ios::binary), remaining_(end - start) { + if (!file_.is_open()) { + THROW_IO_EXCEPTION("Failed to open CSV range: " + file_path); + } + file_.seekg(static_cast(start)); + if (!file_) { + THROW_IO_EXCEPTION("Failed to seek CSV range: " + file_path); + } + setg(buffer_.data(), buffer_.data(), buffer_.data()); + } + + protected: + int_type underflow() override { + if (gptr() < egptr()) { + return traits_type::to_int_type(*gptr()); + } + if (remaining_ == 0) { + return traits_type::eof(); + } + const auto requested = std::min(remaining_, buffer_.size()); + file_.read(buffer_.data(), static_cast(requested)); + const auto read = static_cast(file_.gcount()); + if (read == 0) { + remaining_ = 0; + return traits_type::eof(); + } + remaining_ -= read; + setg(buffer_.data(), buffer_.data(), buffer_.data() + read); + return traits_type::to_int_type(*gptr()); + } + + private: + std::ifstream file_; + size_t remaining_; + std::array buffer_{}; +}; + +class BoundedFileStream final : public std::istream { + public: + BoundedFileStream(const std::string& file_path, size_t start, size_t end) + : std::istream(nullptr), buffer_(file_path, start, end) { + rdbuf(&buffer_); + } + + private: + BoundedFileStreamBuf buffer_; }; } // namespace +std::shared_ptr CsvPartitionPlanCache::GetOrCreate( + const std::vector& file_paths, const CsvReadConfig& config, + int32_t producer_count) { + producer_count = + std::clamp(producer_count, 1, hardware_worker_count()); + std::shared_ptr entry; + { + std::lock_guard lock(mutex); + auto& cached = entries[producer_count]; + if (!cached) { + cached = std::make_shared(); + } + entry = cached; + } + + std::call_once(entry->once, [&] { + auto plan = std::make_shared(); + + // producer_count is a budget for the complete COPY, not for every input + // file. Give every non-empty file one range, then split the currently + // largest range until the global producer budget is reached. + std::vector file_sizes(file_paths.size(), 0); + std::vector range_counts(file_paths.size(), 0); + size_t non_empty_files = 0; + for (size_t i = 0; i < file_paths.size(); ++i) { + std::error_code error; + file_sizes[i] = std::filesystem::file_size(file_paths[i], error); + if (error) { + THROW_IO_EXCEPTION("Failed to get file size: " + file_paths[i]); + } + if (file_sizes[i] > 0) { + range_counts[i] = 1; + ++non_empty_files; + } + } + + size_t assigned_ranges = non_empty_files; + const size_t target_ranges = + std::max(non_empty_files, static_cast(producer_count)); + while (assigned_ranges < target_ranges) { + size_t best = file_paths.size(); + long double best_range_bytes = -1; + for (size_t i = 0; i < file_paths.size(); ++i) { + if (range_counts[i] <= 0 || + static_cast(range_counts[i]) >= file_sizes[i]) { + continue; + } + const auto range_bytes = + static_cast(file_sizes[i]) / range_counts[i]; + if (range_bytes > best_range_bytes) { + best = i; + best_range_bytes = range_bytes; + } + } + if (best == file_paths.size()) { + break; + } + ++range_counts[best]; + ++assigned_ranges; + } + + // Row counting/planning is a separate phase, so it can use the complete + // hardware budget. When there are many input files, H workers scan files + // concurrently. With fewer files, the same H budget is divided among + // intra-file speculative scans. scan_parallel() uses its caller for one + // byte range, therefore the number of active scanners never exceeds H. + const auto scan_budget = hardware_worker_count(); + std::vector scan_threads(file_paths.size(), 0); + if (non_empty_files < static_cast(scan_budget)) { + size_t assigned_scanners = 0; + for (size_t i = 0; i < file_paths.size(); ++i) { + if (file_sizes[i] == 0) { + continue; + } + scan_threads[i] = std::max(1, range_counts[i]); + assigned_scanners += static_cast(scan_threads[i]); + } + CHECK_LE(assigned_scanners, static_cast(scan_budget)); + while (assigned_scanners < static_cast(scan_budget)) { + size_t best = file_paths.size(); + long double best_scanner_bytes = -1; + for (size_t i = 0; i < file_paths.size(); ++i) { + if (scan_threads[i] <= 0 || + static_cast(scan_threads[i]) >= file_sizes[i]) { + continue; + } + const auto scanner_bytes = + static_cast(file_sizes[i]) / scan_threads[i]; + if (scanner_bytes > best_scanner_bytes) { + best = i; + best_scanner_bytes = scanner_bytes; + } + } + if (best == file_paths.size()) { + break; + } + ++scan_threads[best]; + ++assigned_scanners; + } + } else { + for (size_t i = 0; i < file_paths.size(); ++i) { + scan_threads[i] = file_sizes[i] == 0 ? 0 : 1; + } + } + + std::vector scans(file_paths.size()); + std::atomic next_file{0}; + std::atomic scan_cancelled{false}; + std::mutex scan_error_mutex; + std::exception_ptr scan_error; + auto capture_scan_error = [&](std::exception_ptr error) { + bool expected = false; + if (scan_cancelled.compare_exchange_strong(expected, true, + std::memory_order_acq_rel)) { + std::lock_guard lock(scan_error_mutex); + scan_error = std::move(error); + } + }; + auto scan_file = [&](size_t file_index) { + if (file_sizes[file_index] == 0) { + return; + } + scans[file_index] = + CsvRowCountCounter(file_paths[file_index], config.quoting, + config.quote_char, config.double_quote, + config.delimiter, config.use_threads) + .ScanWithRanges(std::max(1, range_counts[file_index]), + std::max(1, scan_threads[file_index])); + }; + + std::vector planners; + if (non_empty_files >= static_cast(scan_budget)) { + planners.reserve(static_cast(scan_budget)); + for (int32_t worker = 0; worker < scan_budget; ++worker) { + planners.emplace_back([&] { + try { + while (!scan_cancelled.load(std::memory_order_acquire)) { + const auto file_index = + next_file.fetch_add(1, std::memory_order_relaxed); + if (file_index >= file_paths.size()) { + break; + } + scan_file(file_index); + } + } catch (...) { capture_scan_error(std::current_exception()); } + }); + } + } else { + planners.reserve(non_empty_files); + for (size_t i = 0; i < file_paths.size(); ++i) { + if (file_sizes[i] == 0) { + continue; + } + planners.emplace_back([&, i] { + try { + if (!scan_cancelled.load(std::memory_order_acquire)) { + scan_file(i); + } + } catch (...) { capture_scan_error(std::current_exception()); } + }); + } + } + for (auto& planner : planners) { + planner.join(); + } + if (scan_error) { + std::rethrow_exception(scan_error); + } + + int64_t total = 0; + for (size_t i = 0; i < file_paths.size(); ++i) { + const auto& scan = scans[i]; + if (scan.row_count < 0 || + scan.row_count > std::numeric_limits::max() - total) { + total = kUnknownRowNum; + } else if (total != kUnknownRowNum) { + total += scan.row_count; + } + + // SKIP is scoped per input file. Distribute it over record-aligned + // ranges so a short first range cannot leak rows that still need to be + // skipped into a later producer. + int64_t remaining_skip = std::max(0, config.skip_rows); + for (size_t range_index = 0; range_index < scan.ranges.size(); + ++range_index) { + const auto& range = scan.ranges[range_index]; + const auto range_end_row = range_index + 1 < scan.ranges.size() + ? scan.ranges[range_index + 1].start_row + : scan.row_count; + CHECK_GE(range_end_row, range.start_row); + const auto range_rows = range_end_row - range.start_row; + const auto range_skip = std::min(remaining_skip, range_rows); + remaining_skip -= range_skip; + plan->tasks.push_back({file_paths[i], range, range_skip}); + } + } + plan->row_count = total; + VLOG(1) << "CSV partition plan: files=" << file_paths.size() + << ", ranges=" << plan->tasks.size() + << ", producers=" << producer_count + << ", scan_budget=" << scan_budget << ", rows=" << total; + entry->plan = std::move(plan); + }); + return entry->plan; +} + struct CsvSupplierRuntime { - explicit CsvSupplierRuntime(const std::string& file_path, - const CsvReadConfig& config) + explicit CsvSupplierRuntime( + const std::string& file_path, const CsvReadConfig& config, + CsvRowCountMode row_count_mode, + std::optional range = std::nullopt) : file_path_(file_path), csv_format_(build_csv_format(config)), selected_column_names_(resolve_selected_column_names(config)), @@ -627,13 +1076,24 @@ struct CsvSupplierRuntime { rows_to_skip_(std::max(0, config.skip_rows)), chunk_size_(resolve_chunk_size(config)), escaping_(config.escaping), - escape_char_(config.escape_char) { + escape_char_(config.escape_char), + quoting_(config.quoting), + quote_char_(config.quote_char), + double_quote_(config.double_quote), + delimiter_(config.delimiter), + use_threads_(config.use_threads), + range_(std::move(range)) { if (selected_column_indices_.empty()) { THROW_SCHEMA_MISMATCH("No columns selected for CSV file: " + file_path_); } - row_num_ = CsvRowCountCounter(file_path, config.quoting, config.quote_char, - config.double_quote, config.delimiter) - .count(); + if (row_count_mode == CsvRowCountMode::kCountOnOpen) { + row_num_ = CsvRowCountCounter(file_path, quoting_, quote_char_, + double_quote_, delimiter_, use_threads_) + .count(); + } + if (range_) { + csv_format_.threading(false); + } reset_reader(); } @@ -696,14 +1156,27 @@ struct CsvSupplierRuntime { return chunk; } - int64_t row_num() const { return row_num_; } + int64_t row_num() const { + if (row_num_ == kUnknownRowNum) { + row_num_ = CsvRowCountCounter(file_path_, quoting_, quote_char_, + double_quote_, delimiter_, use_threads_) + .count(); + } + return row_num_; + } private: void reset_reader() { try { - reader_ = std::make_unique(file_path_, csv_format_); + if (range_) { + range_stream_ = std::make_unique( + file_path_, range_->start, range_->end); + reader_ = std::make_unique(*range_stream_, csv_format_); + } else { + reader_ = std::make_unique(file_path_, csv_format_); + } skip_rows(*reader_); - current_row_number_ = rows_to_skip_; + current_row_number_ = (range_ ? range_->start_row : 0) + rows_to_skip_; } catch (const std::exception& error) { THROW_IO_EXCEPTION("Failed to initialize CSV reader for file: " + file_path_ + ", reason=" + error.what()); @@ -732,8 +1205,15 @@ struct CsvSupplierRuntime { size_t chunk_size_ = kDefaultCsvChunkRows; bool escaping_ = false; char escape_char_ = '\\'; - int64_t row_num_ = 0; + bool quoting_ = false; + char quote_char_ = '"'; + bool double_quote_ = true; + char delimiter_ = ','; + bool use_threads_ = true; + std::optional range_; + mutable int64_t row_num_ = kUnknownRowNum; int64_t current_row_number_ = 0; + std::unique_ptr range_stream_; std::unique_ptr reader_; }; @@ -1003,6 +1483,12 @@ CsvReadConfig build_csv_read_config( config.double_quote = (value == "true" || value == "1" || value == "TRUE"); } + if (csv_options.count("PARALLEL")) { + const auto value = to_lower_copy(csv_options.at("PARALLEL")); + config.use_threads = + !(value == "0" || value == "false" || value == "off" || value == "no"); + } + bool header_row = true; if (csv_options.count("HEADER")) { auto val = to_lower_copy(csv_options.at("HEADER")); @@ -1037,10 +1523,11 @@ std::vector columnMappingsToSelectedCols( } CSVChunkSupplier::CSVChunkSupplier(const std::string& file_path, - CsvReadConfig config) + CsvReadConfig config, + CsvRowCountMode row_count_mode) : file_path_(file_path) { - runtime_ = std::make_unique(file_path, config); - row_num_ = runtime_->row_num(); + runtime_ = + std::make_unique(file_path, config, row_count_mode); VLOG(10) << "Finish init CSVChunkSupplier for file: " << file_path_; } @@ -1053,6 +1540,358 @@ std::shared_ptr CSVChunkSupplier::GetNextChunk() { return runtime_->get_next_chunk(); } +int64_t CSVChunkSupplier::RowNum() const { + if (!runtime_) { + return kUnknownRowNum; + } + return runtime_->row_num(); +} + +namespace { + +class ColumnProjectingChunkSupplier final : public IDataChunkSupplier { + public: + ColumnProjectingChunkSupplier(std::shared_ptr input, + std::vector columns) + : input_(std::move(input)), columns_(std::move(columns)) { + CHECK(input_ != nullptr); + } + + std::shared_ptr GetNextChunk() override { + auto input = input_->GetNextChunk(); + if (!input) { + return nullptr; + } + auto output = std::make_shared(); + for (size_t output_index = 0; output_index < columns_.size(); + ++output_index) { + const auto input_index = columns_[output_index]; + if (input_index < 0 || + static_cast(input_index) >= input->col_num()) { + THROW_INVALID_ARGUMENT_EXCEPTION( + "Chunk projection index is out of range: " + + std::to_string(input_index)); + } + output->set(static_cast(output_index), input->get(input_index)); + } + return output; + } + + int64_t RowNum() const override { return input_->RowNum(); } + + bool SupportsConcurrentGetNext() const override { + return input_->SupportsConcurrentGetNext(); + } + + void Cancel() override { input_->Cancel(); } + + private: + std::shared_ptr input_; + std::vector columns_; +}; + +CsvReadConfig project_csv_config(const CsvReadConfig& config, + const std::vector& columns) { + if (columns.empty()) { + return config; + } + CsvReadConfig projected = config; + const auto input_columns = resolve_selected_column_names(config); + projected.include_columns.clear(); + projected.include_columns.reserve(columns.size()); + for (const auto column : columns) { + if (column < 0 || static_cast(column) >= input_columns.size()) { + THROW_INVALID_ARGUMENT_EXCEPTION( + "CSV projection index is out of range: " + std::to_string(column)); + } + projected.include_columns.push_back( + input_columns[static_cast(column)]); + } + return projected; +} + +class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { + public: + PartitionedCsvChunkSupplier(std::vector file_paths, + CsvReadConfig config, + const ChunkSourceOptions& options, + std::shared_ptr plan_cache) + : file_paths_(std::move(file_paths)), + config_(std::move(config)), + producer_count_(std::clamp(options.producer_count, 1, + hardware_worker_count())), + plan_cache_(std::move(plan_cache)), + queue_(options.queue_capacity) {} + + ~PartitionedCsvChunkSupplier() override { + Cancel(); + JoinWorkers(); + } + + std::shared_ptr GetNextChunk() override { + EnsureScanned(); + StartWorkers(); + std::shared_ptr chunk; + if (queue_.Pop(chunk)) { + return chunk; + } + RethrowError(); + return nullptr; + } + + int64_t RowNum() const override { + const_cast(this)->EnsureScanned(); + return plan_->row_count; + } + + bool SupportsConcurrentGetNext() const override { return true; } + + void Cancel() override { + stop_.store(true, std::memory_order_release); + queue_.Close(); + } + + private: + void EnsureScanned() { + std::call_once(scan_once_, [&] { + CHECK(plan_cache_ != nullptr); + plan_ = plan_cache_->GetOrCreate(file_paths_, config_, producer_count_); + CHECK(plan_ != nullptr); + }); + } + + void StartWorkers() { + std::call_once(workers_once_, [&] { + if (plan_->tasks.empty()) { + queue_.Close(); + return; + } + const auto workers = std::min( + producer_count_, static_cast(plan_->tasks.size())); + active_workers_.store(workers, std::memory_order_relaxed); + workers_.reserve(static_cast(workers)); + for (int32_t worker = 0; worker < workers; ++worker) { + workers_.emplace_back([this] { WorkerMain(); }); + } + }); + } + + void WorkerMain() { + const bool profile_stages = VLOG_IS_ON(1); + try { + while (!stop_.load(std::memory_order_acquire)) { + const auto task_index = + next_task_.fetch_add(1, std::memory_order_relaxed); + if (task_index >= plan_->tasks.size()) { + break; + } + const auto& task = plan_->tasks[task_index]; + CsvReadConfig range_config = config_; + range_config.skip_rows = task.skip_rows; + range_config.use_threads = false; + const auto open_start = profile_stages + ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + CsvSupplierRuntime runtime(task.file_path, range_config, + CsvRowCountMode::kUnknown, task.range); + if (profile_stages) { + parse_time_ns_.fetch_add( + std::chrono::duration_cast( + std::chrono::steady_clock::now() - open_start) + .count(), + std::memory_order_relaxed); + } + while (!stop_.load(std::memory_order_acquire)) { + const auto parse_start = + profile_stages ? std::chrono::steady_clock::now() + : std::chrono::steady_clock::time_point{}; + auto chunk = runtime.get_next_chunk(); + if (profile_stages) { + parse_time_ns_.fetch_add( + std::chrono::duration_cast( + std::chrono::steady_clock::now() - parse_start) + .count(), + std::memory_order_relaxed); + } + if (!chunk || !queue_.Push(std::move(chunk))) { + break; + } + } + } + } catch (...) { SetError(std::current_exception()); } + if (active_workers_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + VLOG(1) << "Partitioned CSV producers: parse_worker_ms=" + << parse_time_ns_.load(std::memory_order_relaxed) / 1000000; + queue_.Close(); + } + } + + void SetError(std::exception_ptr error) { + bool expected = false; + if (has_error_.compare_exchange_strong(expected, true, + std::memory_order_acq_rel)) { + { + std::lock_guard lock(error_mutex_); + first_error_ = std::move(error); + } + Cancel(); + } + } + + void RethrowError() const { + if (!has_error_.load(std::memory_order_acquire)) { + return; + } + std::exception_ptr error; + { + std::lock_guard lock(error_mutex_); + error = first_error_; + } + if (error) { + std::rethrow_exception(error); + } + } + + void JoinWorkers() { + for (auto& worker : workers_) { + if (worker.joinable()) { + worker.join(); + } + } + } + + std::vector file_paths_; + CsvReadConfig config_; + int32_t producer_count_; + std::shared_ptr plan_cache_; + std::shared_ptr plan_; + chunk_pipeline_detail::BoundedQueue> queue_; + std::once_flag scan_once_; + std::once_flag workers_once_; + std::vector workers_; + std::atomic next_task_{0}; + std::atomic active_workers_{0}; + std::atomic stop_{false}; + std::atomic has_error_{false}; + std::atomic parse_time_ns_{0}; + mutable std::mutex error_mutex_; + std::exception_ptr first_error_; +}; + +class ChainedCsvChunkSupplier final : public IDataChunkSupplier { + public: + ChainedCsvChunkSupplier(std::vector file_paths, + CsvReadConfig config) + : file_paths_(std::move(file_paths)), config_(std::move(config)) {} + + std::shared_ptr GetNextChunk() override { + while (next_file_ < file_paths_.size()) { + if (!current_) { + current_ = std::make_shared( + file_paths_[next_file_], config_, CsvRowCountMode::kUnknown); + } + auto chunk = current_->GetNextChunk(); + if (chunk) { + return chunk; + } + current_.reset(); + ++next_file_; + } + return nullptr; + } + + int64_t RowNum() const override { + if (row_num_ != kUnknownRowNum) { + return row_num_; + } + int64_t total = 0; + for (const auto& file_path : file_paths_) { + auto supplier = std::make_shared( + file_path, config_, CsvRowCountMode::kUnknown); + auto rows = supplier->RowNum(); + if (rows < 0 || rows > std::numeric_limits::max() - total) { + return kUnknownRowNum; + } + total += rows; + } + row_num_ = total; + return row_num_; + } + + private: + std::vector file_paths_; + CsvReadConfig config_; + std::shared_ptr current_; + size_t next_file_ = 0; + mutable int64_t row_num_ = kUnknownRowNum; +}; + +} // namespace + +std::shared_ptr IDataChunkSource::Open( + const ChunkSourceOptions& options) const { + auto supplier = Open(); + if (!supplier || options.projected_columns.empty()) { + return supplier; + } + return std::make_shared( + std::move(supplier), options.projected_columns); +} + +CSVChunkSource::CSVChunkSource(std::vector file_paths, + CsvReadConfig config) + : file_paths_(std::move(file_paths)), + config_(std::move(config)), + partition_plan_cache_(std::make_shared()) {} + +std::shared_ptr CSVChunkSource::Open() const { + if (file_paths_.empty()) { + THROW_INVALID_ARGUMENT_EXCEPTION("CSV chunk source has no input paths"); + } + if (file_paths_.size() == 1) { + return std::make_shared(file_paths_.front(), config_, + CsvRowCountMode::kUnknown); + } + return std::make_shared(file_paths_, config_); +} + +std::shared_ptr CSVChunkSource::Open( + const ChunkSourceOptions& options) const { + if (file_paths_.empty()) { + THROW_INVALID_ARGUMENT_EXCEPTION("CSV chunk source has no input paths"); + } + auto open_config = project_csv_config(config_, options.projected_columns); + if (open_config.use_threads && options.parallel_enabled && + !options.preserve_order) { + return std::make_shared( + file_paths_, std::move(open_config), options, partition_plan_cache_); + } + + CsvReadConfig serial_config = std::move(open_config); + serial_config.use_threads = false; + if (file_paths_.size() == 1) { + return std::make_shared(file_paths_.front(), + std::move(serial_config), + CsvRowCountMode::kUnknown); + } + return std::make_shared(file_paths_, + std::move(serial_config)); +} + +int64_t CSVChunkSource::EstimatedBytes() const { + int64_t total = 0; + for (const auto& file_path : file_paths_) { + std::error_code error; + auto size = std::filesystem::file_size(file_path, error); + if (error || size > static_cast( + std::numeric_limits::max() - total)) { + return -1; + } + total += static_cast(size); + } + return total; +} + void fillVertexReaderMeta( label_t v_label, const std::string& v_label_name, const std::string& v_file, const LoadingConfig& loading_config, diff --git a/src/utils/io/read/common/options.cc b/src/utils/io/read/common/options.cc index 81b1575e0..4c6d95e3e 100644 --- a/src/utils/io/read/common/options.cc +++ b/src/utils/io/read/common/options.cc @@ -43,6 +43,7 @@ CsvReadConfig CsvOptionsBuilder::build() const { config.quote_char = parseOpts.quote_char.get(options); config.escaping = parseOpts.escaping.get(options); config.escape_char = parseOpts.escape_char.get(options); + config.use_threads = readOpts.use_threads.get(options); config.skip_rows = readOpts.skip_rows.get(options); int64_t batch_size = readOpts.batch_size.get(options); diff --git a/src/utils/io/read/csv/csv_reader.cc b/src/utils/io/read/csv/csv_reader.cc index f7f7fc23e..fcef40efa 100644 --- a/src/utils/io/read/csv/csv_reader.cc +++ b/src/utils/io/read/csv/csv_reader.cc @@ -485,6 +485,35 @@ CsvReader::CsvReader(std::shared_ptr sharedState, CsvReader::~CsvReader() = default; +std::shared_ptr CsvReader::createChunkSource() { + if (!sharedState_) { + THROW_INVALID_ARGUMENT_EXCEPTION("SharedState is null"); + } + if (!optionsBuilder_) { + THROW_INVALID_ARGUMENT_EXCEPTION("Options builder is null"); + } + // This source is a semantic subset of read(): it deliberately excludes + // post-read filtering and projection, so the bulk path can reopen the same + // parser without changing the normal reader's behavior. Callers fall back + // to read() for every other source shape. + if (sharedState_->skipRows || !sharedState_->projectColumns.empty()) { + return nullptr; + } + + auto config = optionsBuilder_->build(); + if (!optionsBuilder_->projectColumns(config)) { + LOG(WARNING) << "Failed to set column projection, using all columns"; + } + + auto read_config = read_config_for_supplier(config); + + const auto& paths = sharedState_->schema.file.paths; + if (paths.empty()) { + THROW_INVALID_ARGUMENT_EXCEPTION("No file paths provided"); + } + return std::make_shared(paths, std::move(read_config)); +} + void CsvReader::read(std::shared_ptr /*localState*/, execution::Context& ctx) { if (!sharedState_) { diff --git a/tests/storage/test_copy_temp.cc b/tests/storage/test_copy_temp.cc index 4b0ffae81..f643e7274 100644 --- a/tests/storage/test_copy_temp.cc +++ b/tests/storage/test_copy_temp.cc @@ -14,8 +14,10 @@ */ #include +#include #include #include +#include #include "neug/main/connection.h" #include "neug/main/neug_db.h" @@ -24,6 +26,30 @@ namespace neug { namespace test { +class ScopedEnvironmentVariable final { + public: + ScopedEnvironmentVariable(const char* key, const char* value) : key_(key) { + if (const char* previous = std::getenv(key); previous != nullptr) { + previous_ = previous; + } + if (::setenv(key, value, 1) != 0) { + throw std::runtime_error("Failed to set test environment variable"); + } + } + + ~ScopedEnvironmentVariable() { + if (previous_) { + ::setenv(key_.c_str(), previous_->c_str(), 1); + } else { + ::unsetenv(key_.c_str()); + } + } + + private: + std::string key_; + std::optional previous_; +}; + class CopyTempTest : public ::testing::Test { protected: static constexpr const char* DB_DIR = "/tmp/copy_temp_test_db"; @@ -134,6 +160,63 @@ TEST_F(CopyTempTest, NodeDefaultPrimaryKey) { conn->Close(); } +TEST_F(CopyTempTest, PersistentCopyUsesForcedBulkPath) { + auto conn = db_->Connect(); + const std::string people = std::string(CSV_DIR) + "/people.csv"; + const std::string edges = std::string(CSV_DIR) + "/edges.csv"; + const std::string dangling = std::string(CSV_DIR) + "/dangling_edges.csv"; + + ASSERT_TRUE( + conn->Query("CREATE NODE TABLE Person(id INT64, name STRING, age INT64, " + "PRIMARY KEY(id));")); + ASSERT_TRUE(conn->Query( + "CREATE REL TABLE Knows(FROM Person TO Person, weight DOUBLE);")); + ASSERT_TRUE(conn->Query( + "CREATE REL TABLE Dangling(FROM Person TO Person, weight DOUBLE);")); + + { + ScopedEnvironmentVariable force_bulk("NEUG_COPY_BULK_BUILD", "true"); + auto vertices = + conn->Query("COPY Person FROM \"" + people + "\" (header = true);"); + ASSERT_TRUE(vertices) << vertices.error().ToString(); + auto relationships = + conn->Query("COPY Knows FROM \"" + edges + "\" (header = true);"); + ASSERT_TRUE(relationships) << relationships.error().ToString(); + auto dangling_relationships = + conn->Query("COPY Dangling FROM \"" + dangling + "\" (header = true);"); + ASSERT_TRUE(dangling_relationships) + << dangling_relationships.error().ToString(); + } + + auto vertex_count = conn->Query("MATCH (n:Person) RETURN n.id;"); + ASSERT_TRUE(vertex_count) << vertex_count.error().ToString(); + EXPECT_EQ(vertex_count.value().response().row_count(), 4); + auto edge_count = + conn->Query("MATCH (:Person)-[e:Knows]->(:Person) RETURN e.weight;"); + ASSERT_TRUE(edge_count) << edge_count.error().ToString(); + EXPECT_EQ(edge_count.value().response().row_count(), 3); + auto dangling_count = + conn->Query("MATCH (:Person)-[e:Dangling]->(:Person) RETURN e.weight;"); + ASSERT_TRUE(dangling_count) << dangling_count.error().ToString(); + EXPECT_EQ(dangling_count.value().response().row_count(), 1); + + // The same terminal plan must retain the established materialized path + // when bulk build is explicitly disabled. + ASSERT_TRUE( + conn->Query("CREATE NODE TABLE Fallback(id INT64, name STRING, " + "age INT64, PRIMARY KEY(id));")); + { + ScopedEnvironmentVariable disable_bulk("NEUG_COPY_BULK_BUILD", "false"); + auto fallback = + conn->Query("COPY Fallback FROM \"" + people + "\" (header = true);"); + ASSERT_TRUE(fallback) << fallback.error().ToString(); + } + auto fallback_count = conn->Query("MATCH (n:Fallback) RETURN n.id;"); + ASSERT_TRUE(fallback_count) << fallback_count.error().ToString(); + EXPECT_EQ(fallback_count.value().response().row_count(), 4); + conn->Close(); +} + TEST_F(CopyTempTest, NodeWithWhere) { auto conn = db_->Connect(); std::string csv = std::string(CSV_DIR) + "/people.csv"; diff --git a/tests/storage/test_edge_table.cc b/tests/storage/test_edge_table.cc index 9798136f0..db838ff70 100644 --- a/tests/storage/test_edge_table.cc +++ b/tests/storage/test_edge_table.cc @@ -13,8 +13,16 @@ * limitations under the License. */ #include +#include +#include +#include +#include #include +#include +#include #include +#include +#include #include "neug/common/types/value.h" #include "neug/execution/execute/ops/batch/batch_update_utils.h" @@ -68,12 +76,29 @@ class EdgeTableTest : public ::testing::Test { neug::EdgeStrategy::kMultiple, true, true, std::nullopt, "person creates comment edge with two properties"); + schema_.AddEdgeLabel( + "person", "comment", "create_single", {neug::DataTypeId::kInt32}, + {"data"}, neug::EdgeStrategy::kSingle, neug::EdgeStrategy::kMultiple, + true, true, std::nullopt, + "single outgoing edge used to verify ordered bulk loading"); + schema_.AddEdgeLabel( + "person", "comment", "create_single_both", {neug::DataTypeId::kInt32}, + {"data"}, neug::EdgeStrategy::kSingle, neug::EdgeStrategy::kSingle, + true, true, std::nullopt, + "single edge in both directions used to verify one-pass bulk loading"); + schema_.AddEdgeLabel( + "person", "comment", "create_none", {neug::DataTypeId::kInt32}, + {"data"}, neug::EdgeStrategy::kNone, neug::EdgeStrategy::kNone, true, + true, std::nullopt, "edge label with no stored adjacency"); src_label_ = schema_.get_vertex_label_id("person"); dst_label_ = schema_.get_vertex_label_id("comment"); edge_label_empty_ = schema_.get_edge_label_id("create0"); edge_label_int_ = schema_.get_edge_label_id("create1"); edge_label_str_ = schema_.get_edge_label_id("create2"); edge_label_str_int_ = schema_.get_edge_label_id("create3"); + edge_label_single_ = schema_.get_edge_label_id("create_single"); + edge_label_single_both_ = schema_.get_edge_label_id("create_single_both"); + edge_label_none_ = schema_.get_edge_label_id("create_none"); allocator_dir_ = "/tmp/edge_table_test_allocator_" + std::to_string(::getpid()) + "_"; ws.Open(temp_dir_.string()); @@ -130,6 +155,13 @@ class EdgeTableTest : public ::testing::Test { edge_table->BatchAddEdges(src_indexer, dst_indexer, supplier); } + void BatchBuild(std::vector> chunks, + int64_t estimated_bytes = -1) { + auto source = std::make_shared(std::move(chunks), + estimated_bytes); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, std::move(source)); + } + size_t ExpectedBatchInsertCapacity(size_t inserted_edge_num) const { if (inserted_edge_num == 0) { return 0; @@ -237,9 +269,12 @@ class EdgeTableTest : public ::testing::Test { neug::LFIndexer dst_indexer; neug::Schema schema_; neug::label_t src_label_, dst_label_, edge_label_empty_, edge_label_int_, - edge_label_str_, edge_label_str_int_; + edge_label_str_, edge_label_str_int_, edge_label_single_, + edge_label_single_both_, edge_label_none_; std::string allocator_dir_; + const std::filesystem::path& temp_dir() const { return temp_dir_; } + private: std::filesystem::path temp_dir_; neug::CheckpointManager ws; @@ -695,6 +730,1195 @@ TEST_F(EdgeTableTest, TestBatchAddEdgesBundled) { ASSERT_EQ(dsts.size(), edge_num + more_edge_num); } +TEST_F(EdgeTableTest, BatchBuildEdgesBundledWithParallelEncoding) { + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kSrcNum = 100; + constexpr int64_t kDstNum = 100; + constexpr size_t kEdgeNum = 4000; + + auto src_list = generate_random_vertices(kSrcNum, kEdgeNum); + auto dst_list = generate_random_vertices(kDstNum, kEdgeNum); + auto data_list = generate_random_data(kEdgeNum); + auto batches = + convert_to_data_chunks({split_column_to_chunks(src_list, 16), + split_column_to_chunks(dst_list, 16), + split_column_to_chunks(data_list, 16)}); + + InitIndexers(*ckp, kSrcNum, kDstNum); + ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); + OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), kSrcNum, kDstNum); + EXPECT_TRUE(edge_table->CanBatchBuild()); + auto source = std::make_shared(std::move(batches), + 256LL * 1024 * 1024); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + + EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); + EXPECT_EQ(source->OpenCount(), 2); + ASSERT_EQ(source->OpenedProjections().size(), 2); + EXPECT_EQ(source->OpenedProjections()[0], (std::vector{0, 1})); + EXPECT_TRUE(source->OpenedProjections()[1].empty()); + EXPECT_FALSE(edge_table->CanBatchBuild()); + std::vector output_srcs, output_dsts; + OutputOutgoingEndpoints(output_srcs, output_dsts, neug::MAX_TIMESTAMP); + ASSERT_EQ(output_srcs.size(), kEdgeNum); + ASSERT_EQ(output_dsts.size(), kEdgeNum); + + std::vector> expected; + std::vector> actual; + expected.reserve(kEdgeNum); + actual.reserve(kEdgeNum); + for (size_t i = 0; i < kEdgeNum; ++i) { + expected.emplace_back(src_list[i], dst_list[i], data_list[i]); + } + std::vector output_data; + OutputOutgoingEdgeData(output_data, neug::MAX_TIMESTAMP, 0); + ASSERT_EQ(output_data.size(), kEdgeNum); + for (size_t i = 0; i < kEdgeNum; ++i) { + actual.emplace_back(output_srcs[i], output_dsts[i], output_data[i]); + } + std::sort(expected.begin(), expected.end()); + std::sort(actual.begin(), actual.end()); + EXPECT_EQ(actual, expected); +} + +TEST_F(EdgeTableTest, BatchBuildEdgesParallelFillWithoutProperties) { + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kVertexNum = 32; + constexpr size_t kEdgeNum = 2048; + std::vector srcs(kEdgeNum); + std::vector dsts(kEdgeNum); + for (size_t row = 0; row < kEdgeNum; ++row) { + srcs[row] = static_cast(row % kVertexNum); + dsts[row] = static_cast((row * 7) % kVertexNum); + } + auto chunks = convert_to_data_chunks( + {split_column_to_chunks(srcs, 32), split_column_to_chunks(dsts, 32)}); + + InitIndexers(*ckp, kVertexNum, kVertexNum); + ConstructEdgeTable(src_label_, dst_label_, edge_label_empty_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); + auto source = std::make_shared(std::move(chunks), + 1024LL * 1024 * 1024); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + + EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); + std::vector actual_srcs, actual_dsts; + OutputOutgoingEndpoints(actual_srcs, actual_dsts, MAX_TIMESTAMP); + ASSERT_EQ(actual_srcs.size(), kEdgeNum); + ASSERT_EQ(actual_dsts.size(), kEdgeNum); + std::vector> expected; + std::vector> actual; + expected.reserve(kEdgeNum); + actual.reserve(kEdgeNum); + for (size_t row = 0; row < kEdgeNum; ++row) { + expected.emplace_back(srcs[row], dsts[row]); + actual.emplace_back(actual_srcs[row], actual_dsts[row]); + } + std::sort(expected.begin(), expected.end()); + std::sort(actual.begin(), actual.end()); + EXPECT_EQ(actual, expected); +} + +TEST_F(EdgeTableTest, BatchBuildEdgesHandlesEmptyInput) { + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kVertexNum = 16; + InitIndexers(*ckp, kVertexNum, kVertexNum); + ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); + auto source = std::make_shared( + std::vector>{}, 1024LL * 1024 * 1024); + + EXPECT_TRUE(edge_table->CanBatchBuild()); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + + EXPECT_EQ(source->OpenCount(), 2); + EXPECT_EQ(edge_table->EdgeNum(), 0); + // CanBatchBuild() is edge-count based, so a successfully published empty + // graph intentionally remains eligible for another bulk build. + EXPECT_TRUE(edge_table->CanBatchBuild()); + std::vector output_srcs, output_dsts; + OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); + EXPECT_TRUE(output_srcs.empty()); + EXPECT_TRUE(output_dsts.empty()); +} + +TEST_F(EdgeTableTest, BatchBuildEdgesSkipsSourceForNoAdjacency) { + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kVertexNum = 4; + auto chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0, 1, 2}, 2), + split_column_to_chunks(std::vector{1, 2, 3}, 2), + split_column_to_chunks(std::vector{10, 20, 30}, 2)}); + + InitIndexers(*ckp, kVertexNum, kVertexNum); + ConstructEdgeTable(src_label_, dst_label_, edge_label_none_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); + auto source = std::make_shared(std::move(chunks), + 1024LL * 1024 * 1024); + + ASSERT_TRUE(edge_table->CanBatchBuild()); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + + EXPECT_EQ(source->OpenCount(), 0); + EXPECT_TRUE(source->OpenedProjections().empty()); + EXPECT_EQ(edge_table->EdgeNum(), 0); + EXPECT_TRUE(edge_table->CanBatchBuild()); +} + +TEST_F(EdgeTableTest, BatchBuildEdgesParallelFillHandlesSupernodes) { + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kVertexNum = 4; + constexpr size_t kEdgeNum = 8192; + std::vector srcs(kEdgeNum, 0); + std::vector dsts(kEdgeNum, 1); + std::vector data(kEdgeNum); + for (size_t row = 0; row < kEdgeNum; ++row) { + data[row] = static_cast(row); + } + auto chunks = convert_to_data_chunks({split_column_to_chunks(srcs, 64), + split_column_to_chunks(dsts, 64), + split_column_to_chunks(data, 64)}); + + InitIndexers(*ckp, kVertexNum, kVertexNum); + ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); + auto source = std::make_shared(std::move(chunks), + 1024LL * 1024 * 1024); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + + EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); + std::vector output_srcs, output_dsts; + std::vector output_data; + OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); + OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); + ASSERT_EQ(output_srcs.size(), kEdgeNum); + ASSERT_EQ(output_dsts.size(), kEdgeNum); + ASSERT_EQ(output_data.size(), kEdgeNum); + EXPECT_TRUE(std::all_of(output_srcs.begin(), output_srcs.end(), + [](int64_t src) { return src == 0; })); + EXPECT_TRUE(std::all_of(output_dsts.begin(), output_dsts.end(), + [](int64_t dst) { return dst == 1; })); + std::sort(output_data.begin(), output_data.end()); + EXPECT_EQ(output_data, data); + + std::vector incoming_srcs, incoming_dsts; + OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); + ASSERT_EQ(incoming_srcs.size(), kEdgeNum); + ASSERT_EQ(incoming_dsts.size(), kEdgeNum); + EXPECT_TRUE(std::all_of(incoming_srcs.begin(), incoming_srcs.end(), + [](int64_t src) { return src == 0; })); + EXPECT_TRUE(std::all_of(incoming_dsts.begin(), incoming_dsts.end(), + [](int64_t dst) { return dst == 1; })); +} + +TEST_F(EdgeTableTest, SingleEdgeBulkBuildFillsUniqueSlotsAcrossChunks) { + if (std::thread::hardware_concurrency() < 4) { + GTEST_SKIP() << "Concurrent bulk fill needs at least two consumer workers"; + } + + struct SupplierActivity { + std::atomic active{0}; + std::atomic max_active{0}; + }; + + class ConcurrentChunkSupplier final : public IDataChunkSupplier { + public: + ConcurrentChunkSupplier(std::vector> chunks, + std::shared_ptr activity) + : chunks_(std::move(chunks)), activity_(std::move(activity)) {} + + std::shared_ptr GetNextChunk() override { + const auto active = + activity_->active.fetch_add(1, std::memory_order_relaxed) + 1; + auto max_active = activity_->max_active.load(std::memory_order_relaxed); + while (max_active < active && + !activity_->max_active.compare_exchange_weak( + max_active, active, std::memory_order_relaxed)) {} + std::this_thread::sleep_for(std::chrono::microseconds(100)); + const auto index = next_.fetch_add(1, std::memory_order_relaxed); + activity_->active.fetch_sub(1, std::memory_order_relaxed); + return index < chunks_.size() ? chunks_[index] : nullptr; + } + + int64_t RowNum() const override { + int64_t rows = 0; + for (const auto& chunk : chunks_) { + rows += static_cast(chunk->row_num()); + } + return rows; + } + + bool SupportsConcurrentGetNext() const override { return true; } + + private: + std::vector> chunks_; + std::shared_ptr activity_; + std::atomic next_{0}; + }; + + class ConcurrentChunkSource final : public IDataChunkSource { + public: + explicit ConcurrentChunkSource( + std::vector> chunks) + : chunks_(std::move(chunks)) {} + + std::shared_ptr Open() const override { + auto activity = std::make_shared(); + activities_.push_back(activity); + return std::make_shared(chunks_, activity); + } + + std::shared_ptr Open( + const ChunkSourceOptions& options) const override { + opened_options_.push_back(options); + return IDataChunkSource::Open(options); + } + + bool rewindable() const override { return true; } + int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } + + const std::vector& OpenedOptions() const { + return opened_options_; + } + + const std::vector>& Activities() const { + return activities_; + } + + private: + std::vector> chunks_; + mutable std::vector opened_options_; + mutable std::vector> activities_; + }; + + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kSrcNum = 512; + constexpr int64_t kDstNum = 8; + std::vector srcs(kSrcNum); + std::vector dsts(kSrcNum); + std::vector data(kSrcNum); + for (int64_t row = 0; row < kSrcNum; ++row) { + srcs[row] = row; + dsts[row] = row % kDstNum; + data[row] = static_cast(row * 3); + } + auto chunks = convert_to_data_chunks({split_column_to_chunks(srcs, 128), + split_column_to_chunks(dsts, 128), + split_column_to_chunks(data, 128)}); + ASSERT_EQ(chunks.size(), 128); + + InitIndexers(*ckp, kSrcNum, kDstNum); + ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), kSrcNum, kDstNum); + auto source = std::make_shared(std::move(chunks)); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + + ASSERT_EQ(source->OpenedOptions().size(), 2); + EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); + EXPECT_FALSE(source->OpenedOptions()[1].preserve_order); + EXPECT_EQ(source->OpenedOptions()[0].projected_columns, + (std::vector{0, 1})); + EXPECT_TRUE(source->OpenedOptions()[1].projected_columns.empty()); + ASSERT_EQ(source->Activities().size(), 2); + EXPECT_GT(source->Activities()[0]->max_active.load(), 1); + EXPECT_GT(source->Activities()[1]->max_active.load(), 1); + EXPECT_EQ(edge_table->EdgeNum(), kSrcNum); + + std::vector output_srcs, output_dsts; + std::vector output_data; + OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); + OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); + EXPECT_EQ(output_srcs, srcs); + EXPECT_EQ(output_dsts, dsts); + EXPECT_EQ(output_data, data); + + std::vector incoming_srcs, incoming_dsts; + std::vector incoming_data; + OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); + OutputIncomingEdgeData(incoming_data, MAX_TIMESTAMP, 0); + ASSERT_EQ(incoming_srcs.size(), srcs.size()); + ASSERT_EQ(incoming_dsts.size(), dsts.size()); + ASSERT_EQ(incoming_data.size(), data.size()); + std::vector> expected; + std::vector> actual; + expected.reserve(srcs.size()); + actual.reserve(srcs.size()); + for (size_t row = 0; row < srcs.size(); ++row) { + expected.emplace_back(srcs[row], dsts[row], data[row]); + actual.emplace_back(incoming_srcs[row], incoming_dsts[row], + incoming_data[row]); + } + std::sort(expected.begin(), expected.end()); + std::sort(actual.begin(), actual.end()); + EXPECT_EQ(actual, expected); +} + +TEST_F(EdgeTableTest, SingleEdgeBulkBuildSkipsInvalidEndpoints) { + auto ckp = make_checkpoint(workspace()); + InitIndexers(*ckp, 3, 3); + ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), 3, 3); + + const std::vector srcs = {0, 99, 1, 2}; + const std::vector dsts = {0, 1, 99, 2}; + const std::vector data = {10, 20, 30, 40}; + auto chunks = convert_to_data_chunks({split_column_to_chunks(srcs, 4), + split_column_to_chunks(dsts, 4), + split_column_to_chunks(data, 4)}); + auto source = std::make_shared(std::move(chunks), + 1024LL * 1024 * 1024); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + + EXPECT_EQ(source->OpenCount(), 2); + EXPECT_EQ(edge_table->EdgeNum(), 2); + std::vector output_srcs, output_dsts; + std::vector output_data; + OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); + OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); + EXPECT_EQ(output_srcs, (std::vector{0, 2})); + EXPECT_EQ(output_dsts, (std::vector{0, 2})); + EXPECT_EQ(output_data, (std::vector{10, 40})); + + std::vector incoming_srcs, incoming_dsts; + std::vector incoming_data; + OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); + OutputIncomingEdgeData(incoming_data, MAX_TIMESTAMP, 0); + EXPECT_EQ(incoming_srcs, (std::vector{0, 2})); + EXPECT_EQ(incoming_dsts, (std::vector{0, 2})); + EXPECT_EQ(incoming_data, (std::vector{10, 40})); +} + +TEST_F(EdgeTableTest, SingleEdgeBulkBuildDetectsCrossWorkerDuplicate) { + if (std::thread::hardware_concurrency() < 4) { + GTEST_SKIP() + << "Cross-worker duplicate detection needs two consumer workers"; + } + + class OrderedChunkSupplier final : public IDataChunkSupplier { + public: + explicit OrderedChunkSupplier( + std::vector> chunks) + : chunks_(std::move(chunks)) {} + + std::shared_ptr GetNextChunk() override { + if (next_ == chunks_.size()) { + return nullptr; + } + return chunks_[next_++]; + } + + int64_t RowNum() const override { + int64_t rows = 0; + for (const auto& chunk : chunks_) { + rows += static_cast(chunk->row_num()); + } + return rows; + } + + private: + std::vector> chunks_; + size_t next_ = 0; + }; + + class BarrierChunkSupplier final : public IDataChunkSupplier { + public: + explicit BarrierChunkSupplier( + std::vector> chunks) + : chunks_(std::move(chunks)) {} + + std::shared_ptr GetNextChunk() override { + const auto index = next_.fetch_add(1, std::memory_order_relaxed); + if (index >= chunks_.size()) { + return nullptr; + } + if (index < 2) { + std::unique_lock lock(mutex_); + ++arrived_; + cv_.notify_all(); + cv_.wait(lock, [&] { return arrived_ >= 2 || cancelled_; }); + if (cancelled_) { + return nullptr; + } + } + return chunks_[index]; + } + + int64_t RowNum() const override { + int64_t rows = 0; + for (const auto& chunk : chunks_) { + rows += static_cast(chunk->row_num()); + } + return rows; + } + + bool SupportsConcurrentGetNext() const override { return true; } + + void Cancel() override { + { + std::lock_guard lock(mutex_); + cancelled_ = true; + } + cv_.notify_all(); + } + + private: + std::vector> chunks_; + std::atomic next_{0}; + std::mutex mutex_; + std::condition_variable cv_; + size_t arrived_ = 0; + bool cancelled_ = false; + }; + + class CrossWorkerDuplicateSource final : public IDataChunkSource { + public: + explicit CrossWorkerDuplicateSource( + std::vector> chunks) + : chunks_(std::move(chunks)) {} + + std::shared_ptr Open() const override { + return std::make_shared(chunks_); + } + + std::shared_ptr Open( + const ChunkSourceOptions& options) const override { + opened_options_.push_back(options); + if (options.preserve_order) { + return std::make_shared(chunks_); + } + return std::make_shared(chunks_); + } + + bool rewindable() const override { return true; } + int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } + + const std::vector& OpenedOptions() const { + return opened_options_; + } + + private: + std::vector> chunks_; + mutable std::vector opened_options_; + }; + + auto ckp = make_checkpoint(workspace()); + InitIndexers(*ckp, 1, 2); + ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), 1, 2); + + auto chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0, 0}, 2), + split_column_to_chunks(std::vector{0, 1}, 2), + split_column_to_chunks(std::vector{10, 20}, 2)}); + ASSERT_EQ(chunks.size(), 2); + auto source = std::make_shared(std::move(chunks)); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + + ASSERT_EQ(source->OpenedOptions().size(), 2); + EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); + EXPECT_TRUE(source->OpenedOptions()[1].preserve_order); + EXPECT_EQ(edge_table->EdgeNum(), 2); + std::vector output_srcs, output_dsts; + std::vector output_data; + OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); + OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); + EXPECT_EQ(output_srcs, (std::vector{0})); + EXPECT_EQ(output_dsts, (std::vector{1})); + EXPECT_EQ(output_data, (std::vector{20})); +} + +TEST_F(EdgeTableTest, SingleEdgeBulkBuildPreservesLastWriteOrderAcrossChunks) { + class OrderedChunkSupplier final : public IDataChunkSupplier { + public: + explicit OrderedChunkSupplier( + std::vector> chunks) + : chunks_(std::move(chunks)) {} + + std::shared_ptr GetNextChunk() override { + if (next_chunk_ == chunks_.size()) { + return nullptr; + } + return chunks_[next_chunk_++]; + } + + int64_t RowNum() const override { + int64_t rows = 0; + for (const auto& chunk : chunks_) { + rows += static_cast(chunk->row_num()); + } + return rows; + } + + private: + std::vector> chunks_; + size_t next_chunk_ = 0; + }; + + class OrderedChunkSource final : public IDataChunkSource { + public: + explicit OrderedChunkSource(std::vector> chunks) + : chunks_(std::move(chunks)) {} + + std::shared_ptr Open() const override { + ++open_count_; + return std::make_shared(chunks_); + } + + std::shared_ptr Open( + const ChunkSourceOptions& options) const override { + opened_projections_.push_back(options.projected_columns); + opened_preserve_order_.push_back(options.preserve_order); + return IDataChunkSource::Open(options); + } + + bool rewindable() const override { return true; } + int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } + size_t OpenCount() const { return open_count_; } + const std::vector>& OpenedProjections() const { + return opened_projections_; + } + const std::vector& OpenedPreserveOrder() const { + return opened_preserve_order_; + } + + private: + std::vector> chunks_; + mutable size_t open_count_ = 0; + mutable std::vector> opened_projections_; + mutable std::vector opened_preserve_order_; + }; + + auto ckp = make_checkpoint(workspace()); + InitIndexers(*ckp, 1, 3); + ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); + OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), 1, 3); + + const std::vector srcs = {0, 0, 0}; + const std::vector dsts = {0, 1, 2}; + const std::vector data = {10, 20, 30}; + auto chunks = convert_to_data_chunks({split_column_to_chunks(srcs, 3), + split_column_to_chunks(dsts, 3), + split_column_to_chunks(data, 3)}); + ASSERT_EQ(chunks.size(), 3); + auto source = std::make_shared(std::move(chunks)); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + + std::vector output_srcs, output_dsts; + std::vector output_data; + OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); + OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); + ASSERT_EQ(output_srcs, (std::vector{0})); + ASSERT_EQ(output_dsts, (std::vector{2})); + ASSERT_EQ(output_data, (std::vector{30})); + EXPECT_EQ(source->OpenCount(), 2); + ASSERT_EQ(source->OpenedProjections().size(), 2); + EXPECT_EQ(source->OpenedProjections()[0], (std::vector{0, 1})); + EXPECT_TRUE(source->OpenedProjections()[1].empty()); + EXPECT_EQ(source->OpenedPreserveOrder(), (std::vector{false, true})); + EXPECT_EQ(edge_table->EdgeNum(), 3); +} + +TEST_F(EdgeTableTest, SingleOnlyBulkBuildPreservesBothDirections) { + auto ckp = make_checkpoint(workspace()); + InitIndexers(*ckp, 2, 2); + ConstructEdgeTable(src_label_, dst_label_, edge_label_single_both_); + OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), 2, 2); + + auto chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0, 1, 0}, 1), + split_column_to_chunks(std::vector{0, 0, 1}, 1), + split_column_to_chunks(std::vector{10, 20, 30}, 1)}); + auto source = std::make_shared(std::move(chunks), + 1024LL * 1024 * 1024); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + + EXPECT_EQ(source->OpenCount(), 1); + ASSERT_EQ(source->OpenedProjections().size(), 1); + EXPECT_TRUE(source->OpenedProjections()[0].empty()); + EXPECT_EQ(edge_table->EdgeNum(), 3); + + std::vector output_srcs, output_dsts; + std::vector output_data; + OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); + OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); + ASSERT_EQ(output_srcs, (std::vector{0, 1})); + ASSERT_EQ(output_dsts, (std::vector{1, 0})); + ASSERT_EQ(output_data, (std::vector{30, 20})); + + std::vector incoming_srcs, incoming_dsts; + std::vector incoming_data; + OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); + OutputIncomingEdgeData(incoming_data, MAX_TIMESTAMP, 0); + ASSERT_EQ(incoming_srcs, (std::vector{1, 0})); + ASSERT_EQ(incoming_dsts, (std::vector{0, 1})); + ASSERT_EQ(incoming_data, (std::vector{20, 30})); +} + +TEST_F(EdgeTableTest, BatchBuildEdgesFromPartitionedCsvInTwoPasses) { + class EstimatedBytesSource final : public IDataChunkSource { + public: + EstimatedBytesSource(std::shared_ptr source, + int64_t estimated_bytes) + : source_(std::move(source)), estimated_bytes_(estimated_bytes) {} + + std::shared_ptr Open() const override { + ++open_count_; + return source_->Open(); + } + + std::shared_ptr Open( + const ChunkSourceOptions& options) const override { + ++open_count_; + return source_->Open(options); + } + + bool rewindable() const override { return source_->rewindable(); } + int64_t EstimatedBytes() const override { return estimated_bytes_; } + bool ParallelEnabled() const override { return source_->ParallelEnabled(); } + size_t OpenCount() const { return open_count_; } + + private: + std::shared_ptr source_; + int64_t estimated_bytes_; + mutable size_t open_count_ = 0; + }; + + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kSrcNum = 100; + constexpr int64_t kDstNum = 80; + constexpr size_t kEdgeNum = 4000; + std::vector> expected; + expected.reserve(kEdgeNum); + + const auto csv_path = temp_dir() / "parallel-edges.csv"; + { + std::ofstream output(csv_path, std::ios::binary); + ASSERT_TRUE(output.is_open()); + output << "src|dst|data\n"; + for (size_t row = 0; row < kEdgeNum; ++row) { + const auto src = static_cast((row * 17) % kSrcNum); + const auto dst = static_cast((row * 29) % kDstNum); + const auto data = static_cast(row * 3 + 1); + output << src << '|' << dst << '|' << data << '\n'; + expected.emplace_back(src, dst, data); + } + // Both passes must filter the same invalid endpoint rows, otherwise the + // preallocated degree and the serial fill would disagree. + output << "1000|1|7\n1|1000|9\n"; + } + + CsvReadConfig config; + config.delimiter = '|'; + config.skip_rows = 1; + config.chunk_size = 7; + config.column_names = {"src", "dst", "data"}; + config.include_columns = config.column_names; + config.column_types.emplace("src", DataType(DataTypeId::kInt64)); + config.column_types.emplace("dst", DataType(DataTypeId::kInt64)); + config.column_types.emplace("data", DataType(DataTypeId::kInt32)); + + InitIndexers(*ckp, kSrcNum, kDstNum); + ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), kSrcNum, kDstNum); + auto csv_source = std::make_shared( + std::vector{csv_path.string()}, std::move(config)); + // Force the large-file planner while keeping the fixture small. The wrapper + // forwards both Open() calls to the same CSV source, so its cached partition + // plan is reused by the count and fill passes. + auto source = std::make_shared(std::move(csv_source), + 1024LL * 1024 * 1024); + edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + + EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); + EXPECT_EQ(source->OpenCount(), 2); + std::vector output_srcs, output_dsts; + std::vector output_data; + OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); + OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); + ASSERT_EQ(output_srcs.size(), kEdgeNum); + ASSERT_EQ(output_dsts.size(), kEdgeNum); + ASSERT_EQ(output_data.size(), kEdgeNum); + std::vector> actual; + actual.reserve(kEdgeNum); + for (size_t row = 0; row < kEdgeNum; ++row) { + actual.emplace_back(output_srcs[row], output_dsts[row], output_data[row]); + } + std::sort(expected.begin(), expected.end()); + std::sort(actual.begin(), actual.end()); + EXPECT_EQ(actual, expected); + + std::vector incoming_srcs, incoming_dsts; + OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); + EXPECT_EQ(incoming_srcs.size(), kEdgeNum); + EXPECT_EQ(incoming_dsts.size(), kEdgeNum); +} + +TEST_F(EdgeTableTest, FailedBatchBuildDoesNotPublishStagedCsr) { + class ThrowingSupplier final : public IDataChunkSupplier { + public: + explicit ThrowingSupplier(std::shared_ptr chunk) + : chunk_(std::move(chunk)) {} + + std::shared_ptr GetNextChunk() override { + if (chunk_) { + return std::exchange(chunk_, nullptr); + } + throw std::runtime_error("injected bulk edge source failure"); + } + + int64_t RowNum() const override { return 1; } + + private: + std::shared_ptr chunk_; + }; + + class ThrowingSource final : public IDataChunkSource { + public: + explicit ThrowingSource(std::shared_ptr chunk) + : chunk_(std::move(chunk)) {} + + std::shared_ptr Open() const override { + return std::make_shared(chunk_); + } + + bool rewindable() const override { return true; } + + private: + std::shared_ptr chunk_; + }; + + auto ckp = make_checkpoint(workspace()); + InitIndexers(*ckp, 1, 1); + ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); + OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), 1, 1); + auto chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{42}, 1)}); + ASSERT_EQ(chunks.size(), 1); + + EXPECT_THROW(edge_table->BatchBuildEdges( + src_indexer, dst_indexer, + std::make_shared(chunks.front())), + std::runtime_error); + EXPECT_TRUE(edge_table->CanBatchBuild()); + EXPECT_EQ(edge_table->EdgeNum(), 0); +} + +TEST_F(EdgeTableTest, + ConcurrentSingleCheckFailureCancelsSupplierAndDoesNotPublish) { + if (std::thread::hardware_concurrency() < 4) { + GTEST_SKIP() + << "Concurrent uniqueness cancellation needs two consumer workers"; + } + + class ConcurrentSingleCheckThrowingSupplier final + : public IDataChunkSupplier { + public: + ConcurrentSingleCheckThrowingSupplier(std::shared_ptr chunk, + std::atomic& cancel_count) + : chunk_(std::move(chunk)), cancel_count_(cancel_count) {} + + std::shared_ptr GetNextChunk() override { + const auto call = next_call_.fetch_add(1, std::memory_order_relaxed); + if (call == 0) { + return chunk_; + } + if (call == 1) { + throw std::runtime_error("injected concurrent uniqueness failure"); + } + return nullptr; + } + + int64_t RowNum() const override { return 1; } + bool SupportsConcurrentGetNext() const override { return true; } + + void Cancel() override { + cancel_count_.fetch_add(1, std::memory_order_relaxed); + } + + private: + std::shared_ptr chunk_; + std::atomic& cancel_count_; + std::atomic next_call_{0}; + }; + + class ConcurrentSingleCheckThrowingSource final : public IDataChunkSource { + public: + explicit ConcurrentSingleCheckThrowingSource( + std::shared_ptr chunk) + : chunk_(std::move(chunk)) {} + + std::shared_ptr Open() const override { + return std::make_shared( + chunk_, cancel_count_); + } + + std::shared_ptr Open( + const ChunkSourceOptions& options) const override { + ++open_count_; + opened_options_.push_back(options); + return Open(); + } + + bool rewindable() const override { return true; } + int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } + size_t OpenCount() const { return open_count_; } + size_t CancelCount() const { + return cancel_count_.load(std::memory_order_relaxed); + } + const std::vector& OpenedOptions() const { + return opened_options_; + } + + private: + std::shared_ptr chunk_; + mutable size_t open_count_ = 0; + mutable std::atomic cancel_count_{0}; + mutable std::vector opened_options_; + }; + + auto ckp = make_checkpoint(workspace()); + InitIndexers(*ckp, 1, 1); + ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), 1, 1); + auto chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{42}, 1)}); + ASSERT_EQ(chunks.size(), 1); + auto source = + std::make_shared(chunks.front()); + + EXPECT_THROW(edge_table->BatchBuildEdges(src_indexer, dst_indexer, source), + std::runtime_error); + EXPECT_EQ(source->OpenCount(), 1); + EXPECT_EQ(source->CancelCount(), 1); + ASSERT_EQ(source->OpenedOptions().size(), 1); + EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); + EXPECT_EQ(source->OpenedOptions()[0].projected_columns, + (std::vector{0, 1})); + EXPECT_TRUE(edge_table->CanBatchBuild()); + EXPECT_EQ(edge_table->EdgeNum(), 0); +} + +TEST_F(EdgeTableTest, SecondPassFailureDoesNotPublishAllocatedCsr) { + class ThrowAfterChunkSupplier final : public IDataChunkSupplier { + public: + explicit ThrowAfterChunkSupplier(std::shared_ptr chunk) + : chunk_(std::move(chunk)) {} + + std::shared_ptr GetNextChunk() override { + if (chunk_) { + return std::exchange(chunk_, nullptr); + } + throw std::runtime_error("injected second-pass failure"); + } + + int64_t RowNum() const override { return 1; } + + private: + std::shared_ptr chunk_; + }; + + class SecondPassThrowingSource final : public IDataChunkSource { + public: + explicit SecondPassThrowingSource(std::shared_ptr chunk) + : chunk_(std::move(chunk)) {} + + std::shared_ptr Open() const override { + ++open_count_; + if (open_count_ == 1) { + return std::make_shared( + std::vector>{chunk_}); + } + return std::make_shared(chunk_); + } + + bool rewindable() const override { return true; } + size_t OpenCount() const { return open_count_; } + + private: + std::shared_ptr chunk_; + mutable size_t open_count_ = 0; + }; + + auto ckp = make_checkpoint(workspace()); + InitIndexers(*ckp, 1, 1); + ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); + OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), 1, 1); + auto chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{42}, 1)}); + ASSERT_EQ(chunks.size(), 1); + auto source = std::make_shared(chunks.front()); + + EXPECT_THROW(edge_table->BatchBuildEdges(src_indexer, dst_indexer, source), + std::runtime_error); + EXPECT_EQ(source->OpenCount(), 2); + EXPECT_TRUE(edge_table->CanBatchBuild()); + EXPECT_EQ(edge_table->EdgeNum(), 0); +} + +TEST_F(EdgeTableTest, + ConcurrentSecondPassFailureCancelsSupplierAndDoesNotPublish) { + if (std::thread::hardware_concurrency() < 2) { + GTEST_SKIP() << "Concurrent supplier cancellation requires two workers"; + } + + class ConcurrentThrowingSupplier final : public IDataChunkSupplier { + public: + ConcurrentThrowingSupplier(std::shared_ptr chunk, + std::atomic& cancel_count) + : chunk_(std::move(chunk)), cancel_count_(cancel_count) {} + + std::shared_ptr GetNextChunk() override { + const auto call = next_call_.fetch_add(1, std::memory_order_relaxed); + if (call == 0) { + return chunk_; + } + if (call == 1) { + throw std::runtime_error("injected concurrent fill failure"); + } + return nullptr; + } + + int64_t RowNum() const override { return 1; } + bool SupportsConcurrentGetNext() const override { return true; } + + void Cancel() override { + cancel_count_.fetch_add(1, std::memory_order_relaxed); + } + + private: + std::shared_ptr chunk_; + std::atomic& cancel_count_; + std::atomic next_call_{0}; + }; + + class ConcurrentSecondPassThrowingSource final : public IDataChunkSource { + public: + explicit ConcurrentSecondPassThrowingSource( + std::shared_ptr chunk) + : chunk_(std::move(chunk)) {} + + std::shared_ptr Open() const override { + return Open(ChunkSourceOptions()); + } + + std::shared_ptr Open( + const ChunkSourceOptions& /*options*/) const override { + ++open_count_; + if (open_count_ == 1) { + return std::make_shared( + std::vector>{chunk_}); + } + return std::make_shared(chunk_, + cancel_count_); + } + + bool rewindable() const override { return true; } + int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } + size_t OpenCount() const { return open_count_; } + size_t CancelCount() const { + return cancel_count_.load(std::memory_order_relaxed); + } + + private: + std::shared_ptr chunk_; + mutable size_t open_count_ = 0; + mutable std::atomic cancel_count_{0}; + }; + + auto ckp = make_checkpoint(workspace()); + InitIndexers(*ckp, 1, 1); + ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), 1, 1); + auto chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{42}, 1)}); + ASSERT_EQ(chunks.size(), 1); + auto source = + std::make_shared(chunks.front()); + + EXPECT_THROW(edge_table->BatchBuildEdges(src_indexer, dst_indexer, source), + std::runtime_error); + EXPECT_EQ(source->OpenCount(), 2); + EXPECT_EQ(source->CancelCount(), 1); + EXPECT_TRUE(edge_table->CanBatchBuild()); + EXPECT_EQ(edge_table->EdgeNum(), 0); +} + +TEST_F(EdgeTableTest, + NonConcurrentSecondPassFailureCancelsBlockedSupplierAndDoesNotPublish) { + if (std::thread::hardware_concurrency() < 4) { + GTEST_SKIP() << "Bounded pipeline cancellation requires two consumers"; + } + + class BlockingAfterChunkSupplier final : public IDataChunkSupplier { + public: + explicit BlockingAfterChunkSupplier(std::shared_ptr chunk) + : chunk_(std::move(chunk)) {} + + std::shared_ptr GetNextChunk() override { + if (chunk_) { + return std::exchange(chunk_, nullptr); + } + std::unique_lock lock(mutex_); + if (!cancelled_cv_.wait_for(lock, std::chrono::seconds(2), + [&] { return cancelled_; })) { + timed_out_.store(true, std::memory_order_relaxed); + } + return nullptr; + } + + int64_t RowNum() const override { return 1; } + + void Cancel() override { + { + std::lock_guard lock(mutex_); + cancelled_ = true; + } + cancel_count_.fetch_add(1, std::memory_order_relaxed); + cancelled_cv_.notify_all(); + } + + size_t CancelCount() const { + return cancel_count_.load(std::memory_order_relaxed); + } + bool TimedOut() const { return timed_out_.load(std::memory_order_relaxed); } + + private: + std::shared_ptr chunk_; + mutable std::mutex mutex_; + std::condition_variable cancelled_cv_; + bool cancelled_ = false; + std::atomic cancel_count_{0}; + std::atomic timed_out_{false}; + }; + + class NonConcurrentSecondPassFailureSource final : public IDataChunkSource { + public: + NonConcurrentSecondPassFailureSource( + std::shared_ptr valid_chunk, + std::shared_ptr invalid_chunk) + : valid_chunk_(std::move(valid_chunk)), + invalid_chunk_(std::move(invalid_chunk)) {} + + std::shared_ptr Open() const override { + ++open_count_; + if (open_count_ == 1) { + return std::make_shared( + std::vector>{valid_chunk_}); + } + second_pass_supplier_ = + std::make_shared(invalid_chunk_); + return second_pass_supplier_; + } + + bool rewindable() const override { return true; } + int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } + size_t OpenCount() const { return open_count_; } + size_t CancelCount() const { + CHECK(second_pass_supplier_ != nullptr); + return second_pass_supplier_->CancelCount(); + } + bool TimedOut() const { + CHECK(second_pass_supplier_ != nullptr); + return second_pass_supplier_->TimedOut(); + } + + private: + std::shared_ptr valid_chunk_; + std::shared_ptr invalid_chunk_; + mutable size_t open_count_ = 0; + mutable std::shared_ptr second_pass_supplier_; + }; + + auto ckp = make_checkpoint(workspace()); + InitIndexers(*ckp, 1, 1); + ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), 1, 1); + + auto valid_chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{42}, 1)}); + // The source indexer expects int64 endpoints. The uint64 source column makes + // a fill consumer throw after the producer has requested its next chunk. + auto invalid_chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{42}, 1)}); + ASSERT_EQ(valid_chunks.size(), 1); + ASSERT_EQ(invalid_chunks.size(), 1); + auto source = std::make_shared( + valid_chunks.front(), invalid_chunks.front()); + + EXPECT_THROW(edge_table->BatchBuildEdges(src_indexer, dst_indexer, source), + std::exception); + EXPECT_EQ(source->OpenCount(), 2); + EXPECT_EQ(source->CancelCount(), 1); + EXPECT_FALSE(source->TimedOut()); + EXPECT_TRUE(edge_table->CanBatchBuild()); + EXPECT_EQ(edge_table->EdgeNum(), 0); +} + +TEST_F(EdgeTableTest, BatchBuildEdgesUsesIndexerVertexCapacity) { + auto ckp = make_checkpoint(workspace()); + constexpr neug::vid_t kVertexNum = 4097; + constexpr neug::vid_t kCheckpointCapacity = kVertexNum + kVertexNum / 4; + std::vector endpoints = {0}; + std::vector edge_data = {42}; + auto batches = + convert_to_data_chunks({split_column_to_chunks(endpoints, 16), + split_column_to_chunks(endpoints, 16), + split_column_to_chunks(edge_data, 16)}); + + InitIndexers(*ckp, kVertexNum, kVertexNum); + src_indexer.reserve(kCheckpointCapacity); + dst_indexer.reserve(kCheckpointCapacity); + ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); + OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), kVertexNum, + kVertexNum); + BatchBuild(std::move(batches)); + + EXPECT_EQ(edge_table->EdgeNum(), 1); + auto out_csr = edge_table->TakeOutCsr(); + auto in_csr = edge_table->TakeInCsr(); + ASSERT_NE(out_csr, nullptr); + ASSERT_NE(in_csr, nullptr); + EXPECT_EQ(out_csr->size(), kCheckpointCapacity); + EXPECT_EQ(in_csr->size(), kCheckpointCapacity); +} + +TEST_F(EdgeTableTest, BatchAddEdgesUsesIndexerVertexCapacity) { + auto ckp = make_checkpoint(workspace()); + constexpr neug::vid_t kVertexNum = 16; + constexpr neug::vid_t kVertexCapacity = 128; + InitIndexers(*ckp, kVertexNum, kVertexNum); + src_indexer.reserve(kVertexCapacity); + dst_indexer.reserve(kVertexCapacity); + ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); + OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), kVertexNum, + kVertexNum); + + auto chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{42}, 1)}); + BatchInsert(std::move(chunks)); + + auto out_csr = edge_table->TakeOutCsr(); + auto in_csr = edge_table->TakeInCsr(); + ASSERT_NE(out_csr, nullptr); + ASSERT_NE(in_csr, nullptr); + EXPECT_EQ(out_csr->size(), kVertexCapacity); + EXPECT_EQ(in_csr->size(), kVertexCapacity); +} + TEST_F(EdgeTableTest, TestBatchAddEdgesUnbundled) { auto ckp = make_checkpoint(workspace()); int64_t src_num = 100; diff --git a/tests/storage/test_mutable_csr.cc b/tests/storage/test_mutable_csr.cc index 8ab438da3..d4eb98eb4 100644 --- a/tests/storage/test_mutable_csr.cc +++ b/tests/storage/test_mutable_csr.cc @@ -15,6 +15,7 @@ #include #include +#include #include #include #include @@ -551,6 +552,74 @@ class MutableCsrTest : public ::testing::Test { }; TYPED_TEST_SUITE(MutableCsrTest, Datatypes); +TEST(MutableCsrBulkBuildAccessTest, ReservesTwentyPercentPerVertex) { + auto test_dir = make_unique_test_dir("mutable_csr_bulk_build"); + CheckpointManager workspace; + workspace.Open(test_dir.string()); + auto ckp = make_checkpoint(workspace); + MutableCsr csr; + csr.Open(*ckp, ModuleDescriptor(), MemoryLevel::kInMemory); + + MutableCsrBulkBuildAccess writer(csr); + writer.PrepareBuild(3); + std::vector counters; + for (int thread = 0; thread < 4; ++thread) { + counters.emplace_back([&writer] { writer.CountConcurrent(0, 5); }); + } + for (auto& counter : counters) { + counter.join(); + } + writer.CountConcurrent(1, 2); + writer.AllocateFromCounts(); + + EXPECT_EQ(writer.ExpectedDegree(0), 20); + EXPECT_EQ(writer.ExpectedDegree(1), 2); + EXPECT_EQ(writer.ExpectedDegree(2), 0); + EXPECT_EQ(writer.ReservedCapacityForVertex(0), 24); + EXPECT_EQ(writer.ReservedCapacityForVertex(1), 3); + EXPECT_EQ(writer.ReservedCapacityForVertex(2), 0); + + std::vector fillers; + for (int thread = 0; thread < 4; ++thread) { + fillers.emplace_back([&writer] { + const auto begin = writer.ReserveConcurrent(0, 5); + for (int slot = begin; slot < begin + 5; ++slot) { + writer.PutAt(0, slot, static_cast(slot), slot, 0); + } + }); + } + for (auto& filler : fillers) { + filler.join(); + } + for (int i = 0; i < 2; ++i) { + writer.PutSerial(1, static_cast(i), i, 0); + } + writer.Finish(); + EXPECT_EQ(csr.edge_num(), 22); + auto built_edges = csr.get_generic_view(0).get_edges(0); + std::vector built_neighbors; + for (auto it = built_edges.begin(); it != built_edges.end(); ++it) { + built_neighbors.push_back(*it); + } + ASSERT_EQ(built_neighbors.size(), 20); + std::sort(built_neighbors.begin(), built_neighbors.end()); + for (int i = 0; i < 20; ++i) { + EXPECT_EQ(built_neighbors[static_cast(i)], static_cast(i)); + } + + auto before = csr.get_generic_view(0).get_edges(0).start_ptr; + Allocator allocator(MemoryLevel::kInMemory, ""); + for (int i = 20; i < 24; ++i) { + csr.put_edge(0, static_cast(i), i, 0, allocator); + } + auto after = csr.get_generic_view(0).get_edges(0).start_ptr; + EXPECT_EQ(after, before); + EXPECT_EQ(csr.edge_num(), 26); + + workspace.Close(); + std::filesystem::remove_all(test_dir); +} + TYPED_TEST(MutableCsrTest, TestCsrType) { MutableCsr mutable_csr; EXPECT_EQ(mutable_csr.csr_type(), CsrType::kMutable); diff --git a/tests/storage/test_vertex_table.cc b/tests/storage/test_vertex_table.cc index a6b4adeb7..1a5b8b49d 100644 --- a/tests/storage/test_vertex_table.cc +++ b/tests/storage/test_vertex_table.cc @@ -691,6 +691,7 @@ TEST_F(VertexTableTest, VertexTableResizeTest) { EXPECT_EQ(table.VertexNum(), 10000); EXPECT_EQ(table.LidNum(), 10000); + EXPECT_EQ(table.Capacity(), 12500); table.Compact(true); EXPECT_EQ(table.get_vertex_timestamp().InitVertexNum(), 10000); @@ -706,6 +707,29 @@ TEST_F(VertexTableTest, VertexTableResizeTest) { } } +TEST_F(VertexTableTest, BatchBuildVerticesFromRepeatableSource) { + neug::VertexTable table(schema_.get_vertex_schema(v_label_id_)); + auto ckp = make_checkpoint(Workspace()); + OpenVertexTableLegacy(table, ckp, neug::CheckpointManifest(), memory_level_); + + constexpr size_t kVertexNum = 4097; + auto source = + std::make_shared(generate_data_chunks(kVertexNum)); + EXPECT_TRUE(table.CanBatchBuild()); + table.BatchBuildVertices(source); + + EXPECT_EQ(table.VertexNum(), kVertexNum); + EXPECT_EQ(table.LidNum(), kVertexNum); + EXPECT_EQ(table.Capacity(), kVertexNum + kVertexNum / 4); + EXPECT_EQ(source->OpenCount(), 1); + EXPECT_FALSE(table.CanBatchBuild()); + for (int64_t oid : {int64_t{0}, int64_t{127}, int64_t{4096}}) { + neug::vid_t lid; + EXPECT_TRUE(table.get_index(neug::Value::INT64(oid), lid)); + EXPECT_EQ(table.GetOid(lid), neug::Value::INT64(oid)); + } +} + TEST_F(VertexTableTest, VertexTimestampValidVertexNum) { auto ckp = make_checkpoint(Workspace()); neug::VertexTimestamp vts; diff --git a/tests/unittest/utils.h b/tests/unittest/utils.h index 8374ddbcd..ea4113f04 100644 --- a/tests/unittest/utils.h +++ b/tests/unittest/utils.h @@ -18,8 +18,10 @@ #include #include +#include #include #include +#include #include #include #include @@ -71,6 +73,42 @@ class GeneratedChunkSupplier : public neug::IDataChunkSupplier { std::vector> chunks_; }; +class GeneratedChunkSource final : public neug::IDataChunkSource { + public: + explicit GeneratedChunkSource( + std::vector> chunks, + int64_t estimated_bytes = -1) + : chunks_(std::move(chunks)), estimated_bytes_(estimated_bytes) {} + + std::shared_ptr Open() const override { + ++open_count_; + return std::make_shared( + std::vector>(chunks_)); + } + + std::shared_ptr Open( + const neug::ChunkSourceOptions& options) const override { + opened_projections_.push_back(options.projected_columns); + return neug::IDataChunkSource::Open(options); + } + + bool rewindable() const override { return true; } + + int64_t EstimatedBytes() const override { return estimated_bytes_; } + + size_t OpenCount() const { return open_count_; } + + const std::vector>& OpenedProjections() const { + return opened_projections_; + } + + private: + std::vector> chunks_; + int64_t estimated_bytes_; + mutable size_t open_count_ = 0; + mutable std::vector> opened_projections_; +}; + template std::shared_ptr build_value_column_slice( const std::vector& data, size_t begin, size_t end) { diff --git a/tests/utils/test_reader.cc b/tests/utils/test_reader.cc index 90a4db953..c35b85275 100644 --- a/tests/utils/test_reader.cc +++ b/tests/utils/test_reader.cc @@ -15,6 +15,12 @@ #include "test_reader.h" +#include +#include + +#include "neug/storages/loader/chunk_pipeline_utils.h" +#include "neug/storages/loader/loader_utils.h" + namespace neug { namespace test { @@ -45,6 +51,334 @@ TEST_F(ReaderTest, TestBasicCsvRead) { EXPECT_EQ(ctx.row_num(), 3); } +TEST_F(ReaderTest, CsvChunkSourceCanBeReopened) { + createCsvFile("repeatable.csv", "id|name\n1|Alice\n2|Bob\n3|Carol\n"); + std::vector column_names = {"id", "name"}; + std::vector> column_types = { + createInt32Type(), createStringType()}; + auto shared_state = + createSharedState("repeatable.csv", column_names, column_types, + {{"skip_rows", "1"}, {"batch_read", "true"}}); + auto reader = createCsvReader(shared_state); + auto source = reader->createChunkSource(); + + ASSERT_NE(source, nullptr); + EXPECT_TRUE(source->rewindable()); + for (int pass = 0; pass < 2; ++pass) { + auto supplier = source->Open(); + ASSERT_NE(supplier, nullptr); + size_t rows = 0; + while (auto chunk = supplier->GetNextChunk()) { + EXPECT_EQ(chunk->col_num(), 2); + rows += chunk->row_num(); + } + EXPECT_EQ(rows, 3); + } +} + +TEST_F(ReaderTest, CsvChunkSourcePartitionsQuotedRecords) { + createCsvFile( + "partitioned.csv", + "id|name\r\n\r\n1|\"Alice|Smith\"\r\n2|\"line one\nline two\"\r\n" + "3|Carol\r\n4|\"D\"\"Angelo\"\r\n5|Last"); + std::vector column_names = {"id", "name"}; + std::vector> column_types = { + createInt32Type(), createStringType()}; + auto shared_state = createSharedState( + "partitioned.csv", column_names, column_types, + {{"skip_rows", "1"}, {"batch_read", "true"}, {"batch_size", "1"}}); + auto reader = createCsvReader(shared_state); + auto source = reader->createChunkSource(); + ASSERT_NE(source, nullptr); + + auto read_rows = [](const std::shared_ptr& supplier) { + std::vector> rows; + while (auto chunk = supplier->GetNextChunk()) { + EXPECT_EQ(chunk->col_num(), 2); + for (size_t row = 0; row < chunk->row_num(); ++row) { + rows.emplace_back(chunk->get(0)->get_elem(row).GetValue(), + chunk->get(1)->get_elem(row).GetValue()); + } + } + std::sort(rows.begin(), rows.end()); + return rows; + }; + + auto expected = read_rows(source->Open()); + ChunkSourceOptions options; + options.parallel_enabled = true; + options.producer_count = 4; + options.queue_capacity = 8; + options.preserve_order = false; + auto partitioned = source->Open(options); + ASSERT_NE(partitioned, nullptr); + EXPECT_TRUE(partitioned->SupportsConcurrentGetNext()); + EXPECT_EQ(partitioned->RowNum(), 6); // Header is an intentional overcount. + auto actual = read_rows(partitioned); + + ASSERT_EQ(actual, expected); + ASSERT_EQ(actual.size(), 5); + EXPECT_EQ(actual[0], std::make_pair(1, std::string("Alice|Smith"))); + EXPECT_EQ(actual[1], std::make_pair(2, std::string("line one\nline two"))); + EXPECT_EQ(actual[3], std::make_pair(4, std::string("D\"Angelo"))); + EXPECT_EQ(actual[4], std::make_pair(5, std::string("Last"))); +} + +TEST_F(ReaderTest, PartitionedCsvCarriesSkipAcrossRanges) { + std::string long_quoted_field(64 * 1024, 'x'); + for (size_t i = 32; i < long_quoted_field.size(); i += 64) { + long_quoted_field[i] = '\n'; + } + createCsvFile("partition-skip.csv", + "0|\"" + long_quoted_field + + "\"\r\n1|also-skipped\r\n2|kept\r\n3|last\r\n"); + + CsvReadConfig config; + config.delimiter = '|'; + config.quoting = true; + config.skip_rows = 2; + config.chunk_size = 1; + config.column_names = {"id", "name"}; + config.include_columns = config.column_names; + config.column_types.emplace("id", DataType(DataTypeId::kInt32)); + config.column_types.emplace("name", DataType(DataTypeId::kVarchar)); + CSVChunkSource source( + {std::string(ARROW_READER_TEST_DIR) + "/partition-skip.csv"}, config); + + ChunkSourceOptions options; + options.parallel_enabled = true; + options.producer_count = 4; + options.queue_capacity = 4; + options.preserve_order = false; + auto supplier = source.Open(options); + ASSERT_NE(supplier, nullptr); + + std::vector ids; + while (auto chunk = supplier->GetNextChunk()) { + for (size_t row = 0; row < chunk->row_num(); ++row) { + ids.push_back(chunk->get(0)->get_elem(row).GetValue()); + } + } + std::sort(ids.begin(), ids.end()); + EXPECT_EQ(ids, (std::vector{2, 3})); +} + +TEST_F(ReaderTest, PartitionPlannerMatchesParserWhenDoubleQuoteIsFalse) { + std::string multiline = "first \"\"quoted\n"; + multiline.append(64 * 1024, 'z'); + multiline += "\ncontinued\"\" tail"; + createCsvFile("partition-double-quote.csv", + "0|\"" + multiline + "\"\n1|one\n2|two\n"); + + CsvReadConfig config; + config.delimiter = '|'; + config.quoting = true; + config.double_quote = false; + config.chunk_size = 1; + config.column_names = {"id", "name"}; + config.include_columns = config.column_names; + config.column_types.emplace("id", DataType(DataTypeId::kInt32)); + config.column_types.emplace("name", DataType(DataTypeId::kVarchar)); + CSVChunkSource source( + {std::string(ARROW_READER_TEST_DIR) + "/partition-double-quote.csv"}, + config); + + auto read_ids = [](std::shared_ptr supplier) { + std::vector ids; + while (auto chunk = supplier->GetNextChunk()) { + for (size_t row = 0; row < chunk->row_num(); ++row) { + ids.push_back(chunk->get(0)->get_elem(row).GetValue()); + } + } + std::sort(ids.begin(), ids.end()); + return ids; + }; + const auto expected = read_ids(source.Open()); + + ChunkSourceOptions options; + options.parallel_enabled = true; + options.producer_count = 4; + options.queue_capacity = 4; + options.preserve_order = false; + EXPECT_EQ(read_ids(source.Open(options)), expected); + EXPECT_EQ(expected, (std::vector{0, 1, 2})); +} + +TEST_F(ReaderTest, CsvChunkSourcePushesProjectionIntoPartitionedParsing) { + createCsvFile("partition-projection.csv", + "id|ignored|score\n1|Alice|10\n2|Bob|20\n"); + CsvReadConfig config; + config.delimiter = '|'; + config.skip_rows = 1; + config.chunk_size = 1; + config.column_names = {"id", "ignored", "score"}; + config.include_columns = config.column_names; + config.column_types.emplace("id", DataType(DataTypeId::kInt32)); + config.column_types.emplace("ignored", DataType(DataTypeId::kVarchar)); + config.column_types.emplace("score", DataType(DataTypeId::kInt32)); + CSVChunkSource source( + {std::string(ARROW_READER_TEST_DIR) + "/partition-projection.csv"}, + config); + + ChunkSourceOptions options; + options.parallel_enabled = true; + options.producer_count = 2; + options.queue_capacity = 2; + options.preserve_order = false; + options.projected_columns = {2, 0}; + auto supplier = source.Open(options); + ASSERT_NE(supplier, nullptr); + + std::vector> rows; + while (auto chunk = supplier->GetNextChunk()) { + ASSERT_EQ(chunk->col_num(), 2); + for (size_t row = 0; row < chunk->row_num(); ++row) { + rows.emplace_back(chunk->get(0)->get_elem(row).GetValue(), + chunk->get(1)->get_elem(row).GetValue()); + } + } + std::sort(rows.begin(), rows.end()); + EXPECT_EQ(rows, (std::vector>{{10, 1}, {20, 2}})); +} + +TEST_F(ReaderTest, ChunkPipelineAllocationSharesHardwareBudget) { + constexpr int64_t kGiB = 1024LL * 1024 * 1024; + auto allocation = resolve_chunk_pipeline_allocation(kGiB, true, false, 16); + EXPECT_TRUE(allocation.parallel_enabled); + EXPECT_EQ(allocation.producer_count, 8); + EXPECT_EQ(allocation.consumer_count, 8); + EXPECT_LE(allocation.producer_count + allocation.consumer_count, 16); + EXPECT_EQ(allocation.queue_capacity, 16); + + auto ordered = resolve_chunk_pipeline_allocation(kGiB, true, true, 16); + EXPECT_FALSE(ordered.parallel_enabled); + EXPECT_EQ(ordered.producer_count, 1); + EXPECT_EQ(ordered.consumer_count, 1); + + auto disabled = resolve_chunk_pipeline_allocation(kGiB, false, false, 16); + EXPECT_FALSE(disabled.parallel_enabled); + EXPECT_EQ(disabled.producer_count, 1); + EXPECT_EQ(disabled.consumer_count, 1); + + auto small = + resolve_chunk_pipeline_allocation(128LL * 1024 * 1024, true, false, 16); + EXPECT_FALSE(small.parallel_enabled); + EXPECT_EQ(small.producer_count, 1); + EXPECT_EQ(small.consumer_count, 1); + + auto two_workers = resolve_chunk_pipeline_allocation(kGiB, true, false, 2); + EXPECT_TRUE(two_workers.parallel_enabled); + EXPECT_EQ(two_workers.producer_count, 1); + EXPECT_EQ(two_workers.consumer_count, 1); + + auto maximum_size = resolve_chunk_pipeline_allocation( + std::numeric_limits::max(), true, false, 16); + EXPECT_TRUE(maximum_size.parallel_enabled); + EXPECT_EQ(maximum_size.producer_count, 8); + EXPECT_EQ(maximum_size.consumer_count, 8); +} + +TEST_F(ReaderTest, CsvChunkSourceHonorsParallelFalse) { + createCsvFile("serial.csv", "id|name\n1|Alice\n2|Bob\n"); + std::vector column_names = {"id", "name"}; + std::vector> column_types = { + createInt32Type(), createStringType()}; + auto shared_state = createSharedState( + "serial.csv", column_names, column_types, + {{"skip_rows", "1"}, {"batch_read", "true"}, {"parallel", "false"}}); + auto source = createCsvReader(shared_state)->createChunkSource(); + ASSERT_NE(source, nullptr); + EXPECT_FALSE(source->ParallelEnabled()); + + ChunkSourceOptions options; + options.parallel_enabled = true; + options.producer_count = 4; + options.queue_capacity = 8; + options.preserve_order = false; + auto supplier = source->Open(options); + ASSERT_NE(supplier, nullptr); + EXPECT_FALSE(supplier->SupportsConcurrentGetNext()); +} + +TEST_F(ReaderTest, PartitionedCsvSkipsHeaderForEveryFile) { + createCsvFile("part-a.csv", "id|name\n1|Alice\n2|Bob\n"); + createCsvFile("part-b.csv", "id|name\n3|Carol\n4|Dave"); + CsvReadConfig config; + config.delimiter = '|'; + config.quoting = true; + config.escaping = false; + config.skip_rows = 1; + config.chunk_size = 1; + config.column_names = {"id", "name"}; + config.include_columns = config.column_names; + config.column_types.emplace("id", DataType(DataTypeId::kInt32)); + config.column_types.emplace("name", DataType(DataTypeId::kVarchar)); + auto path = [](const char* file) { + return std::string(ARROW_READER_TEST_DIR) + "/" + file; + }; + CSVChunkSource source({path("part-a.csv"), path("part-b.csv")}, config); + + ChunkSourceOptions options; + options.parallel_enabled = true; + options.producer_count = 4; + options.queue_capacity = 8; + options.preserve_order = false; + auto supplier = source.Open(options); + ASSERT_NE(supplier, nullptr); + EXPECT_EQ(supplier->RowNum(), + 6); // Two headers are counted as reserve hints. + + std::vector ids; + while (auto chunk = supplier->GetNextChunk()) { + for (size_t row = 0; row < chunk->row_num(); ++row) { + ids.push_back(chunk->get(0)->get_elem(row).GetValue()); + } + } + std::sort(ids.begin(), ids.end()); + EXPECT_EQ(ids, (std::vector{1, 2, 3, 4})); +} + +TEST_F(ReaderTest, PartitionedCsvPropagatesProducerErrors) { + createCsvFile("partition-error.csv", + "id|name\n1|Alice\nnot-an-int|Broken\n3|Carol\n4|Dave"); + CsvReadConfig config; + config.delimiter = '|'; + config.quoting = true; + config.escaping = false; + config.skip_rows = 1; + config.chunk_size = 1; + config.column_names = {"id", "name"}; + config.include_columns = config.column_names; + config.column_types.emplace("id", DataType(DataTypeId::kInt32)); + config.column_types.emplace("name", DataType(DataTypeId::kVarchar)); + CSVChunkSource source( + {std::string(ARROW_READER_TEST_DIR) + "/partition-error.csv"}, config); + + ChunkSourceOptions options; + options.parallel_enabled = true; + options.producer_count = 4; + options.queue_capacity = 2; + options.preserve_order = false; + auto supplier = source.Open(options); + ASSERT_NE(supplier, nullptr); + EXPECT_ANY_THROW({ + while (supplier->GetNextChunk()) {} + }); +} + +TEST_F(ReaderTest, CsvChunkSourceRejectsPostReadProjection) { + createCsvFile("projected.csv", "id|name\n1|Alice\n"); + std::vector column_names = {"id", "name"}; + std::vector> column_types = { + createInt32Type(), createStringType()}; + auto shared_state = + createSharedState("projected.csv", column_names, column_types, {}); + shared_state->projectColumns = {"id"}; + + auto reader = createCsvReader(shared_state); + EXPECT_EQ(reader->createChunkSource(), nullptr); +} + // Test 2: CSV with different delimiter (tab) TEST_F(ReaderTest, TestCsvWithTabDelimiter) { createCsvFile("test2.csv", "id\tname\tage\n1\tAlice\t95.5\n2\tBob\t87.0\n"); From 78eb6016bcae3dedd675ec5403587cc707b8c54c Mon Sep 17 00:00:00 2001 From: "xiaolei.zl" Date: Fri, 17 Jul 2026 10:10:25 +0800 Subject: [PATCH 2/8] refactor code, reduce dummy code --- .../function/import/csv_read_function.h | 5 +- .../neug/compiler/function/read_function.h | 8 +- .../execute/ops/batch/batch_insert_edge.h | 5 +- .../execute/ops/batch/batch_insert_vertex.h | 5 +- .../execute/ops/batch/batch_update_utils.h | 41 +- include/neug/storages/csr/mutable_csr.h | 60 +- include/neug/storages/graph/edge_table.h | 20 +- include/neug/storages/graph/graph_interface.h | 11 - include/neug/storages/graph/property_graph.h | 10 - include/neug/storages/graph/vertex_table.h | 38 +- .../storages/loader/chunk_pipeline_utils.h | 97 +- include/neug/storages/loader/loader_utils.h | 41 +- include/neug/utils/io/read/csv/csv_reader.h | 3 +- .../execute/ops/batch/batch_insert_edge.cc | 140 +-- .../execute/ops/batch/batch_insert_vertex.cc | 128 +-- .../execute/ops/batch/batch_update_utils.cc | 179 +--- src/storages/csr/mutable_csr.cc | 2 +- src/storages/graph/edge_table.cc | 981 +++++------------- src/storages/graph/graph_interface.cc | 30 +- src/storages/graph/property_graph.cc | 49 +- src/storages/graph/vertex_table.cc | 119 +-- src/storages/loader/loader_utils.cc | 157 +-- src/utils/io/read/csv/csv_reader.cc | 6 +- tests/storage/test_copy_temp.cc | 17 +- tests/storage/test_edge_table.cc | 920 ++++++---------- tests/storage/test_mutable_csr.cc | 7 - tests/storage/test_property_graph.cc | 2 +- tests/storage/test_vertex_table.cc | 10 +- tests/unittest/utils.h | 103 +- tests/utils/test_reader.cc | 105 +- 30 files changed, 971 insertions(+), 2328 deletions(-) diff --git a/include/neug/compiler/function/import/csv_read_function.h b/include/neug/compiler/function/import/csv_read_function.h index 1c9be2f9d..caaeee5fb 100644 --- a/include/neug/compiler/function/import/csv_read_function.h +++ b/include/neug/compiler/function/import/csv_read_function.h @@ -134,7 +134,8 @@ struct CSVReadFunction { } static std::shared_ptr sourceFunc( - std::shared_ptr state) { + std::shared_ptr state, + std::vector projected_columns) { if (!state) { THROW_INVALID_ARGUMENT_EXCEPTION("State is null"); } @@ -156,7 +157,7 @@ struct CSVReadFunction { std::make_unique(source_state); auto reader = std::make_unique( source_state, std::move(options_builder)); - return reader->createChunkSource(); + return reader->createChunkSource(std::move(projected_columns)); } static std::shared_ptr sniffFunc( diff --git a/include/neug/compiler/function/read_function.h b/include/neug/compiler/function/read_function.h index e0253ffc8..298bdc4c6 100644 --- a/include/neug/compiler/function/read_function.h +++ b/include/neug/compiler/function/read_function.h @@ -35,11 +35,11 @@ namespace function { using read_exec_func_t = std::function state)>; -/// Creates a repeatable source for an opt-in bulk ingestion fast path. The -/// function must leave its input state unchanged because callers may fall back -/// to execFunc when the source or destination is not eligible for bulk build. +/// Creates a repeatable source for terminal ingestion. Storage may consume it +/// once through normal BatchAdd or reopen it for a staged bulk build. using read_source_func_t = std::function( - std::shared_ptr state)>; + std::shared_ptr state, + std::vector projected_columns)>; // The function used to sniff/infer file column names and their types from // external data sources. diff --git a/include/neug/execution/execute/ops/batch/batch_insert_edge.h b/include/neug/execution/execute/ops/batch/batch_insert_edge.h index f4219ae84..080e0c1ab 100644 --- a/include/neug/execution/execute/ops/batch/batch_insert_edge.h +++ b/include/neug/execution/execute/ops/batch/batch_insert_edge.h @@ -40,9 +40,8 @@ class BatchInsertEdgeOprBuilder : public IOperatorBuilder { } }; -/// Fuses only a terminal, empty-sink COPY FROM plan. The implementation -/// reverts to the normal reader/Context path unless runtime bulk eligibility -/// is established. +/// Fuses only a terminal, empty-sink COPY FROM plan. Storage chooses staged +/// build or normal BatchAdd from the supplied repeatable source. class BatchInsertEdgeFromSourceOprBuilder : public IOperatorBuilder { public: BatchInsertEdgeFromSourceOprBuilder() = default; diff --git a/include/neug/execution/execute/ops/batch/batch_insert_vertex.h b/include/neug/execution/execute/ops/batch/batch_insert_vertex.h index 798ad59fd..a96e87a0a 100644 --- a/include/neug/execution/execute/ops/batch/batch_insert_vertex.h +++ b/include/neug/execution/execute/ops/batch/batch_insert_vertex.h @@ -39,9 +39,8 @@ class BatchInsertVertexOprBuilder : public IOperatorBuilder { } }; -/// Fuses only a terminal, empty-sink COPY FROM plan. The implementation -/// reverts to the normal reader/Context path unless runtime bulk eligibility -/// is established. +/// Fuses only a terminal, empty-sink COPY FROM plan. Storage chooses staged +/// build or normal BatchAdd from the supplied repeatable source. class BatchInsertVertexFromSourceOprBuilder : public IOperatorBuilder { public: BatchInsertVertexFromSourceOprBuilder() = default; diff --git a/include/neug/execution/execute/ops/batch/batch_update_utils.h b/include/neug/execution/execute/ops/batch/batch_update_utils.h index af53624f0..2cf244c25 100644 --- a/include/neug/execution/execute/ops/batch/batch_update_utils.h +++ b/include/neug/execution/execute/ops/batch/batch_update_utils.h @@ -21,8 +21,12 @@ #include "neug/utils/property/types.h" namespace physical { +class PhysicalPlan; class PropertyMapping; -} +} // namespace physical +namespace common { +class NameOrId; +} // namespace common namespace google { namespace protobuf { template @@ -32,9 +36,14 @@ class RepeatedPtrField; namespace neug { class IDataChunkSupplier; -class IDataChunkSource; class Schema; class StorageReadInterface; +namespace function { +struct ReadFunction; +} +namespace reader { +struct ReadSharedState; +} namespace execution { namespace ops { @@ -66,14 +75,28 @@ std::shared_ptr create_data_chunk_supplier( const Context& ctx, const std::vector>& prop_mappings); -std::shared_ptr create_data_chunk_source( - std::shared_ptr source, - const std::vector>& prop_mappings); +bool resolve_vertex_label_id(const Schema& schema, + const ::common::NameOrId& type, label_t& label_id); + +struct BatchInsertInput { + std::shared_ptr supplier; + Context output; +}; -/// Selects staged bulk builders only for sources large enough to amortize their -/// two bounded parsing passes. Set NEUG_COPY_BULK_BUILD=true/false to force a -/// decision while benchmarking or rolling back. -bool should_use_copy_bulk_build(const IDataChunkSource& source); +struct BatchInsertSource { + std::shared_ptr state; + function::ReadFunction* read_function; +}; + +bool is_terminal_batch_insert(const physical::PhysicalPlan& plan, int op_idx); + +BatchInsertSource build_batch_insert_source(const physical::PhysicalPlan& plan, + int op_idx); + +BatchInsertInput create_batch_insert_input( + const std::shared_ptr& shared_state, + const function::ReadFunction& read_function, + const std::vector>& prop_mappings); std::vector match_files_with_pattern(const std::string& file_path); diff --git a/include/neug/storages/csr/mutable_csr.h b/include/neug/storages/csr/mutable_csr.h index 892456216..9fa2e7726 100644 --- a/include/neug/storages/csr/mutable_csr.h +++ b/include/neug/storages/csr/mutable_csr.h @@ -255,6 +255,16 @@ class MutableCsr : public TypedCsrBase { CsrPrefetchPolicy prefetch_policy_; void refresh_prefetch_policy(); + static int reserved_capacity(int degree) { + CHECK_GE(degree, 0); + if (degree == 0) { + return 0; + } + const auto reserved = + std::ceil(degree * NeugDBConfig::DEFAULT_RESERVE_RATIO); + CHECK_LE(reserved, static_cast(std::numeric_limits::max())); + return static_cast(reserved); + } size_t vertex_capacity() const { if (!degree_list_) { @@ -500,11 +510,9 @@ class MutableCsrBulkBuildAccess { public: using csr_t = MutableCsr; using nbr_t = typename csr_t::nbr_t; - static constexpr bool kSupportsDisjointConcurrentFill = true; static constexpr bool kStoresEdges = true; static constexpr bool kNeedsDegreeCount = true; static constexpr bool kChecksSingleUniqueness = false; - static constexpr bool kNeedsConcurrentGrouping = true; static constexpr bool kTracksInputEdgeCount = false; explicit MutableCsrBulkBuildAccess(csr_t& csr) : csr_(csr) {} @@ -553,7 +561,7 @@ class MutableCsrBulkBuildAccess { size_t total_capacity = 0; for (size_t i = 0; i < vertex_capacity_; ++i) { const auto degree = degrees_[i].load(std::memory_order_relaxed); - const auto reserved_capacity = ReservedCapacity(degree); + const auto reserved_capacity = csr_t::reserved_capacity(degree); // During the fill pass cap_list_ stores the exact expected degree. This // lets range reservations validate both passes without allocating a // second O(V) metadata array. Finish() converts it to runtime capacity. @@ -570,23 +578,13 @@ class MutableCsrBulkBuildAccess { size_t offset = 0; for (size_t i = 0; i < vertex_capacity_; ++i) { - const auto reserved_capacity = ReservedCapacity(capacities_[i]); + const auto reserved_capacity = csr_t::reserved_capacity(capacities_[i]); adj_lists_[i] = reserved_capacity == 0 ? nullptr : nbrs_ + offset; offset += static_cast(reserved_capacity); degrees_[i].store(0, std::memory_order_relaxed); } } - int ExpectedDegree(vid_t src) const { - CHECK_LT(src, vertex_capacity_); - return capacities_[src]; - } - - int ReservedCapacityForVertex(vid_t src) const { - CHECK_LT(src, vertex_capacity_); - return ReservedCapacity(capacities_[src]); - } - int ReserveSerial(vid_t src, int count) { CHECK_LT(src, vertex_capacity_); CHECK_GT(count, 0); @@ -627,24 +625,13 @@ class MutableCsrBulkBuildAccess { CHECK_EQ(degree, capacities_[i]) << "Bulk edge count/fill mismatch for vertex " << i; edge_num += static_cast(degree); - capacities_[i] = ReservedCapacity(degree); + capacities_[i] = csr_t::reserved_capacity(degree); } csr_.edge_num_.store(edge_num, std::memory_order_relaxed); csr_.refresh_prefetch_policy(); } private: - static int ReservedCapacity(int degree) { - CHECK_GE(degree, 0); - if (degree == 0) { - return 0; - } - const auto reserved = - std::ceil(degree * NeugDBConfig::DEFAULT_RESERVE_RATIO); - CHECK_LE(reserved, static_cast(std::numeric_limits::max())); - return static_cast(reserved); - } - void refresh_metadata_ptrs() { vertex_capacity_ = csr_.vertex_capacity(); adj_lists_ = reinterpret_cast( @@ -674,11 +661,9 @@ class SingleMutableCsrBulkBuildAccess { public: using csr_t = SingleMutableCsr; using nbr_t = typename csr_t::nbr_t; - static constexpr bool kSupportsDisjointConcurrentFill = true; static constexpr bool kStoresEdges = true; static constexpr bool kNeedsDegreeCount = false; static constexpr bool kChecksSingleUniqueness = true; - static constexpr bool kNeedsConcurrentGrouping = false; static constexpr bool kTracksInputEdgeCount = true; explicit SingleMutableCsrBulkBuildAccess(csr_t& csr) : csr_(csr) {} @@ -692,8 +677,6 @@ class SingleMutableCsrBulkBuildAccess { nbrs_[i].timestamp.store(INVALID_TIMESTAMP, std::memory_order_relaxed); } csr_.edge_num_.store(0, std::memory_order_relaxed); - edge_count_ = 0; - input_edge_count_precomputed_ = false; } void AllocateFromCounts() {} @@ -713,11 +696,6 @@ class SingleMutableCsrBulkBuildAccess { return previous != INVALID_TIMESTAMP; } - void SetInputEdgeCount(uint64_t count) { - edge_count_ = count; - input_edge_count_precomputed_ = true; - } - void PutSerial(vid_t src, vid_t dst, const EDATA_T& data, timestamp_t ts) { PutConcurrent(src, dst, data, ts); } @@ -732,15 +710,11 @@ class SingleMutableCsrBulkBuildAccess { } void RecordFilledEdges(size_t count) { - if (!input_edge_count_precomputed_) { - edge_count_ += static_cast(count); - } + csr_.edge_num_.fetch_add(static_cast(count), + std::memory_order_relaxed); } - void Finish() { - csr_.edge_num_.store(edge_count_, std::memory_order_relaxed); - csr_.refresh_prefetch_policy(); - } + void Finish() { csr_.refresh_prefetch_policy(); } private: void refresh_ptrs() { @@ -750,8 +724,6 @@ class SingleMutableCsrBulkBuildAccess { } csr_t& csr_; - uint64_t edge_count_ = 0; - bool input_edge_count_precomputed_ = false; size_t vertex_capacity_ = 0; nbr_t* nbrs_ = nullptr; }; diff --git a/include/neug/storages/graph/edge_table.h b/include/neug/storages/graph/edge_table.h index 652d5cfc9..4758471d5 100644 --- a/include/neug/storages/graph/edge_table.h +++ b/include/neug/storages/graph/edge_table.h @@ -136,22 +136,6 @@ class EdgeTable { const IndexerType& dst_indexer, std::shared_ptr supplier); - bool CanBatchBuild() const { - if (!meta_ || !meta_->is_bundled() || EdgeNum() != 0) { - return false; - } - // The staged writer handles mutable CSR layouts and the no-adjacency - // layout. Immutable CSR keeps its established incremental path. - return (meta_->oe_strategy == EdgeStrategy::kNone || meta_->oe_mutable) && - (meta_->ie_strategy == EdgeStrategy::kNone || meta_->ie_mutable); - } - - /// Builds a fresh bundled edge table from a repeatable source and swaps it - /// into place only after both passes succeed. - void BatchBuildEdges(const IndexerType& src_indexer, - const IndexerType& dst_indexer, - std::shared_ptr source); - // Add edges in batch to the edge table. void BatchAddEdges(const std::vector& src_lid_list, const std::vector& dst_lid_list, @@ -207,6 +191,10 @@ class EdgeTable { void DetachInAdjlist(vid_t vid, Allocator& alloc); private: + bool TryBatchBuildEdges(const IndexerType& src_indexer, + const IndexerType& dst_indexer, + const std::shared_ptr& source); + void dropAndCreateNewBundledCSR(Checkpoint& ckp, ColumnBase* prev_data_col); void dropAndCreateNewUnbundledCSR(Checkpoint& ckp, bool delete_property); diff --git a/include/neug/storages/graph/graph_interface.h b/include/neug/storages/graph/graph_interface.h index 5c83ac9b6..d604239a1 100644 --- a/include/neug/storages/graph/graph_interface.h +++ b/include/neug/storages/graph/graph_interface.h @@ -649,19 +649,8 @@ class StorageAPUpdateInterface : public StorageUpdateInterface { Status BatchAddVertices( label_t v_label_id, std::shared_ptr supplier) override; - /// Narrow AP-only bulk-build entry points. They intentionally do not - /// widen StorageInsertInterface: transactional stores keep their existing - /// COPY semantics and the execution fast path explicitly opts into AP. - bool CanBatchBuildVertices(label_t v_label_id) const; - Status BatchBuildVertices(label_t v_label_id, - std::shared_ptr source); Status BatchAddEdges(label_t src_label, label_t dst_label, label_t edge_label, std::shared_ptr supplier) override; - bool CanBatchBuildEdges(label_t src_label, label_t dst_label, - label_t edge_label) const; - Status BatchBuildEdges(label_t src_label, label_t dst_label, - label_t edge_label, - std::shared_ptr source); Status BatchDeleteVertices(label_t v_label_id, const std::vector& vids) override; Status BatchDeleteEdges( diff --git a/include/neug/storages/graph/property_graph.h b/include/neug/storages/graph/property_graph.h index ad8135ba6..4b12e6403 100644 --- a/include/neug/storages/graph/property_graph.h +++ b/include/neug/storages/graph/property_graph.h @@ -293,19 +293,9 @@ class PropertyGraph { Status BatchAddVertices(label_t v_label_id, std::shared_ptr supplier); - bool CanBatchBuildVertices(label_t v_label_id) const; - Status BatchBuildVertices(label_t v_label_id, - std::shared_ptr source); - Status BatchAddEdges(label_t src_label, label_t dst_label, label_t edge_label, std::shared_ptr supplier); - bool CanBatchBuildEdges(label_t src_label, label_t dst_label, - label_t edge_label) const; - Status BatchBuildEdges(label_t src_label, label_t dst_label, - label_t edge_label, - std::shared_ptr source); - Status BatchDeleteVertices(label_t v_label_id, const std::vector& vids); diff --git a/include/neug/storages/graph/vertex_table.h b/include/neug/storages/graph/vertex_table.h index e5600b9bf..1b3d111ef 100644 --- a/include/neug/storages/graph/vertex_table.h +++ b/include/neug/storages/graph/vertex_table.h @@ -18,7 +18,6 @@ #include "neug/common/types/value.h" #include "neug/storages/graph/schema.h" #include "neug/storages/graph/vertex_timestamp.h" -#include "neug/storages/loader/loader_utils.h" #include "neug/storages/module/module.h" #include "neug/utils/indexers.h" #include "neug/utils/property/table.h" @@ -28,6 +27,7 @@ namespace neug { class ModuleBroker; class CheckpointManifest; class Checkpoint; +class IDataChunkSupplier; class VertexTableView; class VertexSet { @@ -257,14 +257,7 @@ class VertexTable { void Compact(timestamp_t ts = MAX_TIMESTAMP); - void insert_vertices(std::shared_ptr suppliers); - - bool CanBatchBuild() const { return Size() == 0; } - - /// Builds an empty table from a repeatable source. CSV sources use a cheap - /// raw row count to reserve before their single typed parse. The existing - /// table is untouched until the staged build succeeds. - void BatchBuildVertices(std::shared_ptr source); + void insert_vertices(std::shared_ptr supplier); const VertexTimestamp& get_vertex_timestamp() const { return *v_ts_; } @@ -294,33 +287,6 @@ class VertexTable { return vids; } - void insert_vertices_preallocated( - std::shared_ptr supplier) { - while (auto chunk = supplier->GetNextChunk()) { - auto& columns = chunk->columns; - const auto& property_names = vertex_schema_->property_names; - CHECK_EQ(columns.size(), property_names.size() + 1) - << "Number of columns in the chunk (" << columns.size() - << ") does not match the number of properties (" - << property_names.size() + 1 << ")."; - auto pk_index = std::get<2>(vertex_schema_->primary_keys[0]); - - std::vector> property_columns; - property_columns.reserve(columns.size() - 1); - for (size_t i = 0; i < columns.size(); ++i) { - if (static_cast(i) != pk_index) { - property_columns.push_back(columns[i]); - } - } - - auto vids = insert_primary_keys(columns[pk_index]); - for (size_t i = 0; i < property_columns.size(); ++i) { - set_properties_from_context_column(table_->get_column_by_id(i), - property_columns[i], vids); - } - } - } - std::shared_ptr ckp_; std::unique_ptr indexer_; std::unique_ptr
table_; diff --git a/include/neug/storages/loader/chunk_pipeline_utils.h b/include/neug/storages/loader/chunk_pipeline_utils.h index fc80d129d..6bdb5fcbe 100644 --- a/include/neug/storages/loader/chunk_pipeline_utils.h +++ b/include/neug/storages/loader/chunk_pipeline_utils.h @@ -22,7 +22,6 @@ #include #include #include -#include #include #include #include @@ -33,56 +32,13 @@ namespace neug { -struct ChunkPipelineOptions { - int32_t consumer_count = 1; - size_t queue_capacity = 2; -}; - -struct ChunkPipelineAllocation { - bool parallel_enabled = false; - int32_t producer_count = 1; - int32_t consumer_count = 1; - size_t queue_capacity = 2; -}; +namespace chunk_pipeline_detail { inline int32_t hardware_worker_count() { auto workers = static_cast(std::thread::hardware_concurrency()); return workers <= 0 ? 1 : workers; } -inline ChunkPipelineAllocation resolve_chunk_pipeline_allocation( - int64_t source_bytes, bool parallel_enabled, bool preserve_order, - int32_t hardware_workers = 0) { - constexpr int64_t kMinParallelBytes = 256LL * 1024 * 1024; - constexpr int64_t kMinPartitionBytes = 64LL * 1024 * 1024; - constexpr size_t kMaxQueuedChunks = 64; - - ChunkPipelineAllocation result; - const auto workers = - hardware_workers > 0 ? hardware_workers : hardware_worker_count(); - if (!parallel_enabled || preserve_order || workers <= 1 || - source_bytes < kMinParallelBytes) { - return result; - } - - // Compute ceil(source_bytes / kMinPartitionBytes) without overflowing when - // source_bytes is close to INT64_MAX. - const auto useful_partitions = std::max( - 1, source_bytes / kMinPartitionBytes + - (source_bytes % kMinPartitionBytes == 0 ? 0 : 1)); - const auto balanced_producers = (workers + 1) / 2; - result.producer_count = static_cast(std::min( - balanced_producers, std::min(useful_partitions, workers - 1))); - result.producer_count = std::max(1, result.producer_count); - result.consumer_count = std::max(1, workers - result.producer_count); - result.parallel_enabled = true; - result.queue_capacity = std::clamp( - static_cast(result.producer_count) * 2, 2, kMaxQueuedChunks); - return result; -} - -namespace chunk_pipeline_detail { - template class BoundedQueue { public: @@ -130,16 +86,15 @@ class BoundedQueue { bool closed_ = false; }; -} // namespace chunk_pipeline_detail - /// Reads chunks from one non-thread-safe supplier and delivers them to a /// bounded pool of consumers. It preserves the first exception, closes the /// queue on cancellation, and joins every thread before rethrowing. template inline void consume_chunk_pipeline_impl(IDataChunkSupplier& supplier, - ChunkPipelineOptions options, + int32_t consumer_count, + size_t queue_capacity, Consume&& consume) { - const auto consumer_count = std::max(1, options.consumer_count); + consumer_count = std::max(1, consumer_count); if (consumer_count == 1) { while (auto chunk = supplier.GetNextChunk()) { consume(0, chunk); @@ -148,7 +103,7 @@ inline void consume_chunk_pipeline_impl(IDataChunkSupplier& supplier, } chunk_pipeline_detail::BoundedQueue> queue( - options.queue_capacity); + queue_capacity); std::atomic cancelled{false}; std::mutex error_mutex; std::exception_ptr first_error; @@ -203,30 +158,13 @@ inline void consume_chunk_pipeline_impl(IDataChunkSupplier& supplier, } } -inline void consume_chunk_pipeline( - IDataChunkSupplier& supplier, ChunkPipelineOptions options, - const std::function&)>& consume) { - consume_chunk_pipeline_impl( - supplier, options, - [&](int32_t /*consumer*/, const std::shared_ptr& chunk) { - consume(chunk); - }); -} - -inline void consume_chunk_pipeline_indexed( - IDataChunkSupplier& supplier, ChunkPipelineOptions options, - const std::function&)>& - consume) { - consume_chunk_pipeline_impl(supplier, options, consume); -} - /// Concurrently pulls chunks from a supplier that owns its producer queue. /// This avoids adding another producer thread and bounded queue between a /// partitioned source and its consumers. template -inline void consume_concurrent_supplier_indexed(IDataChunkSupplier& supplier, - int32_t consumer_count, - Consume&& consume) { +inline void consume_concurrent_supplier_impl(IDataChunkSupplier& supplier, + int32_t consumer_count, + Consume&& consume) { CHECK(supplier.SupportsConcurrentGetNext()); consumer_count = std::max(1, consumer_count); if (consumer_count == 1) { @@ -274,4 +212,23 @@ inline void consume_concurrent_supplier_indexed(IDataChunkSupplier& supplier, } } +} // namespace chunk_pipeline_detail + +/// Delivers chunks to indexed consumers using the cheapest path supported by +/// the supplier: direct concurrent pulls, serial iteration, or one producer +/// feeding a bounded consumer queue. +template +inline void consume_supplier_indexed(IDataChunkSupplier& supplier, + const ChunkSourceOptions& options, + Consume&& consume) { + if (supplier.SupportsConcurrentGetNext()) { + chunk_pipeline_detail::consume_concurrent_supplier_impl( + supplier, options.consumer_count, std::forward(consume)); + return; + } + chunk_pipeline_detail::consume_chunk_pipeline_impl( + supplier, options.consumer_count, options.queue_capacity, + std::forward(consume)); +} + } // namespace neug diff --git a/include/neug/storages/loader/loader_utils.h b/include/neug/storages/loader/loader_utils.h index 436fef4c1..d83acb83c 100644 --- a/include/neug/storages/loader/loader_utils.h +++ b/include/neug/storages/loader/loader_utils.h @@ -68,6 +68,8 @@ CsvReadConfig build_csv_read_config( const std::unordered_map& csv_options, const std::vector& column_types); +class IDataChunkSource; + class IDataChunkSupplier { public: virtual ~IDataChunkSupplier() = default; @@ -79,11 +81,19 @@ class IDataChunkSupplier { /// Stops any background producers and wakes blocked GetNextChunk() calls. virtual void Cancel() {} + + /// Returns the repeatable source backing this supplier, when available. + /// Storage uses this hint to select staged bulk build internally. + virtual std::shared_ptr RepeatableSource() const { + return nullptr; + } }; struct ChunkSourceOptions { - bool parallel_enabled = false; - int32_t producer_count = 1; + /// Zero selects the source's serial supplier. Positive values request that + /// many background parsing workers when input order need not be preserved. + int32_t producer_count = 0; + int32_t consumer_count = 1; size_t queue_capacity = 2; bool preserve_order = true; @@ -100,17 +110,10 @@ class IDataChunkSource { public: virtual ~IDataChunkSource() = default; - /// Opens a new supplier positioned at the beginning of the source. - virtual std::shared_ptr Open() const = 0; - - /// Opens a supplier with execution-specific concurrency settings. Sources - /// that do not implement parsing-time projection receive a generic - /// projection wrapper around Open(). + /// Opens a new supplier positioned at the beginning of the source using the + /// requested projection and concurrency settings. virtual std::shared_ptr Open( - const ChunkSourceOptions& options) const; - - /// Whether Open() can be called more than once with identical contents. - virtual bool rewindable() const = 0; + const ChunkSourceOptions& options = {}) const = 0; /// Returns the source size when cheaply known, otherwise -1. virtual int64_t EstimatedBytes() const { return -1; } @@ -119,6 +122,9 @@ class IDataChunkSource { virtual bool ParallelEnabled() const { return true; } }; +std::shared_ptr make_data_chunk_supplier( + std::shared_ptr source); + inline constexpr int64_t kUnknownRowNum = -1; enum class CsvRowCountMode { @@ -151,24 +157,21 @@ struct CsvPartitionPlanCache; /// raw row-count scan. class CSVChunkSource final : public IDataChunkSource { public: - CSVChunkSource(std::vector file_paths, CsvReadConfig config); + CSVChunkSource(std::vector file_paths, CsvReadConfig config, + std::vector projected_columns = {}); - std::shared_ptr Open() const override; std::shared_ptr Open( - const ChunkSourceOptions& options) const override; - bool rewindable() const override { return true; } + const ChunkSourceOptions& options = {}) const override; int64_t EstimatedBytes() const override; bool ParallelEnabled() const override { return config_.use_threads; } private: std::vector file_paths_; CsvReadConfig config_; + std::vector projected_columns_; std::shared_ptr partition_plan_cache_; }; -using CSVStreamChunkSupplier = CSVChunkSupplier; -using CSVTableChunkSupplier = CSVChunkSupplier; - void fillVertexReaderMeta(label_t v_label, const std::string& v_label_name, const std::string& v_file, const LoadingConfig& loading_config, diff --git a/include/neug/utils/io/read/csv/csv_reader.h b/include/neug/utils/io/read/csv/csv_reader.h index bbbeec875..71bede42f 100644 --- a/include/neug/utils/io/read/csv/csv_reader.h +++ b/include/neug/utils/io/read/csv/csv_reader.h @@ -46,7 +46,8 @@ class CsvReader { /// Creates a repeatable CSV source for direct COPY FROM bulk loading. /// Returns nullptr when the read needs a row filter and must materialize. - std::shared_ptr createChunkSource(); + std::shared_ptr createChunkSource( + std::vector projected_columns = {}); result> inferSchema(); diff --git a/src/execution/execute/ops/batch/batch_insert_edge.cc b/src/execution/execute/ops/batch/batch_insert_edge.cc index a652f2c0f..826dde914 100644 --- a/src/execution/execute/ops/batch/batch_insert_edge.cc +++ b/src/execution/execute/ops/batch/batch_insert_edge.cc @@ -15,15 +15,14 @@ #include "neug/execution/execute/ops/batch/batch_insert_edge.h" #include "neug/compiler/function/read_function.h" -#include "neug/compiler/main/metadata_registry.h" #include "neug/execution/common/context.h" #include "neug/execution/execute/ops/batch/batch_update_utils.h" -#include "neug/execution/execute/ops/batch/data_source.h" #include "neug/storages/graph/graph_interface.h" #include "neug/utils/exception/exception.h" #include "neug/utils/result.h" #include +#include #include #include @@ -38,27 +37,6 @@ namespace ops { namespace { -bool resolve_vertex_label_id(const Schema& schema, const ::common::NameOrId& ni, - label_t& out) { - switch (ni.item_case()) { - case ::common::NameOrId::kId: { - out = ni.id(); - return true; - } - case ::common::NameOrId::kName: { - if (!schema.is_vertex_label_valid(ni.name())) { - LOG(ERROR) << "Unknown vertex type: " << ni.DebugString(); - return false; - } - out = schema.get_vertex_label_id(ni.name()); - return true; - } - default: - LOG(ERROR) << "Unknown vertex type: " << ni.DebugString(); - return false; - } -} - /** Resolve edge + src/dst vertex labels from schema at execution time. */ bool resolve_edge_triplet(const Schema& schema, const physical::EdgeType& edge_type, @@ -113,56 +91,29 @@ class BatchInsertEdgeOpr : public IOperator { public: BatchInsertEdgeOpr( physical::EdgeType edge_type, - std::vector> prop_mappings, - std::vector> src_vertex_bindings, - std::vector> dst_vertex_bindings) - : edge_type_(std::move(edge_type)), - prop_mappings_(std::move(prop_mappings)), - src_vertex_bindings_(std::move(src_vertex_bindings)), - dst_vertex_bindings_(std::move(dst_vertex_bindings)) {} - - std::string get_operator_name() const override { - return "BatchInsertEdgeOpr"; - } - - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; - - private: - physical::EdgeType edge_type_; - std::vector> prop_mappings_, - src_vertex_bindings_, dst_vertex_bindings_; -}; - -class BatchInsertEdgeFromSourceOpr : public IOperator { - public: - BatchInsertEdgeFromSourceOpr( - std::shared_ptr shared_state, - function::ReadFunction* read_function, physical::EdgeType edge_type, std::vector> property_mappings, std::vector> source_mappings, - std::vector> destination_mappings) - : shared_state_(std::move(shared_state)), - read_function_(read_function), - edge_type_(std::move(edge_type)), + std::vector> destination_mappings, + std::optional source = std::nullopt) + : edge_type_(std::move(edge_type)), property_mappings_(std::move(property_mappings)), source_mappings_(std::move(source_mappings)), - destination_mappings_(std::move(destination_mappings)) {} + destination_mappings_(std::move(destination_mappings)), + source_(std::move(source)) {} std::string get_operator_name() const override { - return "BatchInsertEdgeFromSourceOpr"; + return "BatchInsertEdgeOpr"; } neug::result Eval(IStorageInterface& graph, const ParamsMap& params, Context&& ctx, OprTimer* timer) override; private: - std::shared_ptr shared_state_; - function::ReadFunction* read_function_; physical::EdgeType edge_type_; std::vector> property_mappings_; std::vector> source_mappings_; std::vector> destination_mappings_; + std::optional source_; }; neug::result BatchInsertEdgeOpr::Eval( @@ -181,55 +132,19 @@ neug::result BatchInsertEdgeOpr::Eval( "BatchInsertEdge"); } - auto total_mappings = build_total_edge_mappings( - src_vertex_bindings_, dst_vertex_bindings_, prop_mappings_); - auto supplier = create_data_chunk_supplier(ctx, total_mappings); - - RETURN_STATUS_ERROR_IF_NOT_OK( - graph.BatchAddEdges(src_label_id, dst_label_id, edge_label_id, supplier)); - return neug::result(std::move(ctx)); -} - -neug::result BatchInsertEdgeFromSourceOpr::Eval( - IStorageInterface& graph_interface, const ParamsMap& params, Context&& ctx, - OprTimer* timer) { - (void) params; - (void) ctx; - (void) timer; - CHECK(read_function_ != nullptr); - auto& graph = dynamic_cast(graph_interface); - label_t edge_label_id = 0; - label_t src_label_id = 0; - label_t dst_label_id = 0; - if (!resolve_edge_triplet(graph.schema(), edge_type_, edge_label_id, - src_label_id, dst_label_id)) { - RETURN_STATUS_ERROR(StatusCode::ERR_INVALID_ARGUMENT, - "Failed to resolve edge type for " - "BatchInsertEdgeFromSource"); - } - auto mappings = build_total_edge_mappings( source_mappings_, destination_mappings_, property_mappings_); - auto* ap_graph = dynamic_cast(&graph_interface); - if (ap_graph && - ap_graph->CanBatchBuildEdges(src_label_id, dst_label_id, edge_label_id) && - read_function_->sourceFunc) { - auto raw_source = read_function_->sourceFunc(shared_state_); - auto source = create_data_chunk_source(std::move(raw_source), mappings); - if (source && should_use_copy_bulk_build(*source)) { - RETURN_STATUS_ERROR_IF_NOT_OK(ap_graph->BatchBuildEdges( - src_label_id, dst_label_id, edge_label_id, std::move(source))); - return Context{}; - } + BatchInsertInput input; + if (source_) { + input = create_batch_insert_input(source_->state, *source_->read_function, + mappings); + } else { + input.supplier = create_data_chunk_supplier(ctx, mappings); + input.output = std::move(ctx); } - - auto materialized = read_function_->execFunc(shared_state_); - auto supplier = create_data_chunk_supplier(materialized, mappings); RETURN_STATUS_ERROR_IF_NOT_OK(graph.BatchAddEdges( - src_label_id, dst_label_id, edge_label_id, std::move(supplier))); - // Match the empty terminal sink consumed by the fused plan. - materialized.tag_ids.clear(); - return neug::result(std::move(materialized)); + src_label_id, dst_label_id, edge_label_id, std::move(input.supplier))); + return neug::result(std::move(input.output)); } neug::result BatchInsertEdgeOprBuilder::Build( @@ -264,22 +179,16 @@ neug::result BatchInsertEdgeFromSourceOprBuilder::Build( const physical::PhysicalPlan& plan, int op_idx) { (void) schema; ContextMeta result_meta = ctx_meta; - if (op_idx + 3 != plan.plan_size() || - plan.plan(op_idx + 2).opr().sink().tags_size() != 0) { + if (!is_terminal_batch_insert(plan, op_idx)) { return std::make_pair(nullptr, result_meta); } - const auto& source_pb = plan.plan(op_idx).opr().source(); const auto& edge_pb = plan.plan(op_idx + 1).opr().load_edge(); if (!edge_pb.has_edge_type()) { THROW_INTERNAL_EXCEPTION( "BatchInsertEdgeFromSourceOprBuilder: edge type is not set"); } - ReadStateBuilder state_builder; - auto state = state_builder.build(source_pb); - auto catalog = neug::main::MetadataRegistry::getCatalog(); - auto function = catalog->getFunctionWithSignature(source_pb.extension_name()); - auto read_function = function->ptrCast(); + auto source = build_batch_insert_source(plan, op_idx); std::vector> property_mappings; std::vector> source_mappings; @@ -290,12 +199,11 @@ neug::result BatchInsertEdgeFromSourceOprBuilder::Build( destination_mappings); physical::EdgeType edge_type; edge_type.CopyFrom(edge_pb.edge_type()); - return std::make_pair( - std::make_unique( - std::move(state), read_function, std::move(edge_type), - std::move(property_mappings), std::move(source_mappings), - std::move(destination_mappings)), - result_meta); + return std::make_pair(std::make_unique( + std::move(edge_type), std::move(property_mappings), + std::move(source_mappings), + std::move(destination_mappings), std::move(source)), + result_meta); } } // namespace ops diff --git a/src/execution/execute/ops/batch/batch_insert_vertex.cc b/src/execution/execute/ops/batch/batch_insert_vertex.cc index bc2284714..d5024dd0c 100644 --- a/src/execution/execute/ops/batch/batch_insert_vertex.cc +++ b/src/execution/execute/ops/batch/batch_insert_vertex.cc @@ -15,15 +15,13 @@ #include "neug/execution/execute/ops/batch/batch_insert_vertex.h" #include "neug/compiler/function/read_function.h" -#include "neug/compiler/main/metadata_registry.h" #include "neug/execution/common/context.h" #include "neug/execution/execute/ops/batch/batch_update_utils.h" -#include "neug/execution/execute/ops/batch/data_source.h" #include "neug/storages/graph/graph_interface.h" #include "neug/utils/exception/exception.h" #include "neug/utils/result.h" -#include +#include #include #include @@ -34,37 +32,15 @@ class OprTimer; namespace ops { -namespace { - -bool resolve_vertex_label_id(const Schema& schema, - const ::common::NameOrId& type, - label_t& label_id) { - switch (type.item_case()) { - case ::common::NameOrId::kId: - label_id = type.id(); - return true; - case ::common::NameOrId::kName: - if (!schema.is_vertex_label_valid(type.name())) { - LOG(ERROR) << "Unknown vertex type: " << type.DebugString(); - return false; - } - label_id = schema.get_vertex_label_id(type.name()); - return true; - default: - LOG(ERROR) << "Invalid vertex type: " << type.DebugString(); - return false; - } -} - -} // namespace - class BatchInsertVertexOpr : public IOperator { public: BatchInsertVertexOpr( ::common::NameOrId vertex_type, - std::vector> prop_mappings) + std::vector> property_mappings, + std::optional source = std::nullopt) : vertex_type_(std::move(vertex_type)), - prop_mappings_(std::move(prop_mappings)) {} + property_mappings_(std::move(property_mappings)), + source_(std::move(source)) {} std::string get_operator_name() const override { return "BatchInsertVertexOpr"; @@ -74,33 +50,9 @@ class BatchInsertVertexOpr : public IOperator { Context&& ctx, OprTimer* timer) override; private: - ::common::NameOrId vertex_type_; - std::vector> prop_mappings_; -}; - -class BatchInsertVertexFromSourceOpr : public IOperator { - public: - BatchInsertVertexFromSourceOpr( - std::shared_ptr shared_state, - function::ReadFunction* read_function, ::common::NameOrId vertex_type, - std::vector> property_mappings) - : shared_state_(std::move(shared_state)), - read_function_(read_function), - vertex_type_(std::move(vertex_type)), - property_mappings_(std::move(property_mappings)) {} - - std::string get_operator_name() const override { - return "BatchInsertVertexFromSourceOpr"; - } - - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; - - private: - std::shared_ptr shared_state_; - function::ReadFunction* read_function_; ::common::NameOrId vertex_type_; std::vector> property_mappings_; + std::optional source_; }; neug::result BatchInsertVertexOpr::Eval( @@ -114,48 +66,17 @@ neug::result BatchInsertVertexOpr::Eval( RETURN_STATUS_ERROR(StatusCode::ERR_INVALID_ARGUMENT, "Failed to resolve vertex type for BatchInsertVertex"); } - auto supplier = create_data_chunk_supplier(ctx, prop_mappings_); - RETURN_STATUS_ERROR_IF_NOT_OK( - graph.BatchAddVertices(vertex_label_id, supplier)); - return neug::result(std::move(ctx)); -} - -neug::result BatchInsertVertexFromSourceOpr::Eval( - IStorageInterface& graph_interface, const ParamsMap& params, Context&& ctx, - OprTimer* timer) { - (void) params; - (void) ctx; - (void) timer; - CHECK(read_function_ != nullptr); - auto& graph = dynamic_cast(graph_interface); - label_t vertex_label_id = 0; - if (!resolve_vertex_label_id(graph.schema(), vertex_type_, vertex_label_id)) { - RETURN_STATUS_ERROR( - StatusCode::ERR_INVALID_ARGUMENT, - "Failed to resolve vertex type for BatchInsertVertexFromSource"); + BatchInsertInput input; + if (source_) { + input = create_batch_insert_input(source_->state, *source_->read_function, + property_mappings_); + } else { + input.supplier = create_data_chunk_supplier(ctx, property_mappings_); + input.output = std::move(ctx); } - - auto* ap_graph = dynamic_cast(&graph_interface); - if (ap_graph && ap_graph->CanBatchBuildVertices(vertex_label_id) && - read_function_->sourceFunc) { - auto raw_source = read_function_->sourceFunc(shared_state_); - auto source = - create_data_chunk_source(std::move(raw_source), property_mappings_); - if (source && should_use_copy_bulk_build(*source)) { - RETURN_STATUS_ERROR_IF_NOT_OK( - ap_graph->BatchBuildVertices(vertex_label_id, std::move(source))); - // The fused builder only accepts a terminal COPY with an empty sink. - return Context{}; - } - } - - auto materialized = read_function_->execFunc(shared_state_); - auto supplier = create_data_chunk_supplier(materialized, property_mappings_); RETURN_STATUS_ERROR_IF_NOT_OK( - graph.BatchAddVertices(vertex_label_id, std::move(supplier))); - // Match the empty terminal sink consumed by the fused plan. - materialized.tag_ids.clear(); - return neug::result(std::move(materialized)); + graph.BatchAddVertices(vertex_label_id, std::move(input.supplier))); + return neug::result(std::move(input.output)); } neug::result BatchInsertVertexOprBuilder::Build( @@ -183,32 +104,25 @@ neug::result BatchInsertVertexFromSourceOprBuilder::Build( const physical::PhysicalPlan& plan, int op_idx) { (void) schema; ContextMeta result_meta = ctx_meta; - if (op_idx + 3 != plan.plan_size() || - plan.plan(op_idx + 2).opr().sink().tags_size() != 0) { + if (!is_terminal_batch_insert(plan, op_idx)) { return std::make_pair(nullptr, result_meta); } - const auto& source_pb = plan.plan(op_idx).opr().source(); const auto& vertex_pb = plan.plan(op_idx + 1).opr().load_vertex(); if (!vertex_pb.has_vertex_type()) { THROW_INTERNAL_EXCEPTION( "BatchInsertVertexFromSourceOprBuilder: vertex type is not set"); } - ReadStateBuilder state_builder; - auto state = state_builder.build(source_pb); - auto catalog = neug::main::MetadataRegistry::getCatalog(); - auto function = catalog->getFunctionWithSignature(source_pb.extension_name()); - auto read_function = function->ptrCast(); + auto source = build_batch_insert_source(plan, op_idx); std::vector> property_mappings; parse_property_mappings(vertex_pb.property_mappings(), property_mappings); ::common::NameOrId vertex_type; vertex_type.CopyFrom(vertex_pb.vertex_type()); - return std::make_pair( - std::make_unique( - std::move(state), read_function, std::move(vertex_type), - std::move(property_mappings)), - result_meta); + return std::make_pair(std::make_unique( + std::move(vertex_type), + std::move(property_mappings), std::move(source)), + result_meta); } } // namespace ops diff --git a/src/execution/execute/ops/batch/batch_update_utils.cc b/src/execution/execute/ops/batch/batch_update_utils.cc index fc765f408..330fd5f6e 100644 --- a/src/execution/execute/ops/batch/batch_update_utils.cc +++ b/src/execution/execute/ops/batch/batch_update_utils.cc @@ -23,9 +23,7 @@ #include #include -#include #include -#include #include #include #include @@ -33,7 +31,10 @@ #include "neug/common/types/i_context_column.h" #include "neug/common/types/value.h" +#include "neug/compiler/function/read_function.h" +#include "neug/compiler/main/metadata_registry.h" #include "neug/execution/common/context.h" +#include "neug/execution/execute/ops/batch/data_source.h" #include "neug/storages/graph/graph_interface.h" #include "neug/storages/loader/loader_utils.h" #include "neug/utils/string_utils.h" @@ -319,108 +320,6 @@ class MultiChunkSupplier : public IDataChunkSupplier { size_t index_; }; -class ProjectingChunkSupplier final : public IDataChunkSupplier { - public: - ProjectingChunkSupplier(std::shared_ptr input, - std::vector aliases) - : input_(std::move(input)), aliases_(std::move(aliases)) {} - - std::shared_ptr GetNextChunk() override { - auto input = input_->GetNextChunk(); - if (!input) { - return nullptr; - } - auto output = std::make_shared(); - for (size_t index = 0; index < aliases_.size(); ++index) { - auto column = input->get(aliases_[index]); - if (!column) { - THROW_INTERNAL_EXCEPTION("Column not found for tag id: " + - std::to_string(aliases_[index])); - } - output->set(static_cast(index), std::move(column)); - } - return output; - } - - int64_t RowNum() const override { return input_->RowNum(); } - - bool SupportsConcurrentGetNext() const override { - return input_->SupportsConcurrentGetNext(); - } - - void Cancel() override { input_->Cancel(); } - - private: - std::shared_ptr input_; - std::vector aliases_; -}; - -class ProjectingChunkSource final : public IDataChunkSource { - public: - ProjectingChunkSource(std::shared_ptr input, - std::vector aliases) - : input_(std::move(input)), aliases_(std::move(aliases)) {} - - std::shared_ptr Open() const override { - auto supplier = input_->Open(); - if (!supplier) { - return nullptr; - } - return std::make_shared(std::move(supplier), - aliases_); - } - - std::shared_ptr Open( - const ChunkSourceOptions& options) const override { - if (!options.projected_columns.empty()) { - ChunkSourceOptions input_options = options; - input_options.projected_columns.clear(); - input_options.projected_columns.reserve(options.projected_columns.size()); - for (const auto output_column : options.projected_columns) { - if (output_column < 0 || - static_cast(output_column) >= aliases_.size()) { - THROW_INVALID_ARGUMENT_EXCEPTION( - "Projected source column is out of range: " + - std::to_string(output_column)); - } - input_options.projected_columns.push_back( - aliases_[static_cast(output_column)]); - } - // IDataChunkSource::Open(options) guarantees that projected columns are - // returned in the requested order. CSV sources push this into typed - // parsing; other repeatable sources receive the generic wrapper. - return input_->Open(input_options); - } - - auto supplier = input_->Open(options); - if (!supplier) { - return nullptr; - } - return std::make_shared(std::move(supplier), - aliases_); - } - - bool rewindable() const override { return input_->rewindable(); } - - int64_t EstimatedBytes() const override { return input_->EstimatedBytes(); } - - bool ParallelEnabled() const override { return input_->ParallelEnabled(); } - - private: - std::shared_ptr input_; - std::vector aliases_; -}; - -std::vector property_mapping_aliases( - const std::vector>& prop_mappings) { - std::vector aliases; - aliases.reserve(prop_mappings.size()); - for (const auto& mapping : prop_mappings) { - aliases.push_back(mapping.first); - } - return aliases; -} - std::shared_ptr create_data_chunk_supplier( const Context& ctx, const std::vector>& prop_mappings) { @@ -443,33 +342,65 @@ std::shared_ptr create_data_chunk_supplier( return std::make_shared(std::move(projected_chunks)); } -std::shared_ptr create_data_chunk_source( - std::shared_ptr source, - const std::vector>& prop_mappings) { - if (!source) { - return nullptr; +bool resolve_vertex_label_id(const Schema& schema, + const ::common::NameOrId& type, + label_t& label_id) { + switch (type.item_case()) { + case ::common::NameOrId::kId: + label_id = type.id(); + return true; + case ::common::NameOrId::kName: + if (!schema.is_vertex_label_valid(type.name())) { + LOG(ERROR) << "Unknown vertex type: " << type.DebugString(); + return false; + } + label_id = schema.get_vertex_label_id(type.name()); + return true; + default: + LOG(ERROR) << "Invalid vertex type: " << type.DebugString(); + return false; } - return std::make_shared( - std::move(source), property_mapping_aliases(prop_mappings)); } -bool should_use_copy_bulk_build(const IDataChunkSource& source) { - const char* configured = std::getenv("NEUG_COPY_BULK_BUILD"); - if (configured != nullptr) { - std::string value(configured); - std::transform(value.begin(), value.end(), value.begin(), - [](unsigned char ch) { return std::tolower(ch); }); - if (value == "0" || value == "false" || value == "off" || value == "no") { - return false; +bool is_terminal_batch_insert(const physical::PhysicalPlan& plan, int op_idx) { + return op_idx + 3 == plan.plan_size() && + plan.plan(op_idx + 2).opr().sink().tags_size() == 0; +} + +BatchInsertSource build_batch_insert_source(const physical::PhysicalPlan& plan, + int op_idx) { + const auto& source = plan.plan(op_idx).opr().source(); + ReadStateBuilder state_builder; + auto state = state_builder.build(source); + auto catalog = neug::main::MetadataRegistry::getCatalog(); + auto registered_function = + catalog->getFunctionWithSignature(source.extension_name()); + return {std::move(state), + registered_function->ptrCast()}; +} + +BatchInsertInput create_batch_insert_input( + const std::shared_ptr& shared_state, + const function::ReadFunction& read_function, + const std::vector>& prop_mappings) { + if (read_function.sourceFunc) { + std::vector projected_columns; + projected_columns.reserve(prop_mappings.size()); + for (const auto& mapping : prop_mappings) { + projected_columns.push_back(mapping.first); } - if (value == "1" || value == "true" || value == "on" || value == "yes") { - return source.rewindable(); + auto source = + read_function.sourceFunc(shared_state, std::move(projected_columns)); + if (source) { + return {make_data_chunk_supplier(std::move(source)), Context{}}; } - LOG(WARNING) << "Ignore invalid NEUG_COPY_BULK_BUILD=" << configured; } - constexpr int64_t kMinBulkBuildBytes = 256LL * 1024 * 1024; - return source.rewindable() && source.EstimatedBytes() >= kMinBulkBuildBytes; + CHECK(read_function.execFunc != nullptr); + auto output = read_function.execFunc(shared_state); + auto supplier = create_data_chunk_supplier(output, prop_mappings); + output.tag_ids.clear(); + return {std::move(supplier), std::move(output)}; } std::vector match_files_with_pattern( diff --git a/src/storages/csr/mutable_csr.cc b/src/storages/csr/mutable_csr.cc index 5152fcf95..cb4d8609e 100644 --- a/src/storages/csr/mutable_csr.cc +++ b/src/storages/csr/mutable_csr.cc @@ -473,7 +473,7 @@ void MutableCsr::batch_put_edges(const std::vector& src_list, int old_deg = sz_arr[i].load(std::memory_order_relaxed); total_to_move += old_deg; int new_degree = degree[i] + old_deg; - int new_cap = std::ceil(new_degree * NeugDBConfig::DEFAULT_RESERVE_RATIO); + int new_cap = reserved_capacity(new_degree); cap_arr[i] = new_cap; total_to_allocate += new_cap; } diff --git a/src/storages/graph/edge_table.cc b/src/storages/graph/edge_table.cc index 4045aff43..b9ae0052d 100644 --- a/src/storages/graph/edge_table.cc +++ b/src/storages/graph/edge_table.cc @@ -20,17 +20,18 @@ #include "neug/storages/module/module_factory.h" #include -#include #include #include -#include +#include #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -53,33 +54,52 @@ namespace neug { namespace { -size_t vector_bool_storage_bytes(size_t bit_capacity) { - return (bit_capacity + 7) / 8; +bool should_use_bulk_edge_build(const IDataChunkSource& source) { + const char* configured = std::getenv("NEUG_COPY_BULK_BUILD"); + if (configured != nullptr) { + std::string value(configured); + std::transform(value.begin(), value.end(), value.begin(), + [](unsigned char ch) { return std::tolower(ch); }); + if (value == "0" || value == "false" || value == "off" || value == "no") { + return false; + } + if (value == "1" || value == "true" || value == "on" || value == "yes") { + return true; + } + LOG(WARNING) << "Ignore invalid NEUG_COPY_BULK_BUILD=" << configured; + } + + constexpr int64_t kMinBulkBuildBytes = 256LL * 1024 * 1024; + return source.EstimatedBytes() >= kMinBulkBuildBytes; } -uint64_t process_peak_rss_bytes() { - struct rusage usage {}; - if (getrusage(RUSAGE_SELF, &usage) != 0) { - return 0; +ChunkSourceOptions resolve_bulk_edge_source_options(int64_t source_bytes, + bool parallel_enabled, + bool preserve_order) { + constexpr int64_t kMinParallelBytes = 256LL * 1024 * 1024; + constexpr int64_t kMinPartitionBytes = 64LL * 1024 * 1024; + constexpr size_t kMaxQueuedChunks = 64; + + ChunkSourceOptions options; + options.preserve_order = preserve_order; + const auto workers = chunk_pipeline_detail::hardware_worker_count(); + if (!parallel_enabled || preserve_order || workers <= 1 || + source_bytes < kMinParallelBytes) { + return options; } -#if defined(__APPLE__) - return static_cast(usage.ru_maxrss); -#else - return static_cast(usage.ru_maxrss) * 1024; -#endif -} -size_t estimate_edge_fallback_buffer_bytes( - const std::vector& src_lid, const std::vector& dst_lid, - const std::vector& valid_flags, - const std::vector>& bundled_data_cols, - const std::vector>& unbundled_data_chunks) { - return src_lid.capacity() * sizeof(vid_t) + - dst_lid.capacity() * sizeof(vid_t) + - vector_bool_storage_bytes(valid_flags.capacity()) + - bundled_data_cols.capacity() * - sizeof(std::shared_ptr) + - unbundled_data_chunks.capacity() * sizeof(std::shared_ptr); + const auto useful_partitions = std::max( + 1, source_bytes / kMinPartitionBytes + + (source_bytes % kMinPartitionBytes == 0 ? 0 : 1)); + const auto balanced_producers = (workers + 1) / 2; + options.producer_count = static_cast(std::min( + balanced_producers, std::min(useful_partitions, workers - 1))); + options.producer_count = std::max(1, options.producer_count); + options.consumer_count = + std::max(1, workers - options.producer_count); + options.queue_capacity = std::clamp( + static_cast(options.producer_count) * 2, 2, kMaxQueuedChunks); + return options; } } // namespace @@ -431,20 +451,13 @@ void batch_add_bundled_edges_impl( template class EmptyCsrBulkWriter { public: - static constexpr bool kSupportsDisjointConcurrentFill = true; static constexpr bool kStoresEdges = false; static constexpr bool kNeedsDegreeCount = false; static constexpr bool kChecksSingleUniqueness = false; - static constexpr bool kNeedsConcurrentGrouping = false; static constexpr bool kTracksInputEdgeCount = false; void PrepareBuild(vid_t /*vertex_count*/) {} void AllocateFromCounts() {} - int ReserveConcurrent(vid_t /*src*/, int /*count*/) { return 0; } - void PutAt(vid_t /*src*/, int /*slot*/, vid_t /*dst*/, - const EDATA_T& /*data*/, timestamp_t /*ts*/) {} - void PutSerial(vid_t /*src*/, vid_t /*dst*/, const EDATA_T& /*data*/, - timestamp_t /*ts*/) {} void Finish() {} }; @@ -501,15 +514,29 @@ class BulkEdgeDataReader { EmptyType Get(size_t /*row*/) const { return EmptyType(); } }; -struct BulkEdgeEndpointScratch { +constexpr uint32_t kInvalidBulkEdgeIndex = std::numeric_limits::max(); + +struct BulkEdgeGroup { + uint32_t count = 0; + uint32_t head = kInvalidBulkEdgeIndex; + uint32_t tail = kInvalidBulkEdgeIndex; +}; + +struct BulkEdgeWorkerScratch { std::vector src_lids; std::vector dst_lids; + flat_hash_map out_groups; + flat_hash_map in_groups; + std::vector out_next; + std::vector in_next; + bool out_single_duplicate = false; + bool in_single_duplicate = false; }; void index_bulk_edge_endpoints(const std::shared_ptr& chunk, const IndexerType& src_indexer, const IndexerType& dst_indexer, - BulkEdgeEndpointScratch& scratch) { + BulkEdgeWorkerScratch& scratch) { CHECK(chunk != nullptr); CHECK_GE(chunk->col_num(), 2); auto src_column = chunk->get(0); @@ -530,47 +557,11 @@ inline bool is_valid_bulk_edge(vid_t src, vid_t dst) { dst != std::numeric_limits::max(); } -size_t count_valid_bulk_edges(const BulkEdgeEndpointScratch& endpoints) { - size_t valid_edges = 0; - for (size_t row = 0; row < endpoints.src_lids.size(); ++row) { - valid_edges += - is_valid_bulk_edge(endpoints.src_lids[row], endpoints.dst_lids[row]); - } - return valid_edges; -} - -struct BulkEdgeCountScratch { - BulkEdgeEndpointScratch endpoints; - flat_hash_map out_counts; - flat_hash_map in_counts; - bool out_single_duplicate = false; - bool in_single_duplicate = false; - uint64_t valid_edges = 0; -}; - -struct BulkEdgeCountChunkResult { - size_t valid_edges = 0; - size_t out_single_checks = 0; - size_t in_single_checks = 0; -}; - -struct BulkEdgeCountProfile { - size_t out_updates = 0; - size_t in_updates = 0; -}; - struct BulkEdgeCountSummary { - uint64_t valid_edges = 0; bool out_single_duplicate = false; bool in_single_duplicate = false; }; -struct BulkEdgeCountConfig { - bool concurrent = false; - bool check_out_single = false; - bool check_in_single = false; -}; - constexpr size_t kBulkEdgeInitialGroupReserve = 4096; size_t bulk_edge_group_reserve(size_t rows) { @@ -580,49 +571,51 @@ size_t bulk_edge_group_reserve(size_t rows) { return std::min(rows, kBulkEdgeInitialGroupReserve); } +void count_bulk_edge_group(flat_hash_map& groups, + vid_t vertex) { + auto& group = groups[vertex]; + CHECK_LT(group.count, std::numeric_limits::max()); + ++group.count; +} + template -BulkEdgeCountChunkResult count_bulk_edge_chunk( - const BulkEdgeEndpointScratch& endpoints, OutWriter& out, InWriter& in, - BulkEdgeCountScratch& scratch, const BulkEdgeCountConfig& config) { - CHECK_LE(endpoints.src_lids.size(), +void count_bulk_edge_chunk(BulkEdgeWorkerScratch& scratch, OutWriter& out, + InWriter& in, bool concurrent) { + CHECK_LE(scratch.src_lids.size(), static_cast(std::numeric_limits::max())); - BulkEdgeCountChunkResult result; - if (config.concurrent) { + if (concurrent) { if constexpr (OutWriter::kNeedsDegreeCount) { - scratch.out_counts.clear(); - scratch.out_counts.reserve( - bulk_edge_group_reserve(endpoints.src_lids.size())); + scratch.out_groups.clear(); + scratch.out_groups.reserve( + bulk_edge_group_reserve(scratch.src_lids.size())); } if constexpr (InWriter::kNeedsDegreeCount) { - scratch.in_counts.clear(); - scratch.in_counts.reserve( - bulk_edge_group_reserve(endpoints.dst_lids.size())); + scratch.in_groups.clear(); + scratch.in_groups.reserve( + bulk_edge_group_reserve(scratch.dst_lids.size())); } } - for (size_t row = 0; row < endpoints.src_lids.size(); ++row) { - const auto src = endpoints.src_lids[row]; - const auto dst = endpoints.dst_lids[row]; + for (size_t row = 0; row < scratch.src_lids.size(); ++row) { + const auto src = scratch.src_lids[row]; + const auto dst = scratch.dst_lids[row]; if (!is_valid_bulk_edge(src, dst)) { continue; } - ++result.valid_edges; - if (config.concurrent) { + if (concurrent) { if constexpr (OutWriter::kNeedsDegreeCount) { - ++scratch.out_counts[src]; + count_bulk_edge_group(scratch.out_groups, src); } if constexpr (InWriter::kNeedsDegreeCount) { - ++scratch.in_counts[dst]; + count_bulk_edge_group(scratch.in_groups, dst); } if constexpr (OutWriter::kChecksSingleUniqueness) { - if (config.check_out_single && !scratch.out_single_duplicate) { - ++result.out_single_checks; + if (!scratch.out_single_duplicate) { scratch.out_single_duplicate = out.CheckUniqueConcurrent(src); } } if constexpr (InWriter::kChecksSingleUniqueness) { - if (config.check_in_single && !scratch.in_single_duplicate) { - ++result.in_single_checks; + if (!scratch.in_single_duplicate) { scratch.in_single_duplicate = in.CheckUniqueConcurrent(dst); } } @@ -634,240 +627,72 @@ BulkEdgeCountChunkResult count_bulk_edge_chunk( in.CountSerial(dst); } if constexpr (OutWriter::kChecksSingleUniqueness) { - if (config.check_out_single && !scratch.out_single_duplicate) { - ++result.out_single_checks; + if (!scratch.out_single_duplicate) { scratch.out_single_duplicate = out.CheckUniqueSerial(src); } } if constexpr (InWriter::kChecksSingleUniqueness) { - if (config.check_in_single && !scratch.in_single_duplicate) { - ++result.in_single_checks; + if (!scratch.in_single_duplicate) { scratch.in_single_duplicate = in.CheckUniqueSerial(dst); } } } } - if (config.concurrent) { + if (concurrent) { if constexpr (OutWriter::kNeedsDegreeCount) { - for (const auto& [src, count] : scratch.out_counts) { - CHECK_LE(count, static_cast(std::numeric_limits::max())); - out.CountConcurrent(src, static_cast(count)); + for (const auto& [src, group] : scratch.out_groups) { + CHECK_LE(group.count, + static_cast(std::numeric_limits::max())); + out.CountConcurrent(src, static_cast(group.count)); } } if constexpr (InWriter::kNeedsDegreeCount) { - for (const auto& [dst, count] : scratch.in_counts) { - CHECK_LE(count, static_cast(std::numeric_limits::max())); - in.CountConcurrent(dst, static_cast(count)); + for (const auto& [dst, group] : scratch.in_groups) { + CHECK_LE(group.count, + static_cast(std::numeric_limits::max())); + in.CountConcurrent(dst, static_cast(group.count)); } } } - return result; -} - -BulkEdgeCountProfile profile_bulk_edge_count( - const BulkEdgeCountScratch& scratch, - const BulkEdgeCountChunkResult& chunk_result, bool concurrent_count, - bool counts_out_degree, bool counts_in_degree) { - BulkEdgeCountProfile result; - if (counts_out_degree) { - result.out_updates = - concurrent_count ? scratch.out_counts.size() : chunk_result.valid_edges; - } - if (counts_in_degree) { - result.in_updates = - concurrent_count ? scratch.in_counts.size() : chunk_result.valid_edges; - } - return result; -} - -ChunkSourceOptions make_bulk_edge_source_options( - const ChunkPipelineAllocation& allocation, bool preserve_order) { - ChunkSourceOptions source_options; - source_options.parallel_enabled = allocation.parallel_enabled; - source_options.producer_count = allocation.producer_count; - source_options.queue_capacity = allocation.queue_capacity; - source_options.preserve_order = preserve_order; - return source_options; } -struct BulkEdgeCountWorkerStats { - size_t chunks = 0; - size_t edges = 0; - size_t out_updates = 0; - size_t in_updates = 0; - size_t out_single_checks = 0; - size_t in_single_checks = 0; - int64_t work_time_ns = 0; -}; - -using BulkEdgeCountChunk = std::function; +using BulkEdgeCountChunk = std::function; BulkEdgeCountSummary count_bulk_edges( const std::shared_ptr& source, const IndexerType& src_indexer, const IndexerType& dst_indexer, - const ChunkPipelineAllocation& allocation, bool counts_out_degree, - bool counts_in_degree, bool check_out_single, bool check_in_single, + ChunkSourceOptions options, std::vector& scratches, const BulkEdgeCountChunk& count_chunk) { // Degree accumulation and single-slot uniqueness checks are commutative, so // this pass never needs input order even when fill may later fall back to the // ordered last-write-wins path. - auto source_options = make_bulk_edge_source_options(allocation, false); - // The degree pass only needs endpoint OIDs. Push this projection through - // ProjectingChunkSource into CSVChunkSource so edge properties are not typed, - // allocated, and discarded during the first parse. - source_options.projected_columns = {0, 1}; - auto supplier = source->Open(source_options); + // The degree pass only needs endpoint OIDs, so edge properties are not + // parsed, typed, allocated, and discarded during the first pass. + options.projected_columns = {0, 1}; + auto supplier = source->Open(options); CHECK(supplier != nullptr); - const bool profile_stages = VLOG_IS_ON(1); - std::chrono::steady_clock::time_point row_count_start; - std::chrono::steady_clock::time_point row_count_end; - int64_t row_num = 0; - if (profile_stages) { - row_count_start = std::chrono::steady_clock::now(); - row_num = supplier->RowNum(); - row_count_end = std::chrono::steady_clock::now(); - } - const auto worker_count = allocation.consumer_count; - std::vector scratches( - static_cast(worker_count)); - std::vector worker_stats; - if (profile_stages) { - worker_stats.resize(static_cast(worker_count)); - } - const BulkEdgeCountConfig count_config{ - .concurrent = worker_count > 1, - .check_out_single = check_out_single, - .check_in_single = check_in_single, - }; + const auto worker_count = options.consumer_count; + CHECK_EQ(scratches.size(), static_cast(worker_count)); + const bool concurrent = worker_count > 1; auto count = [&](int32_t worker, const std::shared_ptr& chunk) { CHECK_GE(worker, 0); CHECK_LT(worker, worker_count); - const auto start = profile_stages ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; auto& scratch = scratches[static_cast(worker)]; - index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, - scratch.endpoints); - const auto chunk_result = count_chunk(scratch, count_config); - scratch.valid_edges += static_cast(chunk_result.valid_edges); - if (profile_stages) { - const auto work_end = std::chrono::steady_clock::now(); - const auto profile = profile_bulk_edge_count( - scratch, chunk_result, count_config.concurrent, counts_out_degree, - counts_in_degree); - auto& stats = worker_stats[static_cast(worker)]; - ++stats.chunks; - stats.edges += chunk_result.valid_edges; - stats.out_updates += profile.out_updates; - stats.in_updates += profile.in_updates; - stats.out_single_checks += chunk_result.out_single_checks; - stats.in_single_checks += chunk_result.in_single_checks; - stats.work_time_ns += - std::chrono::duration_cast(work_end - start) - .count(); - } + index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); + count_chunk(scratch, concurrent); }; - VLOG(1) << "Bulk edge count pass: bytes=" << source->EstimatedBytes() - << ", producers=" << allocation.producer_count - << ", consumers=" << allocation.consumer_count - << ", queue_capacity=" << allocation.queue_capacity - << ", check_out_single=" << check_out_single - << ", check_in_single=" << check_in_single - << ", direct_supplier_queue=" - << supplier->SupportsConcurrentGetNext(); - const auto consume_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; - if (supplier->SupportsConcurrentGetNext()) { - consume_concurrent_supplier_indexed(*supplier, worker_count, count); - } else if (worker_count == 1) { - while (auto chunk = supplier->GetNextChunk()) { - count(0, chunk); - } - } else { - ChunkPipelineOptions options; - options.consumer_count = worker_count; - options.queue_capacity = allocation.queue_capacity; - consume_chunk_pipeline_indexed(*supplier, options, count); - } - const auto consume_end = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; + consume_supplier_indexed(*supplier, options, count); BulkEdgeCountSummary summary; - for (size_t worker = 0; worker < scratches.size(); ++worker) { - summary.valid_edges += scratches[worker].valid_edges; + for (const auto& scratch : scratches) { summary.out_single_duplicate = - summary.out_single_duplicate || scratches[worker].out_single_duplicate; + summary.out_single_duplicate || scratch.out_single_duplicate; summary.in_single_duplicate = - summary.in_single_duplicate || scratches[worker].in_single_duplicate; - } - if (!profile_stages) { - return summary; - } - size_t out_updates = 0; - size_t in_updates = 0; - size_t out_single_checks = 0; - size_t in_single_checks = 0; - int64_t vid_degree_time_ns = 0; - for (const auto& stats : worker_stats) { - out_updates += stats.out_updates; - in_updates += stats.in_updates; - out_single_checks += stats.out_single_checks; - in_single_checks += stats.in_single_checks; - vid_degree_time_ns += stats.work_time_ns; - } - for (size_t worker = 0; worker < worker_stats.size(); ++worker) { - const auto& stats = worker_stats[worker]; - VLOG(2) << "Bulk edge count worker: worker=" << worker - << ", chunks=" << stats.chunks << ", edges=" << stats.edges - << ", out_atomic_updates=" << stats.out_updates - << ", in_atomic_updates=" << stats.in_updates - << ", out_single_checks=" << stats.out_single_checks - << ", in_single_checks=" << stats.in_single_checks - << ", worker_ms=" << stats.work_time_ns / 1000000; - } - VLOG(1) << "Bulk edge count stages: rows=" << row_num << ", row_count_ms=" - << std::chrono::duration_cast( - row_count_end - row_count_start) - .count() - << ", pipeline_wall_ms=" - << std::chrono::duration_cast( - consume_end - consume_start) - .count() - << ", valid_edges=" << summary.valid_edges - << ", out_atomic_updates=" << out_updates - << ", in_atomic_updates=" << in_updates - << ", out_single_checks=" << out_single_checks - << ", in_single_checks=" << in_single_checks - << ", out_single_duplicate=" << summary.out_single_duplicate - << ", in_single_duplicate=" << summary.in_single_duplicate - << ", vid_degree_worker_ms=" << vid_degree_time_ns / 1000000; + summary.in_single_duplicate || scratch.in_single_duplicate; + } return summary; } -constexpr uint32_t kInvalidBulkEdgeIndex = std::numeric_limits::max(); - -struct BulkEdgeGroup { - uint32_t count = 0; - uint32_t head = kInvalidBulkEdgeIndex; - uint32_t tail = kInvalidBulkEdgeIndex; -}; - -struct BulkEdgeFillScratch { - BulkEdgeEndpointScratch endpoints; - flat_hash_map out_groups; - flat_hash_map in_groups; - std::vector out_next; - std::vector in_next; -}; - -struct BulkEdgeFillResult { - size_t valid_edges = 0; - size_t out_reservations = 0; - size_t in_reservations = 0; -}; - void append_bulk_edge_group(flat_hash_map& groups, std::vector& next, vid_t vertex, uint32_t edge_index) { @@ -885,34 +710,32 @@ void append_bulk_edge_group(flat_hash_map& groups, } template -void group_bulk_edge_chunk(const BulkEdgeEndpointScratch& endpoints, - BulkEdgeFillScratch& scratch) { - CHECK_LE(endpoints.src_lids.size(), +void group_bulk_edge_chunk(BulkEdgeWorkerScratch& scratch) { + CHECK_LE(scratch.src_lids.size(), static_cast(std::numeric_limits::max())); - if constexpr (OutWriter::kNeedsConcurrentGrouping) { + if constexpr (OutWriter::kNeedsDegreeCount) { scratch.out_groups.clear(); scratch.out_groups.reserve( - bulk_edge_group_reserve(endpoints.src_lids.size())); - scratch.out_next.assign(endpoints.src_lids.size(), kInvalidBulkEdgeIndex); + bulk_edge_group_reserve(scratch.src_lids.size())); + scratch.out_next.assign(scratch.src_lids.size(), kInvalidBulkEdgeIndex); } - if constexpr (InWriter::kNeedsConcurrentGrouping) { + if constexpr (InWriter::kNeedsDegreeCount) { scratch.in_groups.clear(); - scratch.in_groups.reserve( - bulk_edge_group_reserve(endpoints.dst_lids.size())); - scratch.in_next.assign(endpoints.dst_lids.size(), kInvalidBulkEdgeIndex); + scratch.in_groups.reserve(bulk_edge_group_reserve(scratch.dst_lids.size())); + scratch.in_next.assign(scratch.dst_lids.size(), kInvalidBulkEdgeIndex); } - for (size_t row = 0; row < endpoints.src_lids.size(); ++row) { - const auto src = endpoints.src_lids[row]; - const auto dst = endpoints.dst_lids[row]; + for (size_t row = 0; row < scratch.src_lids.size(); ++row) { + const auto src = scratch.src_lids[row]; + const auto dst = scratch.dst_lids[row]; if (!is_valid_bulk_edge(src, dst)) { continue; } const auto edge_index = static_cast(row); - if constexpr (OutWriter::kNeedsConcurrentGrouping) { + if constexpr (OutWriter::kNeedsDegreeCount) { append_bulk_edge_group(scratch.out_groups, scratch.out_next, src, edge_index); } - if constexpr (InWriter::kNeedsConcurrentGrouping) { + if constexpr (InWriter::kNeedsDegreeCount) { append_bulk_edge_group(scratch.in_groups, scratch.in_next, dst, edge_index); } @@ -920,13 +743,13 @@ void group_bulk_edge_chunk(const BulkEdgeEndpointScratch& endpoints, } template -void fill_bulk_edge_chunk_serial(const BulkEdgeEndpointScratch& endpoints, +void fill_bulk_edge_chunk_serial(const BulkEdgeWorkerScratch& scratch, const BulkEdgeDataReader& data_reader, OutWriter& out, InWriter& in) { size_t valid_edges = 0; - for (size_t row = 0; row < endpoints.src_lids.size(); ++row) { - const auto src = endpoints.src_lids[row]; - const auto dst = endpoints.dst_lids[row]; + for (size_t row = 0; row < scratch.src_lids.size(); ++row) { + const auto src = scratch.src_lids[row]; + const auto dst = scratch.dst_lids[row]; if (!is_valid_bulk_edge(src, dst)) { continue; } @@ -949,19 +772,19 @@ void fill_bulk_edge_chunk_serial(const BulkEdgeEndpointScratch& endpoints, template void fill_bulk_edge_chunk_concurrent( - const BulkEdgeEndpointScratch& endpoints, + BulkEdgeWorkerScratch& scratch, const BulkEdgeDataReader& data_reader, OutWriter& out, - InWriter& in, BulkEdgeFillScratch& scratch) { - group_bulk_edge_chunk(endpoints, scratch); + InWriter& in) { + group_bulk_edge_chunk(scratch); constexpr bool kDirectOut = - OutWriter::kStoresEdges && !OutWriter::kNeedsConcurrentGrouping; + OutWriter::kStoresEdges && !OutWriter::kNeedsDegreeCount; constexpr bool kDirectIn = - InWriter::kStoresEdges && !InWriter::kNeedsConcurrentGrouping; + InWriter::kStoresEdges && !InWriter::kNeedsDegreeCount; if constexpr (kDirectOut || kDirectIn) { size_t valid_edges = 0; - for (size_t row = 0; row < endpoints.src_lids.size(); ++row) { - const auto src = endpoints.src_lids[row]; - const auto dst = endpoints.dst_lids[row]; + for (size_t row = 0; row < scratch.src_lids.size(); ++row) { + const auto src = scratch.src_lids[row]; + const auto dst = scratch.dst_lids[row]; if (!is_valid_bulk_edge(src, dst)) { continue; } @@ -981,7 +804,7 @@ void fill_bulk_edge_chunk_concurrent( in.RecordFilledEdges(valid_edges); } } - if constexpr (OutWriter::kNeedsConcurrentGrouping) { + if constexpr (OutWriter::kNeedsDegreeCount) { for (const auto& [src, group] : scratch.out_groups) { CHECK_LE(group.count, static_cast(std::numeric_limits::max())); @@ -990,13 +813,13 @@ void fill_bulk_edge_chunk_concurrent( for (uint32_t i = 0; i < group.count; ++i) { CHECK_NE(edge_index, kInvalidBulkEdgeIndex); const auto data = data_reader.Get(edge_index); - out.PutAt(src, slot++, endpoints.dst_lids[edge_index], data, 0); + out.PutAt(src, slot++, scratch.dst_lids[edge_index], data, 0); edge_index = scratch.out_next[edge_index]; } CHECK_EQ(edge_index, kInvalidBulkEdgeIndex); } } - if constexpr (InWriter::kNeedsConcurrentGrouping) { + if constexpr (InWriter::kNeedsDegreeCount) { for (const auto& [dst, group] : scratch.in_groups) { CHECK_LE(group.count, static_cast(std::numeric_limits::max())); @@ -1005,7 +828,7 @@ void fill_bulk_edge_chunk_concurrent( for (uint32_t i = 0; i < group.count; ++i) { CHECK_NE(edge_index, kInvalidBulkEdgeIndex); const auto data = data_reader.Get(edge_index); - in.PutAt(dst, slot++, endpoints.src_lids[edge_index], data, 0); + in.PutAt(dst, slot++, scratch.src_lids[edge_index], data, 0); edge_index = scratch.in_next[edge_index]; } CHECK_EQ(edge_index, kInvalidBulkEdgeIndex); @@ -1013,56 +836,18 @@ void fill_bulk_edge_chunk_concurrent( } } -BulkEdgeFillResult profile_bulk_edge_fill( - const BulkEdgeEndpointScratch& endpoints, - const BulkEdgeFillScratch* concurrent_scratch, bool stores_out_direction, - bool stores_in_direction, bool groups_out_direction, - bool groups_in_direction) { - BulkEdgeFillResult result; - result.valid_edges = count_valid_bulk_edges(endpoints); - if (stores_out_direction) { - result.out_reservations = - concurrent_scratch == nullptr - ? result.valid_edges - : (groups_out_direction ? concurrent_scratch->out_groups.size() - : 0); - } - if (stores_in_direction) { - result.in_reservations = - concurrent_scratch == nullptr - ? result.valid_edges - : (groups_in_direction ? concurrent_scratch->in_groups.size() : 0); - } - return result; -} - -template -using BulkEdgeSerialFillChunk = std::function&)>; - template -using BulkEdgeConcurrentFillChunk = std::function&, - BulkEdgeFillScratch&)>; +using BulkEdgeFillChunk = std::function&, + BulkEdgeWorkerScratch&, bool)>; template struct BulkEdgeBuildOps { - bool stores_out_direction = false; - bool stores_in_direction = false; - bool counts_out_degree = false; - bool counts_in_degree = false; + bool needs_degree_count = false; bool checks_out_single = false; bool checks_in_single = false; - bool groups_out_direction = false; - bool groups_in_direction = false; - bool supports_disjoint_concurrent_fill = false; - std::function prepare_build; BulkEdgeCountChunk count_chunk; - std::function set_input_edge_count; std::function allocate_from_counts; - BulkEdgeSerialFillChunk fill_serial_chunk; - BulkEdgeConcurrentFillChunk fill_concurrent_chunk; - std::function finish; + BulkEdgeFillChunk fill_chunk; }; template @@ -1073,248 +858,69 @@ BulkEdgeBuildOps make_bulk_edge_build_ops(OutWriter& out, OutWriter::kStoresEdges || InWriter::kStoresEdges; constexpr bool kNeedsAnyDegreeCount = OutWriter::kNeedsDegreeCount || InWriter::kNeedsDegreeCount; - constexpr bool kSupportsCountPass = kNeedsAnyDegreeCount || - OutWriter::kChecksSingleUniqueness || - InWriter::kChecksSingleUniqueness; - ops.stores_out_direction = OutWriter::kStoresEdges; - ops.stores_in_direction = InWriter::kStoresEdges; - ops.counts_out_degree = OutWriter::kNeedsDegreeCount; - ops.counts_in_degree = InWriter::kNeedsDegreeCount; + ops.needs_degree_count = kNeedsAnyDegreeCount; ops.checks_out_single = OutWriter::kChecksSingleUniqueness; ops.checks_in_single = InWriter::kChecksSingleUniqueness; - ops.groups_out_direction = OutWriter::kNeedsConcurrentGrouping; - ops.groups_in_direction = InWriter::kNeedsConcurrentGrouping; - ops.prepare_build = [&out, &in](vid_t out_vertices, vid_t in_vertices) { - out.PrepareBuild(out_vertices); - in.PrepareBuild(in_vertices); - }; ops.allocate_from_counts = [&out, &in]() { out.AllocateFromCounts(); in.AllocateFromCounts(); }; - if constexpr (kSupportsCountPass) { - ops.count_chunk = [&out, &in](BulkEdgeCountScratch& scratch, - const BulkEdgeCountConfig& config) { - return count_bulk_edge_chunk(scratch.endpoints, out, in, scratch, config); - }; - } - if constexpr (OutWriter::kTracksInputEdgeCount && - InWriter::kTracksInputEdgeCount) { - ops.set_input_edge_count = [&out, &in](uint64_t count) { - out.SetInputEdgeCount(count); - in.SetInputEdgeCount(count); - }; - } else if constexpr (OutWriter::kTracksInputEdgeCount) { - ops.set_input_edge_count = [&out](uint64_t count) { - out.SetInputEdgeCount(count); - }; - } else if constexpr (InWriter::kTracksInputEdgeCount) { - ops.set_input_edge_count = [&in](uint64_t count) { - in.SetInputEdgeCount(count); + if constexpr (kNeedsAnyDegreeCount) { + ops.count_chunk = [&out, &in](BulkEdgeWorkerScratch& scratch, + bool concurrent) { + count_bulk_edge_chunk(scratch, out, in, concurrent); }; } if constexpr (kStoresAnyDirection) { - ops.supports_disjoint_concurrent_fill = - OutWriter::kSupportsDisjointConcurrentFill && - InWriter::kSupportsDisjointConcurrentFill; - ops.fill_serial_chunk = - [&out, &in](const BulkEdgeEndpointScratch& endpoints, - const BulkEdgeDataReader& data_reader) { - fill_bulk_edge_chunk_serial(endpoints, data_reader, out, in); - }; - if constexpr (OutWriter::kSupportsDisjointConcurrentFill && - InWriter::kSupportsDisjointConcurrentFill) { - ops.fill_concurrent_chunk = - [&out, &in](const BulkEdgeEndpointScratch& endpoints, - const BulkEdgeDataReader& data_reader, - BulkEdgeFillScratch& scratch) { - fill_bulk_edge_chunk_concurrent(endpoints, data_reader, out, in, - scratch); - }; - } + ops.fill_chunk = [&out, &in](const BulkEdgeDataReader& data_reader, + BulkEdgeWorkerScratch& scratch, + bool concurrent) { + if (concurrent) { + fill_bulk_edge_chunk_concurrent(scratch, data_reader, out, in); + } else { + fill_bulk_edge_chunk_serial(scratch, data_reader, out, in); + } + }; } - ops.finish = [&out, &in]() { - out.Finish(); - in.Finish(); - }; return ops; } -struct BulkEdgeFillWorkerStats { - size_t chunks = 0; - size_t edges = 0; - size_t out_reservations = 0; - size_t in_reservations = 0; - int64_t endpoint_time_ns = 0; - int64_t fill_time_ns = 0; -}; - template void fill_bulk_edges(const std::shared_ptr& source, const IndexerType& src_indexer, const IndexerType& dst_indexer, - const ChunkPipelineAllocation& allocation, - bool preserve_order, bool allow_concurrent_fill, + const ChunkSourceOptions& options, + bool allow_concurrent_fill, + std::vector& scratches, const BulkEdgeBuildOps& ops) { - CHECK(!allow_concurrent_fill || ops.supports_disjoint_concurrent_fill); - CHECK(!allow_concurrent_fill || !preserve_order); - const auto source_options = - make_bulk_edge_source_options(allocation, preserve_order); - auto supplier = source->Open(source_options); + CHECK(!allow_concurrent_fill || !options.preserve_order); + auto supplier = source->Open(options); CHECK(supplier != nullptr); - const bool profile_stages = VLOG_IS_ON(1); - const auto pipeline_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; - const auto worker_count = allocation.consumer_count; - std::vector worker_stats; - if (profile_stages) { - worker_stats.resize(static_cast(worker_count)); - } - bool concurrent_fill = false; - - if (allow_concurrent_fill) { - if (worker_count > 1) { - CHECK(static_cast(ops.fill_concurrent_chunk)); - concurrent_fill = true; - std::vector scratches( - static_cast(worker_count)); - auto fill = [&](int32_t worker, const std::shared_ptr& chunk) { - CHECK_GE(worker, 0); - CHECK_LT(worker, worker_count); - auto& scratch = scratches[static_cast(worker)]; - const auto endpoint_start = - profile_stages ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; - index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, - scratch.endpoints); - if (profile_stages) { - worker_stats[static_cast(worker)].endpoint_time_ns += - std::chrono::duration_cast( - std::chrono::steady_clock::now() - endpoint_start) - .count(); - } - const auto fill_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; - const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; - BulkEdgeDataReader data_reader(data_column); - ops.fill_concurrent_chunk(scratch.endpoints, data_reader, scratch); - if (profile_stages) { - const auto fill_end = std::chrono::steady_clock::now(); - const auto result = profile_bulk_edge_fill( - scratch.endpoints, &scratch, ops.stores_out_direction, - ops.stores_in_direction, ops.groups_out_direction, - ops.groups_in_direction); - auto& stats = worker_stats[static_cast(worker)]; - ++stats.chunks; - stats.edges += result.valid_edges; - stats.out_reservations += result.out_reservations; - stats.in_reservations += result.in_reservations; - stats.fill_time_ns += - std::chrono::duration_cast(fill_end - - fill_start) - .count(); - } - }; - - if (supplier->SupportsConcurrentGetNext()) { - consume_concurrent_supplier_indexed(*supplier, worker_count, fill); - } else { - ChunkPipelineOptions options; - options.consumer_count = worker_count; - options.queue_capacity = allocation.queue_capacity; - consume_chunk_pipeline_indexed(*supplier, options, fill); - } - } - } - - if (!concurrent_fill) { - BulkEdgeEndpointScratch endpoints; - while (auto chunk = supplier->GetNextChunk()) { - const auto endpoint_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; - index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, endpoints); - if (profile_stages) { - worker_stats.front().endpoint_time_ns += - std::chrono::duration_cast( - std::chrono::steady_clock::now() - endpoint_start) - .count(); - } - const auto fill_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; + const auto worker_count = options.consumer_count; + CHECK_GE(scratches.size(), static_cast(worker_count)); + if (allow_concurrent_fill && worker_count > 1) { + auto fill = [&](int32_t worker, const std::shared_ptr& chunk) { + CHECK_GE(worker, 0); + CHECK_LT(worker, worker_count); + auto& scratch = scratches[static_cast(worker)]; + index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; BulkEdgeDataReader data_reader(data_column); - ops.fill_serial_chunk(endpoints, data_reader); - if (profile_stages) { - const auto fill_end = std::chrono::steady_clock::now(); - const auto result = profile_bulk_edge_fill( - endpoints, nullptr, ops.stores_out_direction, - ops.stores_in_direction, ops.groups_out_direction, - ops.groups_in_direction); - auto& stats = worker_stats.front(); - ++stats.chunks; - stats.edges += result.valid_edges; - stats.out_reservations += result.out_reservations; - stats.in_reservations += result.in_reservations; - stats.fill_time_ns += - std::chrono::duration_cast(fill_end - - fill_start) - .count(); - } - } - } + ops.fill_chunk(data_reader, scratch, true); + }; - const auto pipeline_end = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; - if (!profile_stages) { + consume_supplier_indexed(*supplier, options, fill); return; } - const auto pipeline_wall_ms = - std::chrono::duration_cast(pipeline_end - - pipeline_start) - .count(); - size_t total_edges = 0; - size_t total_chunks = 0; - size_t out_reservations = 0; - size_t in_reservations = 0; - size_t min_worker_edges = std::numeric_limits::max(); - size_t max_worker_edges = 0; - int64_t endpoint_time_ns = 0; - int64_t fill_time_ns = 0; - for (const auto& stats : worker_stats) { - total_edges += stats.edges; - total_chunks += stats.chunks; - out_reservations += stats.out_reservations; - in_reservations += stats.in_reservations; - min_worker_edges = std::min(min_worker_edges, stats.edges); - max_worker_edges = std::max(max_worker_edges, stats.edges); - endpoint_time_ns += stats.endpoint_time_ns; - fill_time_ns += stats.fill_time_ns; - } - for (size_t worker = 0; worker < worker_stats.size(); ++worker) { - const auto& stats = worker_stats[worker]; - VLOG(2) << "Bulk edge fill worker: worker=" << worker - << ", chunks=" << stats.chunks << ", edges=" << stats.edges - << ", out_range_reservations=" << stats.out_reservations - << ", in_range_reservations=" << stats.in_reservations - << ", endpoint_index_ms=" << stats.endpoint_time_ns / 1000000 - << ", fill_ms=" << stats.fill_time_ns / 1000000; - } - VLOG(1) << "Bulk edge fill stages: pipeline_wall_ms=" << pipeline_wall_ms - << ", concurrent_fill=" << concurrent_fill - << ", fill_workers=" << (concurrent_fill ? worker_count : 1) - << ", chunks=" << total_chunks << ", edges=" << total_edges - << ", out_range_reservations=" << out_reservations - << ", in_range_reservations=" << in_reservations - << ", min_worker_edges=" << min_worker_edges - << ", max_worker_edges=" << max_worker_edges - << ", endpoint_index_worker_ms=" << endpoint_time_ns / 1000000 - << ", csr_fill_worker_ms=" << fill_time_ns / 1000000; + + auto& scratch = scratches.front(); + while (auto chunk = supplier->GetNextChunk()) { + index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); + const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; + BulkEdgeDataReader data_reader(data_column); + ops.fill_chunk(data_reader, scratch, false); + } } vid_t csr_vertex_capacity(const IndexerType& indexer) { @@ -1331,94 +937,43 @@ void build_bundled_edges_with_ops( const std::shared_ptr& source) { // VertexTable owns its reserve policy. Edge storage consumes the resulting // indexer capacity instead of duplicating PropertyGraph::Dump's policy. - ops.prepare_build(csr_vertex_capacity(src_indexer), - csr_vertex_capacity(dst_indexer)); - if (!ops.stores_out_direction && !ops.stores_in_direction) { + if (!ops.fill_chunk) { ops.allocate_from_counts(); - ops.finish(); return; } const auto estimated_bytes = source->EstimatedBytes(); - const auto count_allocation = resolve_chunk_pipeline_allocation( + const auto count_options = resolve_bulk_edge_source_options( estimated_bytes, source->ParallelEnabled(), false); - const bool needs_degree_count = ops.counts_out_degree || ops.counts_in_degree; + const bool needs_degree_count = ops.needs_degree_count; + std::vector scratches; + if (needs_degree_count) { + scratches.resize(static_cast(count_options.consumer_count)); + } const bool has_single_direction = ops.checks_out_single || ops.checks_in_single; - // Single-only layouts retain their one-pass ordered path. Mixed - // Single/Mutable layouts already need the degree pass, so that pass also - // verifies whether fixed Single slots are disjoint across all chunks. - const bool check_single_uniqueness = - has_single_direction && needs_degree_count; - const bool run_count_pass = needs_degree_count || check_single_uniqueness; - VLOG(1) << "Bulk edge count allocation: bytes=" << estimated_bytes - << ", producers=" << count_allocation.producer_count - << ", consumers=" << count_allocation.consumer_count - << ", queue_capacity=" << count_allocation.queue_capacity - << ", degree_count_pass=" << needs_degree_count - << ", single_uniqueness_check=" << check_single_uniqueness - << ", count_pass=" << run_count_pass; BulkEdgeCountSummary count_summary; - if (run_count_pass) { + if (needs_degree_count) { CHECK(static_cast(ops.count_chunk)); - count_summary = count_bulk_edges( - source, src_indexer, dst_indexer, count_allocation, - ops.counts_out_degree, ops.counts_in_degree, - check_single_uniqueness && ops.checks_out_single, - check_single_uniqueness && ops.checks_in_single, ops.count_chunk); - } - const bool single_uniqueness_verified = - has_single_direction && check_single_uniqueness; - if (single_uniqueness_verified) { - CHECK(static_cast(ops.set_input_edge_count)); - ops.set_input_edge_count(count_summary.valid_edges); + count_summary = count_bulk_edges(source, src_indexer, dst_indexer, + count_options, scratches, ops.count_chunk); } const bool single_duplicate = (ops.checks_out_single && count_summary.out_single_duplicate) || (ops.checks_in_single && count_summary.in_single_duplicate); - const bool single_slots_disjoint = - !has_single_direction || - (single_uniqueness_verified && !single_duplicate); - const bool allow_concurrent_fill = - ops.supports_disjoint_concurrent_fill && single_slots_disjoint; const bool preserve_fill_order = - has_single_direction && !single_slots_disjoint; - const auto fill_allocation = + has_single_direction && (!needs_degree_count || single_duplicate); + const bool allow_concurrent_fill = !preserve_fill_order; + const auto fill_options = preserve_fill_order - ? resolve_chunk_pipeline_allocation(estimated_bytes, - source->ParallelEnabled(), true) - : count_allocation; - VLOG(1) << "Bulk edge fill allocation: bytes=" << estimated_bytes - << ", producers=" << fill_allocation.producer_count - << ", consumers=" << fill_allocation.consumer_count - << ", queue_capacity=" << fill_allocation.queue_capacity - << ", single_uniqueness_verified=" << single_uniqueness_verified - << ", out_single_duplicate=" << count_summary.out_single_duplicate - << ", in_single_duplicate=" << count_summary.in_single_duplicate - << ", allow_concurrent_fill=" << allow_concurrent_fill - << ", preserve_order=" << preserve_fill_order; - const bool profile_stages = VLOG_IS_ON(1); - const auto allocate_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; - ops.allocate_from_counts(); - const auto fill_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; - fill_bulk_edges(source, src_indexer, dst_indexer, fill_allocation, - preserve_fill_order, allow_concurrent_fill, ops); - ops.finish(); - if (!profile_stages) { - return; + ? resolve_bulk_edge_source_options(estimated_bytes, + source->ParallelEnabled(), true) + : count_options; + if (scratches.size() < static_cast(fill_options.consumer_count)) { + scratches.resize(static_cast(fill_options.consumer_count)); } - const auto fill_end = std::chrono::steady_clock::now(); - VLOG(1) << "Bulk edge CSR stages: allocate_ms=" - << std::chrono::duration_cast( - fill_start - allocate_start) - .count() - << ", fill_pass_ms=" - << std::chrono::duration_cast(fill_end - - fill_start) - .count(); + ops.allocate_from_counts(); + fill_bulk_edges(source, src_indexer, dst_indexer, fill_options, + allow_concurrent_fill, scratches, ops); } template @@ -1431,9 +986,13 @@ bool build_bundled_edges_typed( with_csr_bulk_writer(out_csr, [&](auto& out) { const bool in_supported = with_csr_bulk_writer(in_csr, [&](auto& in) { + out.PrepareBuild(csr_vertex_capacity(src_indexer)); + in.PrepareBuild(csr_vertex_capacity(dst_indexer)); auto ops = make_bulk_edge_build_ops(out, in); build_bundled_edges_with_ops(ops, src_indexer, dst_indexer, source); + out.Finish(); + in.Finish(); }); built = in_supported; }); @@ -1845,34 +1404,20 @@ std::pair EdgeTable::AddEdge( void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, const IndexerType& dst_indexer, std::shared_ptr supplier) { - const bool profile_stages = VLOG_IS_ON(1); - const auto total_start = std::chrono::steady_clock::now(); - int64_t row_num_time_ns = 0; - int64_t collect_time_ns = 0; - int64_t filter_time_ns = 0; - int64_t ensure_capacity_time_ns = 0; - int64_t write_time_ns = 0; - size_t chunk_count = 0; - size_t input_rows = 0; - size_t cached_property_columns = 0; - size_t peak_buffer_bytes = 0; - const auto peak_rss_start_bytes = - profile_stages ? process_peak_rss_bytes() : 0; + CHECK(supplier != nullptr); + auto source = supplier->RepeatableSource(); + if (source && should_use_bulk_edge_build(*source) && + TryBatchBuildEdges(src_indexer, dst_indexer, source)) { + return; + } + // Keep fallback COPY paths aligned with the vertex table's actual capacity, // while leaving completely unloaded edge tables lazy until persistence. in_csr_->resize(csr_vertex_capacity(dst_indexer)); out_csr_->resize(csr_vertex_capacity(src_indexer)); std::vector src_lid, dst_lid; // Pre-reserve capacity to reduce vector reallocation on large graphs. - const auto row_num_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; auto total_rows = supplier->RowNum(); - if (profile_stages) { - row_num_time_ns = std::chrono::duration_cast( - std::chrono::steady_clock::now() - row_num_start) - .count(); - } if (total_rows > 0) { src_lid.reserve(total_rows); dst_lid.reserve(total_rows); @@ -1883,15 +1428,10 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, std::vector> unbundled_data_chunks; std::vector valid_flags; while (true) { - const auto collect_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; auto chunk = supplier->GetNextChunk(); if (chunk == nullptr) { break; } - ++chunk_count; - input_rows += chunk->row_num(); auto src_col = chunk->get(0); auto dst_col = chunk->get(1); src_indexer.get_index(*src_col, src_lid); @@ -1900,7 +1440,6 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, if (meta_->is_bundled()) { // Bundled: only one property column (index 2). bundled_data_cols.push_back(chunk->get(2)); - ++cached_property_columns; } else { // Unbundled: collect remaining columns as a DataChunk. auto prop_chunk = std::make_shared(); @@ -1908,40 +1447,13 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, auto c = chunk->get(static_cast(i)); if (c) { prop_chunk->set(static_cast(i - 2), c); - ++cached_property_columns; } } unbundled_data_chunks.push_back(prop_chunk); } } - if (profile_stages) { - collect_time_ns += std::chrono::duration_cast( - std::chrono::steady_clock::now() - collect_start) - .count(); - peak_buffer_bytes = std::max( - peak_buffer_bytes, estimate_edge_fallback_buffer_bytes( - src_lid, dst_lid, valid_flags, - bundled_data_cols, unbundled_data_chunks)); - } } - const auto filter_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; filterInvalidEdges(src_lid, dst_lid, valid_flags); - if (profile_stages) { - filter_time_ns = std::chrono::duration_cast( - std::chrono::steady_clock::now() - filter_start) - .count(); - peak_buffer_bytes = std::max( - peak_buffer_bytes, estimate_edge_fallback_buffer_bytes( - src_lid, dst_lid, valid_flags, bundled_data_cols, - unbundled_data_chunks)); - } - const auto valid_rows = src_lid.size(); - const auto invalid_rows = valid_flags.size() - valid_rows; - const auto ensure_capacity_start = - profile_stages ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; size_t new_size = table_idx_.load() + src_lid.size(); if (new_size >= Capacity()) { auto new_cap = new_size; @@ -1950,15 +1462,6 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, } EnsureCapacity(new_cap); } - if (profile_stages) { - ensure_capacity_time_ns = - std::chrono::duration_cast( - std::chrono::steady_clock::now() - ensure_capacity_start) - .count(); - } - const auto write_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; if (meta_->is_bundled()) { batch_add_bundled_edges_impl(out_csr_.get(), in_csr_.get(), meta_, src_lid, dst_lid, bundled_data_cols, valid_flags); @@ -1970,48 +1473,26 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, src_lid, dst_lid, oe_csr, ie_csr, table_.get(), table_idx_, capacity_, meta_->properties, unbundled_data_chunks, valid_flags); } - if (profile_stages) { - write_time_ns = std::chrono::duration_cast( - std::chrono::steady_clock::now() - write_start) - .count(); - const auto total_time_ns = - std::chrono::duration_cast( - std::chrono::steady_clock::now() - total_start) - .count(); - const auto peak_rss_end_bytes = process_peak_rss_bytes(); - VLOG(1) << "Fallback edge load stages: input_rows=" << input_rows - << ", valid_rows=" << valid_rows - << ", invalid_rows=" << invalid_rows << ", chunks=" << chunk_count - << ", cached_property_columns=" << cached_property_columns - << ", peak_buffer_bytes=" << peak_buffer_bytes - << ", process_peak_rss_start_bytes=" << peak_rss_start_bytes - << ", process_peak_rss_end_bytes=" << peak_rss_end_bytes - << ", total_ms=" << total_time_ns / 1000000 - << ", row_num_ms=" << row_num_time_ns / 1000000 - << ", collect_lookup_ms=" << collect_time_ns / 1000000 - << ", filter_ms=" << filter_time_ns / 1000000 - << ", ensure_capacity_ms=" << ensure_capacity_time_ns / 1000000 - << ", write_ms=" << write_time_ns / 1000000; - } -} - -void EdgeTable::BatchBuildEdges(const IndexerType& src_indexer, - const IndexerType& dst_indexer, - std::shared_ptr source) { - if (!source) { - THROW_INVALID_ARGUMENT_EXCEPTION( - "BatchBuildEdges requires a non-null data source"); +} + +bool EdgeTable::TryBatchBuildEdges( + const IndexerType& src_indexer, const IndexerType& dst_indexer, + const std::shared_ptr& source) { + if (!source || !meta_ || !meta_->is_bundled() || !out_csr_ || !in_csr_ || + out_csr_->edge_num() != 0 || in_csr_->edge_num() != 0 || + (meta_->oe_strategy != EdgeStrategy::kNone && !meta_->oe_mutable) || + (meta_->ie_strategy != EdgeStrategy::kNone && !meta_->ie_mutable)) { + return false; } - CHECK(source->rewindable()); - CHECK(CanBatchBuild()) - << "Bulk edge build requires an empty bundled edge table"; EdgeTable staged(meta_); staged.Init(ckp_, memory_level_); - CHECK(build_bundled_edges(staged.out_csr_.get(), staged.in_csr_.get(), meta_, - src_indexer, dst_indexer, source)) - << "Bulk edge build does not support this CSR layout"; + if (!build_bundled_edges(staged.out_csr_.get(), staged.in_csr_.get(), meta_, + src_indexer, dst_indexer, source)) { + return false; + } Swap(staged); + return true; } void EdgeTable::BatchAddEdges( diff --git a/src/storages/graph/graph_interface.cc b/src/storages/graph/graph_interface.cc index 8c2150924..15b4ce0db 100644 --- a/src/storages/graph/graph_interface.cc +++ b/src/storages/graph/graph_interface.cc @@ -117,18 +117,8 @@ Status StorageAPUpdateInterface::DeleteEdges(label_t src_label, vid_t src, Status StorageAPUpdateInterface::BatchAddVertices( label_t v_label_id, std::shared_ptr supplier) { - return graph_.BatchAddVertices(v_label_id, std::move(supplier)); -} - -bool StorageAPUpdateInterface::CanBatchBuildVertices(label_t v_label_id) const { - return graph_.CanBatchBuildVertices(v_label_id); -} - -Status StorageAPUpdateInterface::BatchBuildVertices( - label_t v_label_id, std::shared_ptr source) { - auto status = graph_.BatchBuildVertices(v_label_id, std::move(source)); + auto status = graph_.BatchAddVertices(v_label_id, std::move(supplier)); if (status.ok()) { - // The staged table swap replaces raw pointers cached by GraphView. mut_view_.Rebuild(graph_); } return status; @@ -137,23 +127,9 @@ Status StorageAPUpdateInterface::BatchBuildVertices( Status StorageAPUpdateInterface::BatchAddEdges( label_t src_label, label_t dst_label, label_t edge_label, std::shared_ptr supplier) { - return graph_.BatchAddEdges(src_label, dst_label, edge_label, - std::move(supplier)); -} - -bool StorageAPUpdateInterface::CanBatchBuildEdges(label_t src_label, - label_t dst_label, - label_t edge_label) const { - return graph_.CanBatchBuildEdges(src_label, dst_label, edge_label); -} - -Status StorageAPUpdateInterface::BatchBuildEdges( - label_t src_label, label_t dst_label, label_t edge_label, - std::shared_ptr source) { - auto status = graph_.BatchBuildEdges(src_label, dst_label, edge_label, - std::move(source)); + auto status = graph_.BatchAddEdges(src_label, dst_label, edge_label, + std::move(supplier)); if (status.ok()) { - // The staged CSR swap replaces raw pointers cached by GraphView. mut_view_.Rebuild(graph_); } return status; diff --git a/src/storages/graph/property_graph.cc b/src/storages/graph/property_graph.cc index 635974784..1eb225923 100644 --- a/src/storages/graph/property_graph.cc +++ b/src/storages/graph/property_graph.cc @@ -137,24 +137,7 @@ Status PropertyGraph::EnsureCapacity(label_t src_label, label_t dst_label, Status PropertyGraph::BatchAddVertices( label_t v_label, std::shared_ptr supplier) { RETURN_IF_NOT_OK(vertex_label_check(v_label)); - vertex_tables_[v_label].insert_vertices(supplier); - return neug::Status::OK(); -} - -bool PropertyGraph::CanBatchBuildVertices(label_t v_label) const { - return vertex_label_check(v_label).ok() && - vertex_tables_[v_label].CanBatchBuild(); -} - -Status PropertyGraph::BatchBuildVertices( - label_t v_label, std::shared_ptr source) { - RETURN_IF_NOT_OK(vertex_label_check(v_label)); - if (!source || !source->rewindable() || !CanBatchBuildVertices(v_label)) { - return Status(StatusCode::ERR_NOT_SUPPORTED, - "Bulk vertex build requires an empty label and a " - "repeatable source."); - } - vertex_tables_[v_label].BatchBuildVertices(std::move(source)); + vertex_tables_[v_label].insert_vertices(std::move(supplier)); return neug::Status::OK(); } @@ -166,35 +149,7 @@ Status PropertyGraph::BatchAddEdges( assert(edge_tables_.count(index) > 0); edge_tables_.at(index).BatchAddEdges( vertex_tables_.at(src_v_label).get_indexer(), - vertex_tables_.at(dst_v_label).get_indexer(), supplier); - return neug::Status::OK(); -} - -bool PropertyGraph::CanBatchBuildEdges(label_t src_v_label, label_t dst_v_label, - label_t e_label) const { - if (!edge_triplet_check(src_v_label, dst_v_label, e_label).ok()) { - return false; - } - const auto index = - schema_.generate_edge_label(src_v_label, dst_v_label, e_label); - auto it = edge_tables_.find(index); - return it != edge_tables_.end() && it->second.CanBatchBuild(); -} - -Status PropertyGraph::BatchBuildEdges( - label_t src_v_label, label_t dst_v_label, label_t e_label, - std::shared_ptr source) { - RETURN_IF_NOT_OK(edge_triplet_check(src_v_label, dst_v_label, e_label)); - if (!source || !source->rewindable() || - !CanBatchBuildEdges(src_v_label, dst_v_label, e_label)) { - return Status(StatusCode::ERR_NOT_SUPPORTED, - "Bulk edge build requires an empty bundled edge table and " - "a repeatable source."); - } - auto index = schema_.generate_edge_label(src_v_label, dst_v_label, e_label); - edge_tables_.at(index).BatchBuildEdges( - vertex_tables_.at(src_v_label).get_indexer(), - vertex_tables_.at(dst_v_label).get_indexer(), std::move(source)); + vertex_tables_.at(dst_v_label).get_indexer(), std::move(supplier)); return neug::Status::OK(); } diff --git a/src/storages/graph/vertex_table.cc b/src/storages/graph/vertex_table.cc index 031bafc65..db8518b43 100644 --- a/src/storages/graph/vertex_table.cc +++ b/src/storages/graph/vertex_table.cc @@ -15,9 +15,8 @@ #include "neug/storages/graph/vertex_table.h" -#include - #include "neug/storages/checkpoint_manifest.h" +#include "neug/storages/loader/loader_utils.h" #include "neug/storages/module/module_broker.h" #include "neug/storages/module/module_factory.h" #include "neug/storages/module_descriptor.h" @@ -49,15 +48,7 @@ void VertexTable::Init(std::shared_ptr ckp, MemoryLevel level) { void VertexTable::insert_vertices( std::shared_ptr supplier) { - const bool profile_stages = VLOG_IS_ON(1); - const auto total_start = std::chrono::steady_clock::now(); - int64_t row_num_time_ns = 0; - int64_t reserve_time_ns = 0; - int64_t get_chunk_time_ns = 0; - int64_t pk_time_ns = 0; - int64_t property_time_ns = 0; - size_t chunk_count = 0; - size_t loaded_rows = 0; + CHECK(supplier != nullptr); auto reserve_checkpoint_headroom = [this](size_t required_size) { const size_t headroom = required_size / 4; if (headroom > std::numeric_limits::max() - required_size) { @@ -69,15 +60,7 @@ void VertexTable::insert_vertices( } }; - const auto row_num_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; auto row_nums = supplier->RowNum(); - if (profile_stages) { - row_num_time_ns = std::chrono::duration_cast( - std::chrono::steady_clock::now() - row_num_start) - .count(); - } if (row_nums < 0) { VLOG(1) << "Row number from supplier is unknown, skip pre-reserve."; row_nums = 0; @@ -89,26 +72,9 @@ void VertexTable::insert_vertices( if (row_count > std::numeric_limits::max() - indexer_->size()) { THROW_RUNTIME_ERROR("Vertex row count overflow"); } - const auto pre_reserve_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; reserve_checkpoint_headroom(indexer_->size() + row_count); - if (profile_stages) { - reserve_time_ns += std::chrono::duration_cast( - std::chrono::steady_clock::now() - pre_reserve_start) - .count(); - } while (true) { - const auto get_chunk_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; auto chunk = supplier->GetNextChunk(); - if (profile_stages) { - get_chunk_time_ns += - std::chrono::duration_cast( - std::chrono::steady_clock::now() - get_chunk_start) - .count(); - } if (chunk == nullptr) { break; } @@ -136,96 +102,15 @@ void VertexTable::insert_vertices( THROW_RUNTIME_ERROR("Vertex row count overflow"); } size_t new_size = indexer_->size() + chunk_rows; - const auto chunk_reserve_start = - profile_stages ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; reserve_checkpoint_headroom(new_size); - if (profile_stages) { - reserve_time_ns += - std::chrono::duration_cast( - std::chrono::steady_clock::now() - chunk_reserve_start) - .count(); - } - - const auto pk_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; auto vids = insert_primary_keys(pk_col); - if (profile_stages) { - pk_time_ns += std::chrono::duration_cast( - std::chrono::steady_clock::now() - pk_start) - .count(); - } - - const auto property_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; for (size_t i = 0; i < prop_cols.size(); ++i) { auto col = table_->get_column_by_id(i); set_properties_from_context_column(col, prop_cols[i], vids); } - if (profile_stages) { - property_time_ns += std::chrono::duration_cast( - std::chrono::steady_clock::now() - property_start) - .count(); - ++chunk_count; - loaded_rows += chunk_rows; - } VLOG(10) << "Inserted " << chunk_rows << " vertices, current vertex num: " << VertexNum(); } - if (profile_stages) { - const auto total_time_ns = - std::chrono::duration_cast( - std::chrono::steady_clock::now() - total_start) - .count(); - VLOG(1) << "Vertex load stages: rows=" << loaded_rows - << ", chunks=" << chunk_count - << ", total_ms=" << total_time_ns / 1000000 - << ", row_num_ms=" << row_num_time_ns / 1000000 - << ", reserve_ms=" << reserve_time_ns / 1000000 - << ", get_chunk_ms=" << get_chunk_time_ns / 1000000 - << ", insert_pk_ms=" << pk_time_ns / 1000000 - << ", set_properties_ms=" << property_time_ns / 1000000; - } -} - -void VertexTable::BatchBuildVertices(std::shared_ptr source) { - if (!source) { - THROW_INVALID_ARGUMENT_EXCEPTION( - "BatchBuildVertices requires a non-null data source"); - } - CHECK(source->rewindable()); - CHECK(CanBatchBuild()) - << "Bulk vertex build requires an empty destination table"; - - auto supplier = source->Open(); - if (!supplier) { - THROW_INTERNAL_EXCEPTION("Data source returned a null supplier"); - } - const auto row_num = supplier->RowNum(); - - VertexTable staged(vertex_schema_); - staged.Init(ckp_, memory_level_); - if (row_num >= 0) { - if (static_cast(row_num) > std::numeric_limits::max()) { - THROW_RUNTIME_ERROR("Vertex row count exceeds addressable capacity"); - } - const auto row_count = static_cast(row_num); - // Keep the same headroom that PropertyGraph::DisassembleTo requires before - // writing a checkpoint. This avoids a second resize after the bulk build. - const size_t checkpoint_headroom = row_count / 4; - if (checkpoint_headroom > std::numeric_limits::max() - row_count) { - THROW_RUNTIME_ERROR("Vertex bulk load capacity overflow"); - } - staged.EnsureCapacity(row_count + checkpoint_headroom); - staged.insert_vertices_preallocated(std::move(supplier)); - } else { - // Unknown-cardinality suppliers still stream once and grow only when an - // actual chunk proves that more capacity is needed. - staged.insert_vertices(std::move(supplier)); - } - Swap(staged); } void VertexTable::Close() { diff --git a/src/storages/loader/loader_utils.cc b/src/storages/loader/loader_utils.cc index b8ae84ffe..9f6f24479 100644 --- a/src/storages/loader/loader_utils.cc +++ b/src/storages/loader/loader_utils.cc @@ -26,7 +26,6 @@ #include #include #include -#include #include #include #include @@ -855,8 +854,8 @@ class BoundedFileStream final : public std::istream { std::shared_ptr CsvPartitionPlanCache::GetOrCreate( const std::vector& file_paths, const CsvReadConfig& config, int32_t producer_count) { - producer_count = - std::clamp(producer_count, 1, hardware_worker_count()); + producer_count = std::clamp( + producer_count, 1, chunk_pipeline_detail::hardware_worker_count()); std::shared_ptr entry; { std::lock_guard lock(mutex); @@ -918,7 +917,7 @@ std::shared_ptr CsvPartitionPlanCache::GetOrCreate( // concurrently. With fewer files, the same H budget is divided among // intra-file speculative scans. scan_parallel() uses its caller for one // byte range, therefore the number of active scanners never exceeds H. - const auto scan_budget = hardware_worker_count(); + const auto scan_budget = chunk_pipeline_detail::hardware_worker_count(); std::vector scan_threads(file_paths.size(), 0); if (non_empty_files < static_cast(scan_budget)) { size_t assigned_scanners = 0; @@ -1549,45 +1548,58 @@ int64_t CSVChunkSupplier::RowNum() const { namespace { -class ColumnProjectingChunkSupplier final : public IDataChunkSupplier { +class SourceBackedChunkSupplier final : public IDataChunkSupplier { public: - ColumnProjectingChunkSupplier(std::shared_ptr input, - std::vector columns) - : input_(std::move(input)), columns_(std::move(columns)) { - CHECK(input_ != nullptr); + explicit SourceBackedChunkSupplier(std::shared_ptr source) + : source_(std::move(source)) { + CHECK(source_ != nullptr); } std::shared_ptr GetNextChunk() override { - auto input = input_->GetNextChunk(); - if (!input) { - return nullptr; - } - auto output = std::make_shared(); - for (size_t output_index = 0; output_index < columns_.size(); - ++output_index) { - const auto input_index = columns_[output_index]; - if (input_index < 0 || - static_cast(input_index) >= input->col_num()) { - THROW_INVALID_ARGUMENT_EXCEPTION( - "Chunk projection index is out of range: " + - std::to_string(input_index)); - } - output->set(static_cast(output_index), input->get(input_index)); - } - return output; + return OpenSupplier()->GetNextChunk(); } - int64_t RowNum() const override { return input_->RowNum(); } + int64_t RowNum() const override { return OpenSupplier()->RowNum(); } bool SupportsConcurrentGetNext() const override { - return input_->SupportsConcurrentGetNext(); + return OpenSupplier()->SupportsConcurrentGetNext(); + } + + void Cancel() override { + std::shared_ptr supplier; + { + std::lock_guard lock(mutex_); + cancelled_ = true; + supplier = supplier_; + } + if (supplier) { + supplier->Cancel(); + } } - void Cancel() override { input_->Cancel(); } + std::shared_ptr RepeatableSource() const override { + return source_; + } private: - std::shared_ptr input_; - std::vector columns_; + std::shared_ptr OpenSupplier() const { + std::lock_guard lock(mutex_); + if (!supplier_) { + supplier_ = source_->Open(); + if (!supplier_) { + THROW_INTERNAL_EXCEPTION("Data source returned a null supplier"); + } + if (cancelled_) { + supplier_->Cancel(); + } + } + return supplier_; + } + + std::shared_ptr source_; + mutable std::mutex mutex_; + mutable std::shared_ptr supplier_; + mutable bool cancelled_ = false; }; CsvReadConfig project_csv_config(const CsvReadConfig& config, @@ -1610,6 +1622,27 @@ CsvReadConfig project_csv_config(const CsvReadConfig& config, return projected; } +std::vector compose_projection( + const std::vector& source_columns, + const std::vector& output_columns) { + if (source_columns.empty()) { + return output_columns; + } + if (output_columns.empty()) { + return source_columns; + } + std::vector result; + result.reserve(output_columns.size()); + for (const auto column : output_columns) { + if (column < 0 || static_cast(column) >= source_columns.size()) { + THROW_INVALID_ARGUMENT_EXCEPTION( + "Projected source column is out of range: " + std::to_string(column)); + } + result.push_back(source_columns[static_cast(column)]); + } + return result; +} + class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { public: PartitionedCsvChunkSupplier(std::vector file_paths, @@ -1618,8 +1651,9 @@ class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { std::shared_ptr plan_cache) : file_paths_(std::move(file_paths)), config_(std::move(config)), - producer_count_(std::clamp(options.producer_count, 1, - hardware_worker_count())), + producer_count_(std::clamp( + options.producer_count, 1, + chunk_pipeline_detail::hardware_worker_count())), plan_cache_(std::move(plan_cache)), queue_(options.queue_capacity) {} @@ -1677,7 +1711,6 @@ class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { } void WorkerMain() { - const bool profile_stages = VLOG_IS_ON(1); try { while (!stop_.load(std::memory_order_acquire)) { const auto task_index = @@ -1689,30 +1722,10 @@ class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { CsvReadConfig range_config = config_; range_config.skip_rows = task.skip_rows; range_config.use_threads = false; - const auto open_start = profile_stages - ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; CsvSupplierRuntime runtime(task.file_path, range_config, CsvRowCountMode::kUnknown, task.range); - if (profile_stages) { - parse_time_ns_.fetch_add( - std::chrono::duration_cast( - std::chrono::steady_clock::now() - open_start) - .count(), - std::memory_order_relaxed); - } while (!stop_.load(std::memory_order_acquire)) { - const auto parse_start = - profile_stages ? std::chrono::steady_clock::now() - : std::chrono::steady_clock::time_point{}; auto chunk = runtime.get_next_chunk(); - if (profile_stages) { - parse_time_ns_.fetch_add( - std::chrono::duration_cast( - std::chrono::steady_clock::now() - parse_start) - .count(), - std::memory_order_relaxed); - } if (!chunk || !queue_.Push(std::move(chunk))) { break; } @@ -1720,8 +1733,6 @@ class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { } } catch (...) { SetError(std::current_exception()); } if (active_workers_.fetch_sub(1, std::memory_order_acq_rel) == 1) { - VLOG(1) << "Partitioned CSV producers: parse_worker_ms=" - << parse_time_ns_.load(std::memory_order_relaxed) / 1000000; queue_.Close(); } } @@ -1773,7 +1784,6 @@ class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { std::atomic active_workers_{0}; std::atomic stop_{false}; std::atomic has_error_{false}; - std::atomic parse_time_ns_{0}; mutable std::mutex error_mutex_; std::exception_ptr first_error_; }; @@ -1828,40 +1838,31 @@ class ChainedCsvChunkSupplier final : public IDataChunkSupplier { } // namespace -std::shared_ptr IDataChunkSource::Open( - const ChunkSourceOptions& options) const { - auto supplier = Open(); - if (!supplier || options.projected_columns.empty()) { - return supplier; +std::shared_ptr make_data_chunk_supplier( + std::shared_ptr source) { + if (!source) { + return nullptr; } - return std::make_shared( - std::move(supplier), options.projected_columns); + return std::make_shared(std::move(source)); } CSVChunkSource::CSVChunkSource(std::vector file_paths, - CsvReadConfig config) + CsvReadConfig config, + std::vector projected_columns) : file_paths_(std::move(file_paths)), config_(std::move(config)), + projected_columns_(std::move(projected_columns)), partition_plan_cache_(std::make_shared()) {} -std::shared_ptr CSVChunkSource::Open() const { - if (file_paths_.empty()) { - THROW_INVALID_ARGUMENT_EXCEPTION("CSV chunk source has no input paths"); - } - if (file_paths_.size() == 1) { - return std::make_shared(file_paths_.front(), config_, - CsvRowCountMode::kUnknown); - } - return std::make_shared(file_paths_, config_); -} - std::shared_ptr CSVChunkSource::Open( const ChunkSourceOptions& options) const { if (file_paths_.empty()) { THROW_INVALID_ARGUMENT_EXCEPTION("CSV chunk source has no input paths"); } - auto open_config = project_csv_config(config_, options.projected_columns); - if (open_config.use_threads && options.parallel_enabled && + auto open_config = project_csv_config( + config_, + compose_projection(projected_columns_, options.projected_columns)); + if (open_config.use_threads && options.producer_count > 0 && !options.preserve_order) { return std::make_shared( file_paths_, std::move(open_config), options, partition_plan_cache_); diff --git a/src/utils/io/read/csv/csv_reader.cc b/src/utils/io/read/csv/csv_reader.cc index fcef40efa..1c9bccb94 100644 --- a/src/utils/io/read/csv/csv_reader.cc +++ b/src/utils/io/read/csv/csv_reader.cc @@ -485,7 +485,8 @@ CsvReader::CsvReader(std::shared_ptr sharedState, CsvReader::~CsvReader() = default; -std::shared_ptr CsvReader::createChunkSource() { +std::shared_ptr CsvReader::createChunkSource( + std::vector projected_columns) { if (!sharedState_) { THROW_INVALID_ARGUMENT_EXCEPTION("SharedState is null"); } @@ -511,7 +512,8 @@ std::shared_ptr CsvReader::createChunkSource() { if (paths.empty()) { THROW_INVALID_ARGUMENT_EXCEPTION("No file paths provided"); } - return std::make_shared(paths, std::move(read_config)); + return std::make_shared(paths, std::move(read_config), + std::move(projected_columns)); } void CsvReader::read(std::shared_ptr /*localState*/, diff --git a/tests/storage/test_copy_temp.cc b/tests/storage/test_copy_temp.cc index f643e7274..8eeac087a 100644 --- a/tests/storage/test_copy_temp.cc +++ b/tests/storage/test_copy_temp.cc @@ -160,7 +160,7 @@ TEST_F(CopyTempTest, NodeDefaultPrimaryKey) { conn->Close(); } -TEST_F(CopyTempTest, PersistentCopyUsesForcedBulkPath) { +TEST_F(CopyTempTest, PersistentCopySelectsEdgeBuildPath) { auto conn = db_->Connect(); const std::string people = std::string(CSV_DIR) + "/people.csv"; const std::string edges = std::string(CSV_DIR) + "/edges.csv"; @@ -200,20 +200,19 @@ TEST_F(CopyTempTest, PersistentCopyUsesForcedBulkPath) { ASSERT_TRUE(dangling_count) << dangling_count.error().ToString(); EXPECT_EQ(dangling_count.value().response().row_count(), 1); - // The same terminal plan must retain the established materialized path - // when bulk build is explicitly disabled. - ASSERT_TRUE( - conn->Query("CREATE NODE TABLE Fallback(id INT64, name STRING, " - "age INT64, PRIMARY KEY(id));")); + // Disabling staged edge build keeps the terminal plan on normal BatchAdd. + ASSERT_TRUE(conn->Query( + "CREATE REL TABLE Fallback(FROM Person TO Person, weight DOUBLE);")); { ScopedEnvironmentVariable disable_bulk("NEUG_COPY_BULK_BUILD", "false"); auto fallback = - conn->Query("COPY Fallback FROM \"" + people + "\" (header = true);"); + conn->Query("COPY Fallback FROM \"" + edges + "\" (header = true);"); ASSERT_TRUE(fallback) << fallback.error().ToString(); } - auto fallback_count = conn->Query("MATCH (n:Fallback) RETURN n.id;"); + auto fallback_count = + conn->Query("MATCH (:Person)-[e:Fallback]->(:Person) RETURN e.weight;"); ASSERT_TRUE(fallback_count) << fallback_count.error().ToString(); - EXPECT_EQ(fallback_count.value().response().row_count(), 4); + EXPECT_EQ(fallback_count.value().response().row_count(), 3); conn->Close(); } diff --git a/tests/storage/test_edge_table.cc b/tests/storage/test_edge_table.cc index db838ff70..b18d4e7ce 100644 --- a/tests/storage/test_edge_table.cc +++ b/tests/storage/test_edge_table.cc @@ -90,6 +90,10 @@ class EdgeTableTest : public ::testing::Test { "person", "comment", "create_none", {neug::DataTypeId::kInt32}, {"data"}, neug::EdgeStrategy::kNone, neug::EdgeStrategy::kNone, true, true, std::nullopt, "edge label with no stored adjacency"); + schema_.AddEdgeLabel( + "person", "comment", "create_out_none", {neug::DataTypeId::kInt32}, + {"data"}, neug::EdgeStrategy::kNone, neug::EdgeStrategy::kMultiple, + true, true, std::nullopt, "edge label with only incoming adjacency"); src_label_ = schema_.get_vertex_label_id("person"); dst_label_ = schema_.get_vertex_label_id("comment"); edge_label_empty_ = schema_.get_edge_label_id("create0"); @@ -99,6 +103,7 @@ class EdgeTableTest : public ::testing::Test { edge_label_single_ = schema_.get_edge_label_id("create_single"); edge_label_single_both_ = schema_.get_edge_label_id("create_single_both"); edge_label_none_ = schema_.get_edge_label_id("create_none"); + edge_label_out_none_ = schema_.get_edge_label_id("create_out_none"); allocator_dir_ = "/tmp/edge_table_test_allocator_" + std::to_string(::getpid()) + "_"; ws.Open(temp_dir_.string()); @@ -155,11 +160,15 @@ class EdgeTableTest : public ::testing::Test { edge_table->BatchAddEdges(src_indexer, dst_indexer, supplier); } + void BatchBuild(std::shared_ptr source) { + edge_table->BatchAddEdges(src_indexer, dst_indexer, + make_data_chunk_supplier(std::move(source))); + } + void BatchBuild(std::vector> chunks, - int64_t estimated_bytes = -1) { - auto source = std::make_shared(std::move(chunks), - estimated_bytes); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, std::move(source)); + int64_t estimated_bytes = kForceBulkBuildBytes) { + BatchBuild(std::make_shared(std::move(chunks), + estimated_bytes)); } size_t ExpectedBatchInsertCapacity(size_t inserted_edge_num) const { @@ -270,7 +279,7 @@ class EdgeTableTest : public ::testing::Test { neug::Schema schema_; neug::label_t src_label_, dst_label_, edge_label_empty_, edge_label_int_, edge_label_str_, edge_label_str_int_, edge_label_single_, - edge_label_single_both_, edge_label_none_; + edge_label_single_both_, edge_label_none_, edge_label_out_none_; std::string allocator_dir_; const std::filesystem::path& temp_dir() const { return temp_dir_; } @@ -747,17 +756,15 @@ TEST_F(EdgeTableTest, BatchBuildEdgesBundledWithParallelEncoding) { InitIndexers(*ckp, kSrcNum, kDstNum); ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), kSrcNum, kDstNum); - EXPECT_TRUE(edge_table->CanBatchBuild()); auto source = std::make_shared(std::move(batches), 256LL * 1024 * 1024); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + BatchBuild(source); EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); EXPECT_EQ(source->OpenCount(), 2); ASSERT_EQ(source->OpenedProjections().size(), 2); EXPECT_EQ(source->OpenedProjections()[0], (std::vector{0, 1})); EXPECT_TRUE(source->OpenedProjections()[1].empty()); - EXPECT_FALSE(edge_table->CanBatchBuild()); std::vector output_srcs, output_dsts; OutputOutgoingEndpoints(output_srcs, output_dsts, neug::MAX_TIMESTAMP); ASSERT_EQ(output_srcs.size(), kEdgeNum); @@ -798,8 +805,8 @@ TEST_F(EdgeTableTest, BatchBuildEdgesParallelFillWithoutProperties) { ConstructEdgeTable(src_label_, dst_label_, edge_label_empty_); OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); auto source = std::make_shared(std::move(chunks), - 1024LL * 1024 * 1024); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + kForceBulkBuildBytes); + BatchBuild(source); EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); std::vector actual_srcs, actual_dsts; @@ -826,16 +833,12 @@ TEST_F(EdgeTableTest, BatchBuildEdgesHandlesEmptyInput) { ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); auto source = std::make_shared( - std::vector>{}, 1024LL * 1024 * 1024); + std::vector>{}, kForceBulkBuildBytes); - EXPECT_TRUE(edge_table->CanBatchBuild()); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + BatchBuild(source); EXPECT_EQ(source->OpenCount(), 2); EXPECT_EQ(edge_table->EdgeNum(), 0); - // CanBatchBuild() is edge-count based, so a successfully published empty - // graph intentionally remains eligible for another bulk build. - EXPECT_TRUE(edge_table->CanBatchBuild()); std::vector output_srcs, output_dsts; OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); EXPECT_TRUE(output_srcs.empty()); @@ -854,15 +857,47 @@ TEST_F(EdgeTableTest, BatchBuildEdgesSkipsSourceForNoAdjacency) { ConstructEdgeTable(src_label_, dst_label_, edge_label_none_); OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); auto source = std::make_shared(std::move(chunks), - 1024LL * 1024 * 1024); + kForceBulkBuildBytes); - ASSERT_TRUE(edge_table->CanBatchBuild()); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + BatchBuild(source); EXPECT_EQ(source->OpenCount(), 0); EXPECT_TRUE(source->OpenedProjections().empty()); EXPECT_EQ(edge_table->EdgeNum(), 0); - EXPECT_TRUE(edge_table->CanBatchBuild()); +} + +TEST_F(EdgeTableTest, SecondBatchFallsBackWhenIncomingCsrIsNotEmpty) { + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kVertexNum = 3; + InitIndexers(*ckp, kVertexNum, kVertexNum); + ConstructEdgeTable(src_label_, dst_label_, edge_label_out_none_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); + + auto first_chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{10}, 1)}); + auto first_source = std::make_shared( + std::move(first_chunks), kForceBulkBuildBytes); + BatchBuild(first_source); + EXPECT_EQ(first_source->OpenCount(), 2); + + auto second_chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{1}, 1), + split_column_to_chunks(std::vector{1}, 1), + split_column_to_chunks(std::vector{20}, 1)}); + auto second_source = std::make_shared( + std::move(second_chunks), kForceBulkBuildBytes); + BatchBuild(second_source); + EXPECT_EQ(second_source->OpenCount(), 1); + + std::vector incoming_srcs, incoming_dsts; + std::vector incoming_data; + OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); + OutputIncomingEdgeData(incoming_data, MAX_TIMESTAMP, 0); + EXPECT_EQ(incoming_srcs, (std::vector{0, 1})); + EXPECT_EQ(incoming_dsts, (std::vector{0, 1})); + EXPECT_EQ(incoming_data, (std::vector{10, 20})); } TEST_F(EdgeTableTest, BatchBuildEdgesParallelFillHandlesSupernodes) { @@ -883,8 +918,8 @@ TEST_F(EdgeTableTest, BatchBuildEdgesParallelFillHandlesSupernodes) { ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); auto source = std::make_shared(std::move(chunks), - 1024LL * 1024 * 1024); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + kForceBulkBuildBytes); + BatchBuild(source); EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); std::vector output_srcs, output_dsts; @@ -921,76 +956,6 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildFillsUniqueSlotsAcrossChunks) { std::atomic max_active{0}; }; - class ConcurrentChunkSupplier final : public IDataChunkSupplier { - public: - ConcurrentChunkSupplier(std::vector> chunks, - std::shared_ptr activity) - : chunks_(std::move(chunks)), activity_(std::move(activity)) {} - - std::shared_ptr GetNextChunk() override { - const auto active = - activity_->active.fetch_add(1, std::memory_order_relaxed) + 1; - auto max_active = activity_->max_active.load(std::memory_order_relaxed); - while (max_active < active && - !activity_->max_active.compare_exchange_weak( - max_active, active, std::memory_order_relaxed)) {} - std::this_thread::sleep_for(std::chrono::microseconds(100)); - const auto index = next_.fetch_add(1, std::memory_order_relaxed); - activity_->active.fetch_sub(1, std::memory_order_relaxed); - return index < chunks_.size() ? chunks_[index] : nullptr; - } - - int64_t RowNum() const override { - int64_t rows = 0; - for (const auto& chunk : chunks_) { - rows += static_cast(chunk->row_num()); - } - return rows; - } - - bool SupportsConcurrentGetNext() const override { return true; } - - private: - std::vector> chunks_; - std::shared_ptr activity_; - std::atomic next_{0}; - }; - - class ConcurrentChunkSource final : public IDataChunkSource { - public: - explicit ConcurrentChunkSource( - std::vector> chunks) - : chunks_(std::move(chunks)) {} - - std::shared_ptr Open() const override { - auto activity = std::make_shared(); - activities_.push_back(activity); - return std::make_shared(chunks_, activity); - } - - std::shared_ptr Open( - const ChunkSourceOptions& options) const override { - opened_options_.push_back(options); - return IDataChunkSource::Open(options); - } - - bool rewindable() const override { return true; } - int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } - - const std::vector& OpenedOptions() const { - return opened_options_; - } - - const std::vector>& Activities() const { - return activities_; - } - - private: - std::vector> chunks_; - mutable std::vector opened_options_; - mutable std::vector> activities_; - }; - auto ckp = make_checkpoint(workspace()); constexpr int64_t kSrcNum = 512; constexpr int64_t kDstNum = 8; @@ -1010,8 +975,34 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildFillsUniqueSlotsAcrossChunks) { InitIndexers(*ckp, kSrcNum, kDstNum); ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); OpenEdgeTableInMemory(ckp, CheckpointManifest(), kSrcNum, kDstNum); - auto source = std::make_shared(std::move(chunks)); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + auto shared_chunks = + std::make_shared>>( + std::move(chunks)); + std::vector> activities; + auto source = std::make_shared( + [shared_chunks, &activities, kSrcNum](const ChunkSourceOptions&, size_t) { + auto activity = std::make_shared(); + activities.push_back(activity); + auto next = std::make_shared>(0); + return std::make_shared( + [shared_chunks, activity, next] { + const auto active = + activity->active.fetch_add(1, std::memory_order_relaxed) + 1; + auto max_active = + activity->max_active.load(std::memory_order_relaxed); + while (max_active < active && + !activity->max_active.compare_exchange_weak( + max_active, active, std::memory_order_relaxed)) {} + std::this_thread::sleep_for(std::chrono::microseconds(100)); + const auto index = next->fetch_add(1, std::memory_order_relaxed); + activity->active.fetch_sub(1, std::memory_order_relaxed); + return index < shared_chunks->size() ? (*shared_chunks)[index] + : nullptr; + }, + kSrcNum, true); + }, + kForceBulkBuildBytes); + BatchBuild(source); ASSERT_EQ(source->OpenedOptions().size(), 2); EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); @@ -1019,9 +1010,9 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildFillsUniqueSlotsAcrossChunks) { EXPECT_EQ(source->OpenedOptions()[0].projected_columns, (std::vector{0, 1})); EXPECT_TRUE(source->OpenedOptions()[1].projected_columns.empty()); - ASSERT_EQ(source->Activities().size(), 2); - EXPECT_GT(source->Activities()[0]->max_active.load(), 1); - EXPECT_GT(source->Activities()[1]->max_active.load(), 1); + ASSERT_EQ(activities.size(), 2); + EXPECT_GT(activities[0]->max_active.load(), 1); + EXPECT_GT(activities[1]->max_active.load(), 1); EXPECT_EQ(edge_table->EdgeNum(), kSrcNum); std::vector output_srcs, output_dsts; @@ -1066,8 +1057,8 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildSkipsInvalidEndpoints) { split_column_to_chunks(dsts, 4), split_column_to_chunks(data, 4)}); auto source = std::make_shared(std::move(chunks), - 1024LL * 1024 * 1024); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + kForceBulkBuildBytes); + BatchBuild(source); EXPECT_EQ(source->OpenCount(), 2); EXPECT_EQ(edge_table->EdgeNum(), 2); @@ -1094,113 +1085,14 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildDetectsCrossWorkerDuplicate) { << "Cross-worker duplicate detection needs two consumer workers"; } - class OrderedChunkSupplier final : public IDataChunkSupplier { - public: - explicit OrderedChunkSupplier( - std::vector> chunks) - : chunks_(std::move(chunks)) {} - - std::shared_ptr GetNextChunk() override { - if (next_ == chunks_.size()) { - return nullptr; - } - return chunks_[next_++]; - } - - int64_t RowNum() const override { - int64_t rows = 0; - for (const auto& chunk : chunks_) { - rows += static_cast(chunk->row_num()); - } - return rows; - } - - private: - std::vector> chunks_; - size_t next_ = 0; - }; - - class BarrierChunkSupplier final : public IDataChunkSupplier { - public: - explicit BarrierChunkSupplier( - std::vector> chunks) - : chunks_(std::move(chunks)) {} - - std::shared_ptr GetNextChunk() override { - const auto index = next_.fetch_add(1, std::memory_order_relaxed); - if (index >= chunks_.size()) { - return nullptr; - } - if (index < 2) { - std::unique_lock lock(mutex_); - ++arrived_; - cv_.notify_all(); - cv_.wait(lock, [&] { return arrived_ >= 2 || cancelled_; }); - if (cancelled_) { - return nullptr; - } - } - return chunks_[index]; - } - - int64_t RowNum() const override { - int64_t rows = 0; - for (const auto& chunk : chunks_) { - rows += static_cast(chunk->row_num()); - } - return rows; - } - - bool SupportsConcurrentGetNext() const override { return true; } - - void Cancel() override { - { - std::lock_guard lock(mutex_); - cancelled_ = true; - } - cv_.notify_all(); - } - - private: - std::vector> chunks_; - std::atomic next_{0}; + struct BarrierState { + std::atomic next{0}; std::mutex mutex_; std::condition_variable cv_; size_t arrived_ = 0; bool cancelled_ = false; }; - class CrossWorkerDuplicateSource final : public IDataChunkSource { - public: - explicit CrossWorkerDuplicateSource( - std::vector> chunks) - : chunks_(std::move(chunks)) {} - - std::shared_ptr Open() const override { - return std::make_shared(chunks_); - } - - std::shared_ptr Open( - const ChunkSourceOptions& options) const override { - opened_options_.push_back(options); - if (options.preserve_order) { - return std::make_shared(chunks_); - } - return std::make_shared(chunks_); - } - - bool rewindable() const override { return true; } - int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } - - const std::vector& OpenedOptions() const { - return opened_options_; - } - - private: - std::vector> chunks_; - mutable std::vector opened_options_; - }; - auto ckp = make_checkpoint(workspace()); InitIndexers(*ckp, 1, 2); ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); @@ -1211,8 +1103,48 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildDetectsCrossWorkerDuplicate) { split_column_to_chunks(std::vector{0, 1}, 2), split_column_to_chunks(std::vector{10, 20}, 2)}); ASSERT_EQ(chunks.size(), 2); - auto source = std::make_shared(std::move(chunks)); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + auto shared_chunks = + std::make_shared>>( + std::move(chunks)); + auto source = std::make_shared( + [shared_chunks](const ChunkSourceOptions& options, size_t) { + if (options.preserve_order) { + return std::shared_ptr( + std::make_shared(*shared_chunks, true)); + } + auto state = std::make_shared(); + return std::shared_ptr( + std::make_shared( + [shared_chunks, state] { + const auto index = + state->next.fetch_add(1, std::memory_order_relaxed); + if (index >= shared_chunks->size()) { + return std::shared_ptr{}; + } + if (index < 2) { + std::unique_lock lock(state->mutex_); + ++state->arrived_; + state->cv_.notify_all(); + state->cv_.wait(lock, [&] { + return state->arrived_ >= 2 || state->cancelled_; + }); + if (state->cancelled_) { + return std::shared_ptr{}; + } + } + return (*shared_chunks)[index]; + }, + 2, true, + [state] { + { + std::lock_guard lock(state->mutex_); + state->cancelled_ = true; + } + state->cv_.notify_all(); + })); + }, + kForceBulkBuildBytes); + BatchBuild(source); ASSERT_EQ(source->OpenedOptions().size(), 2); EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); @@ -1228,66 +1160,6 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildDetectsCrossWorkerDuplicate) { } TEST_F(EdgeTableTest, SingleEdgeBulkBuildPreservesLastWriteOrderAcrossChunks) { - class OrderedChunkSupplier final : public IDataChunkSupplier { - public: - explicit OrderedChunkSupplier( - std::vector> chunks) - : chunks_(std::move(chunks)) {} - - std::shared_ptr GetNextChunk() override { - if (next_chunk_ == chunks_.size()) { - return nullptr; - } - return chunks_[next_chunk_++]; - } - - int64_t RowNum() const override { - int64_t rows = 0; - for (const auto& chunk : chunks_) { - rows += static_cast(chunk->row_num()); - } - return rows; - } - - private: - std::vector> chunks_; - size_t next_chunk_ = 0; - }; - - class OrderedChunkSource final : public IDataChunkSource { - public: - explicit OrderedChunkSource(std::vector> chunks) - : chunks_(std::move(chunks)) {} - - std::shared_ptr Open() const override { - ++open_count_; - return std::make_shared(chunks_); - } - - std::shared_ptr Open( - const ChunkSourceOptions& options) const override { - opened_projections_.push_back(options.projected_columns); - opened_preserve_order_.push_back(options.preserve_order); - return IDataChunkSource::Open(options); - } - - bool rewindable() const override { return true; } - int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } - size_t OpenCount() const { return open_count_; } - const std::vector>& OpenedProjections() const { - return opened_projections_; - } - const std::vector& OpenedPreserveOrder() const { - return opened_preserve_order_; - } - - private: - std::vector> chunks_; - mutable size_t open_count_ = 0; - mutable std::vector> opened_projections_; - mutable std::vector opened_preserve_order_; - }; - auto ckp = make_checkpoint(workspace()); InitIndexers(*ckp, 1, 3); ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); @@ -1300,8 +1172,15 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildPreservesLastWriteOrderAcrossChunks) { split_column_to_chunks(dsts, 3), split_column_to_chunks(data, 3)}); ASSERT_EQ(chunks.size(), 3); - auto source = std::make_shared(std::move(chunks)); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + auto shared_chunks = + std::make_shared>>( + std::move(chunks)); + auto source = std::make_shared( + [shared_chunks](const ChunkSourceOptions&, size_t) { + return std::make_shared(*shared_chunks, true); + }, + kForceBulkBuildBytes); + BatchBuild(source); std::vector output_srcs, output_dsts; std::vector output_data; @@ -1311,10 +1190,12 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildPreservesLastWriteOrderAcrossChunks) { ASSERT_EQ(output_dsts, (std::vector{2})); ASSERT_EQ(output_data, (std::vector{30})); EXPECT_EQ(source->OpenCount(), 2); - ASSERT_EQ(source->OpenedProjections().size(), 2); - EXPECT_EQ(source->OpenedProjections()[0], (std::vector{0, 1})); - EXPECT_TRUE(source->OpenedProjections()[1].empty()); - EXPECT_EQ(source->OpenedPreserveOrder(), (std::vector{false, true})); + ASSERT_EQ(source->OpenedOptions().size(), 2); + EXPECT_EQ(source->OpenedOptions()[0].projected_columns, + (std::vector{0, 1})); + EXPECT_TRUE(source->OpenedOptions()[1].projected_columns.empty()); + EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); + EXPECT_TRUE(source->OpenedOptions()[1].preserve_order); EXPECT_EQ(edge_table->EdgeNum(), 3); } @@ -1329,8 +1210,8 @@ TEST_F(EdgeTableTest, SingleOnlyBulkBuildPreservesBothDirections) { split_column_to_chunks(std::vector{0, 0, 1}, 1), split_column_to_chunks(std::vector{10, 20, 30}, 1)}); auto source = std::make_shared(std::move(chunks), - 1024LL * 1024 * 1024); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + kForceBulkBuildBytes); + BatchBuild(source); EXPECT_EQ(source->OpenCount(), 1); ASSERT_EQ(source->OpenedProjections().size(), 1); @@ -1355,34 +1236,6 @@ TEST_F(EdgeTableTest, SingleOnlyBulkBuildPreservesBothDirections) { } TEST_F(EdgeTableTest, BatchBuildEdgesFromPartitionedCsvInTwoPasses) { - class EstimatedBytesSource final : public IDataChunkSource { - public: - EstimatedBytesSource(std::shared_ptr source, - int64_t estimated_bytes) - : source_(std::move(source)), estimated_bytes_(estimated_bytes) {} - - std::shared_ptr Open() const override { - ++open_count_; - return source_->Open(); - } - - std::shared_ptr Open( - const ChunkSourceOptions& options) const override { - ++open_count_; - return source_->Open(options); - } - - bool rewindable() const override { return source_->rewindable(); } - int64_t EstimatedBytes() const override { return estimated_bytes_; } - bool ParallelEnabled() const override { return source_->ParallelEnabled(); } - size_t OpenCount() const { return open_count_; } - - private: - std::shared_ptr source_; - int64_t estimated_bytes_; - mutable size_t open_count_ = 0; - }; - auto ckp = make_checkpoint(workspace()); constexpr int64_t kSrcNum = 100; constexpr int64_t kDstNum = 80; @@ -1425,9 +1278,12 @@ TEST_F(EdgeTableTest, BatchBuildEdgesFromPartitionedCsvInTwoPasses) { // Force the large-file planner while keeping the fixture small. The wrapper // forwards both Open() calls to the same CSV source, so its cached partition // plan is reused by the count and fill passes. - auto source = std::make_shared(std::move(csv_source), - 1024LL * 1024 * 1024); - edge_table->BatchBuildEdges(src_indexer, dst_indexer, source); + auto source = std::make_shared( + [csv_source](const ChunkSourceOptions& options, size_t) { + return csv_source->Open(options); + }, + kForceBulkBuildBytes, csv_source->ParallelEnabled()); + BatchBuild(source); EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); EXPECT_EQ(source->OpenCount(), 2); @@ -1454,39 +1310,6 @@ TEST_F(EdgeTableTest, BatchBuildEdgesFromPartitionedCsvInTwoPasses) { } TEST_F(EdgeTableTest, FailedBatchBuildDoesNotPublishStagedCsr) { - class ThrowingSupplier final : public IDataChunkSupplier { - public: - explicit ThrowingSupplier(std::shared_ptr chunk) - : chunk_(std::move(chunk)) {} - - std::shared_ptr GetNextChunk() override { - if (chunk_) { - return std::exchange(chunk_, nullptr); - } - throw std::runtime_error("injected bulk edge source failure"); - } - - int64_t RowNum() const override { return 1; } - - private: - std::shared_ptr chunk_; - }; - - class ThrowingSource final : public IDataChunkSource { - public: - explicit ThrowingSource(std::shared_ptr chunk) - : chunk_(std::move(chunk)) {} - - std::shared_ptr Open() const override { - return std::make_shared(chunk_); - } - - bool rewindable() const override { return true; } - - private: - std::shared_ptr chunk_; - }; - auto ckp = make_checkpoint(workspace()); InitIndexers(*ckp, 1, 1); ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); @@ -1497,11 +1320,20 @@ TEST_F(EdgeTableTest, FailedBatchBuildDoesNotPublishStagedCsr) { split_column_to_chunks(std::vector{42}, 1)}); ASSERT_EQ(chunks.size(), 1); - EXPECT_THROW(edge_table->BatchBuildEdges( - src_indexer, dst_indexer, - std::make_shared(chunks.front())), - std::runtime_error); - EXPECT_TRUE(edge_table->CanBatchBuild()); + auto source = std::make_shared( + [chunk = chunks.front()](const ChunkSourceOptions&, size_t) { + auto remaining = std::make_shared>(chunk); + return std::make_shared( + [remaining] { + if (*remaining) { + return std::exchange(*remaining, std::shared_ptr{}); + } + throw std::runtime_error("injected bulk edge source failure"); + }, + 1); + }, + kForceBulkBuildBytes); + EXPECT_THROW(BatchBuild(source), std::runtime_error); EXPECT_EQ(edge_table->EdgeNum(), 0); } @@ -1512,72 +1344,6 @@ TEST_F(EdgeTableTest, << "Concurrent uniqueness cancellation needs two consumer workers"; } - class ConcurrentSingleCheckThrowingSupplier final - : public IDataChunkSupplier { - public: - ConcurrentSingleCheckThrowingSupplier(std::shared_ptr chunk, - std::atomic& cancel_count) - : chunk_(std::move(chunk)), cancel_count_(cancel_count) {} - - std::shared_ptr GetNextChunk() override { - const auto call = next_call_.fetch_add(1, std::memory_order_relaxed); - if (call == 0) { - return chunk_; - } - if (call == 1) { - throw std::runtime_error("injected concurrent uniqueness failure"); - } - return nullptr; - } - - int64_t RowNum() const override { return 1; } - bool SupportsConcurrentGetNext() const override { return true; } - - void Cancel() override { - cancel_count_.fetch_add(1, std::memory_order_relaxed); - } - - private: - std::shared_ptr chunk_; - std::atomic& cancel_count_; - std::atomic next_call_{0}; - }; - - class ConcurrentSingleCheckThrowingSource final : public IDataChunkSource { - public: - explicit ConcurrentSingleCheckThrowingSource( - std::shared_ptr chunk) - : chunk_(std::move(chunk)) {} - - std::shared_ptr Open() const override { - return std::make_shared( - chunk_, cancel_count_); - } - - std::shared_ptr Open( - const ChunkSourceOptions& options) const override { - ++open_count_; - opened_options_.push_back(options); - return Open(); - } - - bool rewindable() const override { return true; } - int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } - size_t OpenCount() const { return open_count_; } - size_t CancelCount() const { - return cancel_count_.load(std::memory_order_relaxed); - } - const std::vector& OpenedOptions() const { - return opened_options_; - } - - private: - std::shared_ptr chunk_; - mutable size_t open_count_ = 0; - mutable std::atomic cancel_count_{0}; - mutable std::vector opened_options_; - }; - auto ckp = make_checkpoint(workspace()); InitIndexers(*ckp, 1, 1); ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); @@ -1587,62 +1353,42 @@ TEST_F(EdgeTableTest, split_column_to_chunks(std::vector{0}, 1), split_column_to_chunks(std::vector{42}, 1)}); ASSERT_EQ(chunks.size(), 1); - auto source = - std::make_shared(chunks.front()); - - EXPECT_THROW(edge_table->BatchBuildEdges(src_indexer, dst_indexer, source), - std::runtime_error); + std::atomic cancel_count{0}; + auto source = std::make_shared( + [chunk = chunks.front(), &cancel_count](const ChunkSourceOptions&, + size_t) { + auto next_call = std::make_shared>(0); + return std::make_shared( + [chunk, next_call]() -> std::shared_ptr { + const auto call = + next_call->fetch_add(1, std::memory_order_relaxed); + if (call == 0) { + return chunk; + } + if (call == 1) { + throw std::runtime_error( + "injected concurrent uniqueness failure"); + } + return nullptr; + }, + 1, true, + [&cancel_count] { + cancel_count.fetch_add(1, std::memory_order_relaxed); + }); + }, + kForceBulkBuildBytes); + + EXPECT_THROW(BatchBuild(source), std::runtime_error); EXPECT_EQ(source->OpenCount(), 1); - EXPECT_EQ(source->CancelCount(), 1); + EXPECT_EQ(cancel_count.load(std::memory_order_relaxed), 1); ASSERT_EQ(source->OpenedOptions().size(), 1); EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); EXPECT_EQ(source->OpenedOptions()[0].projected_columns, (std::vector{0, 1})); - EXPECT_TRUE(edge_table->CanBatchBuild()); EXPECT_EQ(edge_table->EdgeNum(), 0); } TEST_F(EdgeTableTest, SecondPassFailureDoesNotPublishAllocatedCsr) { - class ThrowAfterChunkSupplier final : public IDataChunkSupplier { - public: - explicit ThrowAfterChunkSupplier(std::shared_ptr chunk) - : chunk_(std::move(chunk)) {} - - std::shared_ptr GetNextChunk() override { - if (chunk_) { - return std::exchange(chunk_, nullptr); - } - throw std::runtime_error("injected second-pass failure"); - } - - int64_t RowNum() const override { return 1; } - - private: - std::shared_ptr chunk_; - }; - - class SecondPassThrowingSource final : public IDataChunkSource { - public: - explicit SecondPassThrowingSource(std::shared_ptr chunk) - : chunk_(std::move(chunk)) {} - - std::shared_ptr Open() const override { - ++open_count_; - if (open_count_ == 1) { - return std::make_shared( - std::vector>{chunk_}); - } - return std::make_shared(chunk_); - } - - bool rewindable() const override { return true; } - size_t OpenCount() const { return open_count_; } - - private: - std::shared_ptr chunk_; - mutable size_t open_count_ = 0; - }; - auto ckp = make_checkpoint(workspace()); InitIndexers(*ckp, 1, 1); ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); @@ -1652,12 +1398,28 @@ TEST_F(EdgeTableTest, SecondPassFailureDoesNotPublishAllocatedCsr) { split_column_to_chunks(std::vector{0}, 1), split_column_to_chunks(std::vector{42}, 1)}); ASSERT_EQ(chunks.size(), 1); - auto source = std::make_shared(chunks.front()); - - EXPECT_THROW(edge_table->BatchBuildEdges(src_indexer, dst_indexer, source), - std::runtime_error); + auto source = std::make_shared( + [chunk = chunks.front()](const ChunkSourceOptions&, size_t open_index) { + if (open_index == 0) { + return std::shared_ptr( + std::make_shared( + std::vector>{chunk})); + } + auto remaining = std::make_shared>(chunk); + return std::shared_ptr( + std::make_shared( + [remaining]() -> std::shared_ptr { + if (*remaining) { + return std::exchange(*remaining, nullptr); + } + throw std::runtime_error("injected second-pass failure"); + }, + 1)); + }, + kForceBulkBuildBytes); + + EXPECT_THROW(BatchBuild(source), std::runtime_error); EXPECT_EQ(source->OpenCount(), 2); - EXPECT_TRUE(edge_table->CanBatchBuild()); EXPECT_EQ(edge_table->EdgeNum(), 0); } @@ -1667,70 +1429,6 @@ TEST_F(EdgeTableTest, GTEST_SKIP() << "Concurrent supplier cancellation requires two workers"; } - class ConcurrentThrowingSupplier final : public IDataChunkSupplier { - public: - ConcurrentThrowingSupplier(std::shared_ptr chunk, - std::atomic& cancel_count) - : chunk_(std::move(chunk)), cancel_count_(cancel_count) {} - - std::shared_ptr GetNextChunk() override { - const auto call = next_call_.fetch_add(1, std::memory_order_relaxed); - if (call == 0) { - return chunk_; - } - if (call == 1) { - throw std::runtime_error("injected concurrent fill failure"); - } - return nullptr; - } - - int64_t RowNum() const override { return 1; } - bool SupportsConcurrentGetNext() const override { return true; } - - void Cancel() override { - cancel_count_.fetch_add(1, std::memory_order_relaxed); - } - - private: - std::shared_ptr chunk_; - std::atomic& cancel_count_; - std::atomic next_call_{0}; - }; - - class ConcurrentSecondPassThrowingSource final : public IDataChunkSource { - public: - explicit ConcurrentSecondPassThrowingSource( - std::shared_ptr chunk) - : chunk_(std::move(chunk)) {} - - std::shared_ptr Open() const override { - return Open(ChunkSourceOptions()); - } - - std::shared_ptr Open( - const ChunkSourceOptions& /*options*/) const override { - ++open_count_; - if (open_count_ == 1) { - return std::make_shared( - std::vector>{chunk_}); - } - return std::make_shared(chunk_, - cancel_count_); - } - - bool rewindable() const override { return true; } - int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } - size_t OpenCount() const { return open_count_; } - size_t CancelCount() const { - return cancel_count_.load(std::memory_order_relaxed); - } - - private: - std::shared_ptr chunk_; - mutable size_t open_count_ = 0; - mutable std::atomic cancel_count_{0}; - }; - auto ckp = make_checkpoint(workspace()); InitIndexers(*ckp, 1, 1); ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); @@ -1740,14 +1438,40 @@ TEST_F(EdgeTableTest, split_column_to_chunks(std::vector{0}, 1), split_column_to_chunks(std::vector{42}, 1)}); ASSERT_EQ(chunks.size(), 1); - auto source = - std::make_shared(chunks.front()); - - EXPECT_THROW(edge_table->BatchBuildEdges(src_indexer, dst_indexer, source), - std::runtime_error); + std::atomic cancel_count{0}; + auto source = std::make_shared( + [chunk = chunks.front(), &cancel_count](const ChunkSourceOptions&, + size_t open_index) { + if (open_index == 0) { + return std::shared_ptr( + std::make_shared( + std::vector>{chunk})); + } + auto next_call = std::make_shared>(0); + return std::shared_ptr( + std::make_shared( + [chunk, next_call]() -> std::shared_ptr { + const auto call = + next_call->fetch_add(1, std::memory_order_relaxed); + if (call == 0) { + return chunk; + } + if (call == 1) { + throw std::runtime_error( + "injected concurrent fill failure"); + } + return nullptr; + }, + 1, true, + [&cancel_count] { + cancel_count.fetch_add(1, std::memory_order_relaxed); + })); + }, + kForceBulkBuildBytes); + + EXPECT_THROW(BatchBuild(source), std::runtime_error); EXPECT_EQ(source->OpenCount(), 2); - EXPECT_EQ(source->CancelCount(), 1); - EXPECT_TRUE(edge_table->CanBatchBuild()); + EXPECT_EQ(cancel_count.load(std::memory_order_relaxed), 1); EXPECT_EQ(edge_table->EdgeNum(), 0); } @@ -1757,84 +1481,13 @@ TEST_F(EdgeTableTest, GTEST_SKIP() << "Bounded pipeline cancellation requires two consumers"; } - class BlockingAfterChunkSupplier final : public IDataChunkSupplier { - public: - explicit BlockingAfterChunkSupplier(std::shared_ptr chunk) - : chunk_(std::move(chunk)) {} - - std::shared_ptr GetNextChunk() override { - if (chunk_) { - return std::exchange(chunk_, nullptr); - } - std::unique_lock lock(mutex_); - if (!cancelled_cv_.wait_for(lock, std::chrono::seconds(2), - [&] { return cancelled_; })) { - timed_out_.store(true, std::memory_order_relaxed); - } - return nullptr; - } - - int64_t RowNum() const override { return 1; } - - void Cancel() override { - { - std::lock_guard lock(mutex_); - cancelled_ = true; - } - cancel_count_.fetch_add(1, std::memory_order_relaxed); - cancelled_cv_.notify_all(); - } - - size_t CancelCount() const { - return cancel_count_.load(std::memory_order_relaxed); - } - bool TimedOut() const { return timed_out_.load(std::memory_order_relaxed); } - - private: - std::shared_ptr chunk_; - mutable std::mutex mutex_; - std::condition_variable cancelled_cv_; - bool cancelled_ = false; - std::atomic cancel_count_{0}; - std::atomic timed_out_{false}; - }; - - class NonConcurrentSecondPassFailureSource final : public IDataChunkSource { - public: - NonConcurrentSecondPassFailureSource( - std::shared_ptr valid_chunk, - std::shared_ptr invalid_chunk) - : valid_chunk_(std::move(valid_chunk)), - invalid_chunk_(std::move(invalid_chunk)) {} - - std::shared_ptr Open() const override { - ++open_count_; - if (open_count_ == 1) { - return std::make_shared( - std::vector>{valid_chunk_}); - } - second_pass_supplier_ = - std::make_shared(invalid_chunk_); - return second_pass_supplier_; - } - - bool rewindable() const override { return true; } - int64_t EstimatedBytes() const override { return 1024LL * 1024 * 1024; } - size_t OpenCount() const { return open_count_; } - size_t CancelCount() const { - CHECK(second_pass_supplier_ != nullptr); - return second_pass_supplier_->CancelCount(); - } - bool TimedOut() const { - CHECK(second_pass_supplier_ != nullptr); - return second_pass_supplier_->TimedOut(); - } - - private: - std::shared_ptr valid_chunk_; - std::shared_ptr invalid_chunk_; - mutable size_t open_count_ = 0; - mutable std::shared_ptr second_pass_supplier_; + struct BlockingSupplierState { + std::shared_ptr chunk; + std::mutex mutex; + std::condition_variable cancelled_cv; + bool cancelled = false; + std::atomic cancel_count{0}; + std::atomic timed_out{false}; }; auto ckp = make_checkpoint(workspace()); @@ -1854,15 +1507,49 @@ TEST_F(EdgeTableTest, split_column_to_chunks(std::vector{42}, 1)}); ASSERT_EQ(valid_chunks.size(), 1); ASSERT_EQ(invalid_chunks.size(), 1); - auto source = std::make_shared( - valid_chunks.front(), invalid_chunks.front()); - - EXPECT_THROW(edge_table->BatchBuildEdges(src_indexer, dst_indexer, source), - std::exception); + std::shared_ptr second_pass_state; + auto source = std::make_shared( + [valid_chunk = valid_chunks.front(), + invalid_chunk = invalid_chunks.front(), + &second_pass_state](const ChunkSourceOptions&, size_t open_index) { + if (open_index == 0) { + return std::shared_ptr( + std::make_shared( + std::vector>{valid_chunk})); + } + second_pass_state = std::make_shared(); + second_pass_state->chunk = invalid_chunk; + return std::shared_ptr( + std::make_shared( + [state = second_pass_state] { + if (state->chunk) { + return std::exchange(state->chunk, nullptr); + } + std::unique_lock lock(state->mutex); + if (!state->cancelled_cv.wait_for( + lock, std::chrono::seconds(2), + [&] { return state->cancelled; })) { + state->timed_out.store(true, std::memory_order_relaxed); + } + return std::shared_ptr{}; + }, + 1, false, + [state = second_pass_state] { + { + std::lock_guard lock(state->mutex); + state->cancelled = true; + } + state->cancel_count.fetch_add(1, std::memory_order_relaxed); + state->cancelled_cv.notify_all(); + })); + }, + kForceBulkBuildBytes); + + EXPECT_THROW(BatchBuild(source), std::exception); EXPECT_EQ(source->OpenCount(), 2); - EXPECT_EQ(source->CancelCount(), 1); - EXPECT_FALSE(source->TimedOut()); - EXPECT_TRUE(edge_table->CanBatchBuild()); + ASSERT_NE(second_pass_state, nullptr); + EXPECT_EQ(second_pass_state->cancel_count.load(std::memory_order_relaxed), 1); + EXPECT_FALSE(second_pass_state->timed_out.load(std::memory_order_relaxed)); EXPECT_EQ(edge_table->EdgeNum(), 0); } @@ -1940,7 +1627,10 @@ TEST_F(EdgeTableTest, TestBatchAddEdgesUnbundled) { this->OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), src_num, dst_num); this->ExpectUnbundledStats(0, 0); - this->BatchInsert(std::move(batches)); + auto source = std::make_shared(std::move(batches), + kForceBulkBuildBytes); + this->BatchBuild(source); + EXPECT_EQ(source->OpenCount(), 1); EXPECT_EQ(this->edge_table->EdgeNum(), edge_num); this->ExpectUnbundledStats(edge_num, ExpectedBatchInsertCapacity(edge_num)); diff --git a/tests/storage/test_mutable_csr.cc b/tests/storage/test_mutable_csr.cc index d4eb98eb4..c5a3b456b 100644 --- a/tests/storage/test_mutable_csr.cc +++ b/tests/storage/test_mutable_csr.cc @@ -572,13 +572,6 @@ TEST(MutableCsrBulkBuildAccessTest, ReservesTwentyPercentPerVertex) { writer.CountConcurrent(1, 2); writer.AllocateFromCounts(); - EXPECT_EQ(writer.ExpectedDegree(0), 20); - EXPECT_EQ(writer.ExpectedDegree(1), 2); - EXPECT_EQ(writer.ExpectedDegree(2), 0); - EXPECT_EQ(writer.ReservedCapacityForVertex(0), 24); - EXPECT_EQ(writer.ReservedCapacityForVertex(1), 3); - EXPECT_EQ(writer.ReservedCapacityForVertex(2), 0); - std::vector fillers; for (int thread = 0; thread < 4; ++thread) { fillers.emplace_back([&writer] { diff --git a/tests/storage/test_property_graph.cc b/tests/storage/test_property_graph.cc index cd087e08a..0c0396e96 100644 --- a/tests/storage/test_property_graph.cc +++ b/tests/storage/test_property_graph.cc @@ -144,4 +144,4 @@ TEST_F(PropertyGraphTest, TestOpenAndBulkInsert) { } } -} // namespace neug \ No newline at end of file +} // namespace neug diff --git a/tests/storage/test_vertex_table.cc b/tests/storage/test_vertex_table.cc index 1a5b8b49d..ae0664aec 100644 --- a/tests/storage/test_vertex_table.cc +++ b/tests/storage/test_vertex_table.cc @@ -707,22 +707,20 @@ TEST_F(VertexTableTest, VertexTableResizeTest) { } } -TEST_F(VertexTableTest, BatchBuildVerticesFromRepeatableSource) { +TEST_F(VertexTableTest, InsertVerticesFromRepeatableSource) { neug::VertexTable table(schema_.get_vertex_schema(v_label_id_)); auto ckp = make_checkpoint(Workspace()); OpenVertexTableLegacy(table, ckp, neug::CheckpointManifest(), memory_level_); constexpr size_t kVertexNum = 4097; - auto source = - std::make_shared(generate_data_chunks(kVertexNum)); - EXPECT_TRUE(table.CanBatchBuild()); - table.BatchBuildVertices(source); + auto source = std::make_shared( + generate_data_chunks(kVertexNum), kForceBulkBuildBytes); + table.insert_vertices(make_data_chunk_supplier(source)); EXPECT_EQ(table.VertexNum(), kVertexNum); EXPECT_EQ(table.LidNum(), kVertexNum); EXPECT_EQ(table.Capacity(), kVertexNum + kVertexNum / 4); EXPECT_EQ(source->OpenCount(), 1); - EXPECT_FALSE(table.CanBatchBuild()); for (int64_t oid : {int64_t{0}, int64_t{127}, int64_t{4096}}) { neug::vid_t lid; EXPECT_TRUE(table.get_index(neug::Value::INT64(oid), lid)); diff --git a/tests/unittest/utils.h b/tests/unittest/utils.h index ea4113f04..2ec50a8fc 100644 --- a/tests/unittest/utils.h +++ b/tests/unittest/utils.h @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -43,14 +44,20 @@ #include "neug/utils/property/table.h" #include "neug/utils/property/types.h" +inline constexpr int64_t kForceBulkBuildBytes = 1LL << 30; + class GeneratedChunkSupplier : public neug::IDataChunkSupplier { public: explicit GeneratedChunkSupplier( - std::vector>&& chunks) - : chunks_(std::move(chunks)) {} + std::vector> chunks, + bool preserve_order = false) + : chunks_(std::move(chunks)), preserve_order_(preserve_order) {} ~GeneratedChunkSupplier() override = default; std::shared_ptr GetNextChunk() override { + if (preserve_order_) { + return next_ < chunks_.size() ? chunks_[next_++] : nullptr; + } if (chunks_.empty()) { return nullptr; } @@ -71,6 +78,38 @@ class GeneratedChunkSupplier : public neug::IDataChunkSupplier { private: std::vector> chunks_; + bool preserve_order_; + size_t next_ = 0; +}; + +class TestChunkSupplier final : public neug::IDataChunkSupplier { + public: + using NextChunk = std::function()>; + using CancelCallback = std::function; + + TestChunkSupplier(NextChunk next_chunk, int64_t row_num, + bool concurrent = false, CancelCallback cancel = {}) + : next_chunk_(std::move(next_chunk)), + cancel_(std::move(cancel)), + row_num_(row_num), + concurrent_(concurrent) {} + + std::shared_ptr GetNextChunk() override { + return next_chunk_(); + } + int64_t RowNum() const override { return row_num_; } + bool SupportsConcurrentGetNext() const override { return concurrent_; } + void Cancel() override { + if (cancel_) { + cancel_(); + } + } + + private: + NextChunk next_chunk_; + CancelCallback cancel_; + int64_t row_num_; + bool concurrent_; }; class GeneratedChunkSource final : public neug::IDataChunkSource { @@ -80,20 +119,29 @@ class GeneratedChunkSource final : public neug::IDataChunkSource { int64_t estimated_bytes = -1) : chunks_(std::move(chunks)), estimated_bytes_(estimated_bytes) {} - std::shared_ptr Open() const override { - ++open_count_; - return std::make_shared( - std::vector>(chunks_)); - } - std::shared_ptr Open( const neug::ChunkSourceOptions& options) const override { + ++open_count_; opened_projections_.push_back(options.projected_columns); - return neug::IDataChunkSource::Open(options); + auto chunks = chunks_; + if (!options.projected_columns.empty()) { + for (auto& chunk : chunks) { + auto projected = std::make_shared(); + for (size_t output = 0; output < options.projected_columns.size(); + ++output) { + const auto input = options.projected_columns[output]; + if (input < 0 || static_cast(input) >= chunk->col_num()) { + throw std::out_of_range("Chunk projection index is out of range"); + } + projected->set(static_cast(output), chunk->get(input)); + } + chunk = std::move(projected); + } + } + return std::make_shared(std::move(chunks), + options.preserve_order); } - bool rewindable() const override { return true; } - int64_t EstimatedBytes() const override { return estimated_bytes_; } size_t OpenCount() const { return open_count_; } @@ -109,6 +157,39 @@ class GeneratedChunkSource final : public neug::IDataChunkSource { mutable std::vector> opened_projections_; }; +class TestChunkSource final : public neug::IDataChunkSource { + public: + using Factory = std::function( + const neug::ChunkSourceOptions&, size_t)>; + + explicit TestChunkSource(Factory factory, int64_t estimated_bytes = -1, + bool parallel_enabled = true) + : factory_(std::move(factory)), + estimated_bytes_(estimated_bytes), + parallel_enabled_(parallel_enabled) {} + + std::shared_ptr Open( + const neug::ChunkSourceOptions& options) const override { + opened_options_.push_back(options); + return factory_(options, open_count_++); + } + + int64_t EstimatedBytes() const override { return estimated_bytes_; } + bool ParallelEnabled() const override { return parallel_enabled_; } + + size_t OpenCount() const { return open_count_; } + const std::vector& OpenedOptions() const { + return opened_options_; + } + + private: + Factory factory_; + int64_t estimated_bytes_; + bool parallel_enabled_; + mutable size_t open_count_ = 0; + mutable std::vector opened_options_; +}; + template std::shared_ptr build_value_column_slice( const std::vector& data, size_t begin, size_t end) { diff --git a/tests/utils/test_reader.cc b/tests/utils/test_reader.cc index c35b85275..8e43a3794 100644 --- a/tests/utils/test_reader.cc +++ b/tests/utils/test_reader.cc @@ -16,7 +16,7 @@ #include "test_reader.h" #include -#include +#include #include "neug/storages/loader/chunk_pipeline_utils.h" #include "neug/storages/loader/loader_utils.h" @@ -24,6 +24,21 @@ namespace neug { namespace test { +namespace { + +ChunkSourceOptions parallel_source_options( + int32_t producer_count = 4, size_t queue_capacity = 8, + std::vector projected_columns = {}) { + return { + .producer_count = producer_count, + .queue_capacity = queue_capacity, + .preserve_order = false, + .projected_columns = std::move(projected_columns), + }; +} + +} // namespace + // Test 1: Basic CSV reading with default options TEST_F(ReaderTest, TestBasicCsvRead) { // Create test CSV file @@ -63,7 +78,6 @@ TEST_F(ReaderTest, CsvChunkSourceCanBeReopened) { auto source = reader->createChunkSource(); ASSERT_NE(source, nullptr); - EXPECT_TRUE(source->rewindable()); for (int pass = 0; pass < 2; ++pass) { auto supplier = source->Open(); ASSERT_NE(supplier, nullptr); @@ -105,12 +119,7 @@ TEST_F(ReaderTest, CsvChunkSourcePartitionsQuotedRecords) { }; auto expected = read_rows(source->Open()); - ChunkSourceOptions options; - options.parallel_enabled = true; - options.producer_count = 4; - options.queue_capacity = 8; - options.preserve_order = false; - auto partitioned = source->Open(options); + auto partitioned = source->Open(parallel_source_options()); ASSERT_NE(partitioned, nullptr); EXPECT_TRUE(partitioned->SupportsConcurrentGetNext()); EXPECT_EQ(partitioned->RowNum(), 6); // Header is an intentional overcount. @@ -145,12 +154,7 @@ TEST_F(ReaderTest, PartitionedCsvCarriesSkipAcrossRanges) { CSVChunkSource source( {std::string(ARROW_READER_TEST_DIR) + "/partition-skip.csv"}, config); - ChunkSourceOptions options; - options.parallel_enabled = true; - options.producer_count = 4; - options.queue_capacity = 4; - options.preserve_order = false; - auto supplier = source.Open(options); + auto supplier = source.Open(parallel_source_options(4, 4)); ASSERT_NE(supplier, nullptr); std::vector ids; @@ -195,12 +199,7 @@ TEST_F(ReaderTest, PartitionPlannerMatchesParserWhenDoubleQuoteIsFalse) { }; const auto expected = read_ids(source.Open()); - ChunkSourceOptions options; - options.parallel_enabled = true; - options.producer_count = 4; - options.queue_capacity = 4; - options.preserve_order = false; - EXPECT_EQ(read_ids(source.Open(options)), expected); + EXPECT_EQ(read_ids(source.Open(parallel_source_options(4, 4))), expected); EXPECT_EQ(expected, (std::vector{0, 1, 2})); } @@ -220,13 +219,7 @@ TEST_F(ReaderTest, CsvChunkSourcePushesProjectionIntoPartitionedParsing) { {std::string(ARROW_READER_TEST_DIR) + "/partition-projection.csv"}, config); - ChunkSourceOptions options; - options.parallel_enabled = true; - options.producer_count = 2; - options.queue_capacity = 2; - options.preserve_order = false; - options.projected_columns = {2, 0}; - auto supplier = source.Open(options); + auto supplier = source.Open(parallel_source_options(2, 2, {2, 0})); ASSERT_NE(supplier, nullptr); std::vector> rows; @@ -241,43 +234,6 @@ TEST_F(ReaderTest, CsvChunkSourcePushesProjectionIntoPartitionedParsing) { EXPECT_EQ(rows, (std::vector>{{10, 1}, {20, 2}})); } -TEST_F(ReaderTest, ChunkPipelineAllocationSharesHardwareBudget) { - constexpr int64_t kGiB = 1024LL * 1024 * 1024; - auto allocation = resolve_chunk_pipeline_allocation(kGiB, true, false, 16); - EXPECT_TRUE(allocation.parallel_enabled); - EXPECT_EQ(allocation.producer_count, 8); - EXPECT_EQ(allocation.consumer_count, 8); - EXPECT_LE(allocation.producer_count + allocation.consumer_count, 16); - EXPECT_EQ(allocation.queue_capacity, 16); - - auto ordered = resolve_chunk_pipeline_allocation(kGiB, true, true, 16); - EXPECT_FALSE(ordered.parallel_enabled); - EXPECT_EQ(ordered.producer_count, 1); - EXPECT_EQ(ordered.consumer_count, 1); - - auto disabled = resolve_chunk_pipeline_allocation(kGiB, false, false, 16); - EXPECT_FALSE(disabled.parallel_enabled); - EXPECT_EQ(disabled.producer_count, 1); - EXPECT_EQ(disabled.consumer_count, 1); - - auto small = - resolve_chunk_pipeline_allocation(128LL * 1024 * 1024, true, false, 16); - EXPECT_FALSE(small.parallel_enabled); - EXPECT_EQ(small.producer_count, 1); - EXPECT_EQ(small.consumer_count, 1); - - auto two_workers = resolve_chunk_pipeline_allocation(kGiB, true, false, 2); - EXPECT_TRUE(two_workers.parallel_enabled); - EXPECT_EQ(two_workers.producer_count, 1); - EXPECT_EQ(two_workers.consumer_count, 1); - - auto maximum_size = resolve_chunk_pipeline_allocation( - std::numeric_limits::max(), true, false, 16); - EXPECT_TRUE(maximum_size.parallel_enabled); - EXPECT_EQ(maximum_size.producer_count, 8); - EXPECT_EQ(maximum_size.consumer_count, 8); -} - TEST_F(ReaderTest, CsvChunkSourceHonorsParallelFalse) { createCsvFile("serial.csv", "id|name\n1|Alice\n2|Bob\n"); std::vector column_names = {"id", "name"}; @@ -290,12 +246,7 @@ TEST_F(ReaderTest, CsvChunkSourceHonorsParallelFalse) { ASSERT_NE(source, nullptr); EXPECT_FALSE(source->ParallelEnabled()); - ChunkSourceOptions options; - options.parallel_enabled = true; - options.producer_count = 4; - options.queue_capacity = 8; - options.preserve_order = false; - auto supplier = source->Open(options); + auto supplier = source->Open(parallel_source_options()); ASSERT_NE(supplier, nullptr); EXPECT_FALSE(supplier->SupportsConcurrentGetNext()); } @@ -318,12 +269,7 @@ TEST_F(ReaderTest, PartitionedCsvSkipsHeaderForEveryFile) { }; CSVChunkSource source({path("part-a.csv"), path("part-b.csv")}, config); - ChunkSourceOptions options; - options.parallel_enabled = true; - options.producer_count = 4; - options.queue_capacity = 8; - options.preserve_order = false; - auto supplier = source.Open(options); + auto supplier = source.Open(parallel_source_options()); ASSERT_NE(supplier, nullptr); EXPECT_EQ(supplier->RowNum(), 6); // Two headers are counted as reserve hints. @@ -354,12 +300,7 @@ TEST_F(ReaderTest, PartitionedCsvPropagatesProducerErrors) { CSVChunkSource source( {std::string(ARROW_READER_TEST_DIR) + "/partition-error.csv"}, config); - ChunkSourceOptions options; - options.parallel_enabled = true; - options.producer_count = 4; - options.queue_capacity = 2; - options.preserve_order = false; - auto supplier = source.Open(options); + auto supplier = source.Open(parallel_source_options(4, 2)); ASSERT_NE(supplier, nullptr); EXPECT_ANY_THROW({ while (supplier->GetNextChunk()) {} From fc674ed121e30c95eeb9f45ee6d1606f60fec081 Mon Sep 17 00:00:00 2001 From: "xiaolei.zl" Date: Fri, 17 Jul 2026 17:53:20 +0800 Subject: [PATCH 3/8] refactor and reduce code --- include/neug/storages/csr/mutable_csr.h | 254 +---- include/neug/storages/graph/edge_table.h | 3 +- include/neug/storages/loader/loader_utils.h | 2 + .../execute/ops/batch/batch_insert_edge.cc | 32 +- src/storages/csr/mutable_csr.cc | 13 +- src/storages/graph/edge_table.cc | 658 +----------- .../loader/bundled_edge_csr_loader.cc | 909 +++++++++++++++++ src/storages/loader/bundled_edge_csr_loader.h | 55 ++ src/storages/loader/loader_utils.cc | 578 ++++++----- tests/storage/test_edge_table.cc | 934 +++++++----------- tests/storage/test_mutable_csr.cc | 62 -- tests/utils/test_reader.cc | 392 ++++---- tests/utils/test_reader.h | 5 - 13 files changed, 1882 insertions(+), 2015 deletions(-) create mode 100644 src/storages/loader/bundled_edge_csr_loader.cc create mode 100644 src/storages/loader/bundled_edge_csr_loader.h diff --git a/include/neug/storages/csr/mutable_csr.h b/include/neug/storages/csr/mutable_csr.h index 9fa2e7726..881213686 100644 --- a/include/neug/storages/csr/mutable_csr.h +++ b/include/neug/storages/csr/mutable_csr.h @@ -23,7 +23,6 @@ #include #include #include -#include #include #include #include @@ -32,7 +31,6 @@ #include #include -#include "neug/config.h" #include "neug/storages/allocators.h" #include "neug/storages/container/i_container.h" #include "neug/storages/csr/csr_base.h" @@ -45,6 +43,10 @@ namespace neug { +namespace internal { +class BundledEdgeCsrLoader; +} + // std::atomic must have the same size as int on supported platforms // so that the degree_list buffer (persisted as int[]) can be safely // reinterpreted as atomic[] for concurrent access. @@ -52,11 +54,11 @@ static_assert( sizeof(std::atomic) == sizeof(int), "atomic must have the same size as int on supported platforms"); -template -class MutableCsrBulkBuildAccess; +namespace mutable_csr_detail { -template -class SingleMutableCsrBulkBuildAccess; +int capacity_with_reserve(int degree); + +} // namespace mutable_csr_detail template class MutableCsr : public TypedCsrBase { @@ -243,7 +245,7 @@ class MutableCsr : public TypedCsrBase { } private: - friend class MutableCsrBulkBuildAccess; + friend class internal::BundledEdgeCsrLoader; std::unique_ptr locks_; std::shared_ptr adj_list_buffer_; @@ -255,16 +257,6 @@ class MutableCsr : public TypedCsrBase { CsrPrefetchPolicy prefetch_policy_; void refresh_prefetch_policy(); - static int reserved_capacity(int degree) { - CHECK_GE(degree, 0); - if (degree == 0) { - return 0; - } - const auto reserved = - std::ceil(degree * NeugDBConfig::DEFAULT_RESERVE_RATIO); - CHECK_LE(reserved, static_cast(std::numeric_limits::max())); - return static_cast(reserved); - } size_t vertex_capacity() const { if (!degree_list_) { @@ -402,7 +394,7 @@ class SingleMutableCsr : public TypedCsrBase { } private: - friend class SingleMutableCsrBulkBuildAccess; + friend class internal::BundledEdgeCsrLoader; std::shared_ptr nbr_list_; std::atomic edge_num_{0}; @@ -502,230 +494,4 @@ class EmptyCsr : public TypedCsrBase { } }; -/// Internal append-only access used while a fresh CSR is being bulk built. -/// The CSR is not published until both passes complete, so this bypasses the -/// incremental-growth path used by transactional updates. -template -class MutableCsrBulkBuildAccess { - public: - using csr_t = MutableCsr; - using nbr_t = typename csr_t::nbr_t; - static constexpr bool kStoresEdges = true; - static constexpr bool kNeedsDegreeCount = true; - static constexpr bool kChecksSingleUniqueness = false; - static constexpr bool kTracksInputEdgeCount = false; - - explicit MutableCsrBulkBuildAccess(csr_t& csr) : csr_(csr) {} - - void PrepareBuild(vid_t vertex_count) { - CHECK(csr_.adj_list_buffer_ != nullptr); - CHECK(csr_.degree_list_ != nullptr); - CHECK(csr_.cap_list_ != nullptr); - CHECK(csr_.nbr_list_ != nullptr); - csr_.adj_list_buffer_->Resize(static_cast(vertex_count) * - sizeof(nbr_t*)); - csr_.degree_list_->Resize(static_cast(vertex_count) * sizeof(int)); - csr_.cap_list_->Resize(static_cast(vertex_count) * sizeof(int)); - csr_.locks_ = std::make_unique(vertex_count); - refresh_metadata_ptrs(); - CHECK(adj_lists_ != nullptr || vertex_count == 0); - CHECK(degrees_ != nullptr || vertex_count == 0); - CHECK(capacities_ != nullptr || vertex_count == 0); - for (vid_t i = 0; i < vertex_count; ++i) { - adj_lists_[i] = nullptr; - degrees_[i].store(0, std::memory_order_relaxed); - capacities_[i] = 0; - } - csr_.edge_num_.store(0, std::memory_order_relaxed); - csr_.unsorted_since_ = 0; - } - - void CountSerial(vid_t src, int count = 1) { - CHECK_LT(src, vertex_capacity_); - CHECK_GT(count, 0); - auto degree = degrees_[src].load(std::memory_order_relaxed); - CHECK_LE(degree, std::numeric_limits::max() - count); - degrees_[src].store(degree + count, std::memory_order_relaxed); - } - - void CountConcurrent(vid_t src, int count = 1) { - CHECK_LT(src, vertex_capacity_); - CHECK_GT(count, 0); - const auto degree = - degrees_[src].fetch_add(count, std::memory_order_relaxed); - CHECK_LE(degree, std::numeric_limits::max() - count); - } - - void AllocateFromCounts() { - refresh_metadata_ptrs(); - size_t total_capacity = 0; - for (size_t i = 0; i < vertex_capacity_; ++i) { - const auto degree = degrees_[i].load(std::memory_order_relaxed); - const auto reserved_capacity = csr_t::reserved_capacity(degree); - // During the fill pass cap_list_ stores the exact expected degree. This - // lets range reservations validate both passes without allocating a - // second O(V) metadata array. Finish() converts it to runtime capacity. - capacities_[i] = degree; - CHECK_LE(static_cast(reserved_capacity), - std::numeric_limits::max() - total_capacity); - total_capacity += static_cast(reserved_capacity); - } - CHECK_LE(total_capacity, - std::numeric_limits::max() / sizeof(nbr_t)); - csr_.nbr_list_->Resize(total_capacity * sizeof(nbr_t)); - nbrs_ = reinterpret_cast(csr_.nbr_list_->GetData()); - CHECK(nbrs_ != nullptr || total_capacity == 0); - - size_t offset = 0; - for (size_t i = 0; i < vertex_capacity_; ++i) { - const auto reserved_capacity = csr_t::reserved_capacity(capacities_[i]); - adj_lists_[i] = reserved_capacity == 0 ? nullptr : nbrs_ + offset; - offset += static_cast(reserved_capacity); - degrees_[i].store(0, std::memory_order_relaxed); - } - } - - int ReserveSerial(vid_t src, int count) { - CHECK_LT(src, vertex_capacity_); - CHECK_GT(count, 0); - const auto slot = degrees_[src].load(std::memory_order_relaxed); - CHECK_LE(slot, capacities_[src] - count); - degrees_[src].store(slot + count, std::memory_order_relaxed); - return slot; - } - - int ReserveConcurrent(vid_t src, int count) { - CHECK_LT(src, vertex_capacity_); - CHECK_GT(count, 0); - CHECK_LE(count, capacities_[src]); - const auto slot = degrees_[src].fetch_add(count, std::memory_order_relaxed); - CHECK_LE(slot, capacities_[src] - count); - return slot; - } - - void PutAt(vid_t src, int slot, vid_t dst, const EDATA_T& data, - timestamp_t ts) { - CHECK_LT(src, vertex_capacity_); - CHECK_GE(slot, 0); - CHECK_LT(slot, capacities_[src]); - auto& nbr = adj_lists_[src][slot]; - nbr.neighbor = dst; - nbr.data = data; - nbr.timestamp.store(ts, std::memory_order_relaxed); - } - - void PutSerial(vid_t src, vid_t dst, const EDATA_T& data, timestamp_t ts) { - PutAt(src, ReserveSerial(src, 1), dst, data, ts); - } - - void Finish() { - uint64_t edge_num = 0; - for (size_t i = 0; i < vertex_capacity_; ++i) { - const auto degree = degrees_[i].load(std::memory_order_relaxed); - CHECK_EQ(degree, capacities_[i]) - << "Bulk edge count/fill mismatch for vertex " << i; - edge_num += static_cast(degree); - capacities_[i] = csr_t::reserved_capacity(degree); - } - csr_.edge_num_.store(edge_num, std::memory_order_relaxed); - csr_.refresh_prefetch_policy(); - } - - private: - void refresh_metadata_ptrs() { - vertex_capacity_ = csr_.vertex_capacity(); - adj_lists_ = reinterpret_cast( - csr_.adj_list_buffer_ == nullptr ? nullptr - : csr_.adj_list_buffer_->GetData()); - degrees_ = reinterpret_cast*>( - csr_.degree_list_ == nullptr ? nullptr : csr_.degree_list_->GetData()); - capacities_ = reinterpret_cast( - csr_.cap_list_ == nullptr ? nullptr : csr_.cap_list_->GetData()); - } - - csr_t& csr_; - size_t vertex_capacity_ = 0; - nbr_t** adj_lists_ = nullptr; - std::atomic* degrees_ = nullptr; - int* capacities_ = nullptr; - nbr_t* nbrs_ = nullptr; -}; - -/// Append-only writer for single-edge CSR layouts. A preceding degree pass can -/// use the timestamp field as a build-private seen marker. If every endpoint is -/// unique, workers can then fill fixed vertex slots concurrently without -/// locks. Duplicate endpoints require the ordered serial fill path to preserve -/// the pre-existing last-write-wins behavior. -template -class SingleMutableCsrBulkBuildAccess { - public: - using csr_t = SingleMutableCsr; - using nbr_t = typename csr_t::nbr_t; - static constexpr bool kStoresEdges = true; - static constexpr bool kNeedsDegreeCount = false; - static constexpr bool kChecksSingleUniqueness = true; - static constexpr bool kTracksInputEdgeCount = true; - - explicit SingleMutableCsrBulkBuildAccess(csr_t& csr) : csr_(csr) {} - - void PrepareBuild(vid_t vertex_count) { - CHECK(csr_.nbr_list_ != nullptr); - csr_.nbr_list_->Resize(static_cast(vertex_count) * sizeof(nbr_t)); - refresh_ptrs(); - CHECK(nbrs_ != nullptr || vertex_count == 0); - for (vid_t i = 0; i < vertex_count; ++i) { - nbrs_[i].timestamp.store(INVALID_TIMESTAMP, std::memory_order_relaxed); - } - csr_.edge_num_.store(0, std::memory_order_relaxed); - } - - void AllocateFromCounts() {} - - bool CheckUniqueSerial(vid_t src) { - CHECK_LT(src, vertex_capacity_); - auto& timestamp = nbrs_[src].timestamp; - const auto previous = timestamp.load(std::memory_order_relaxed); - timestamp.store(0, std::memory_order_relaxed); - return previous != INVALID_TIMESTAMP; - } - - bool CheckUniqueConcurrent(vid_t src) { - CHECK_LT(src, vertex_capacity_); - const auto previous = - nbrs_[src].timestamp.exchange(0, std::memory_order_relaxed); - return previous != INVALID_TIMESTAMP; - } - - void PutSerial(vid_t src, vid_t dst, const EDATA_T& data, timestamp_t ts) { - PutConcurrent(src, dst, data, ts); - } - - void PutConcurrent(vid_t src, vid_t dst, const EDATA_T& data, - timestamp_t ts) { - CHECK_LT(src, vertex_capacity_); - auto& nbr = nbrs_[src]; - nbr.neighbor = dst; - nbr.data = data; - nbr.timestamp.store(ts, std::memory_order_relaxed); - } - - void RecordFilledEdges(size_t count) { - csr_.edge_num_.fetch_add(static_cast(count), - std::memory_order_relaxed); - } - - void Finish() { csr_.refresh_prefetch_policy(); } - - private: - void refresh_ptrs() { - vertex_capacity_ = csr_.vertex_capacity(); - nbrs_ = reinterpret_cast( - csr_.nbr_list_ == nullptr ? nullptr : csr_.nbr_list_->GetData()); - } - - csr_t& csr_; - size_t vertex_capacity_ = 0; - nbr_t* nbrs_ = nullptr; -}; - } // namespace neug diff --git a/include/neug/storages/graph/edge_table.h b/include/neug/storages/graph/edge_table.h index 4758471d5..52d78f8e9 100644 --- a/include/neug/storages/graph/edge_table.h +++ b/include/neug/storages/graph/edge_table.h @@ -193,7 +193,8 @@ class EdgeTable { private: bool TryBatchBuildEdges(const IndexerType& src_indexer, const IndexerType& dst_indexer, - const std::shared_ptr& source); + const std::shared_ptr& source, + vid_t src_vertex_capacity, vid_t dst_vertex_capacity); void dropAndCreateNewBundledCSR(Checkpoint& ckp, ColumnBase* prev_data_col); void dropAndCreateNewUnbundledCSR(Checkpoint& ckp, bool delete_property); diff --git a/include/neug/storages/loader/loader_utils.h b/include/neug/storages/loader/loader_utils.h index d83acb83c..68f1a9d13 100644 --- a/include/neug/storages/loader/loader_utils.h +++ b/include/neug/storages/loader/loader_utils.h @@ -169,6 +169,8 @@ class CSVChunkSource final : public IDataChunkSource { std::vector file_paths_; CsvReadConfig config_; std::vector projected_columns_; + // Source-local cache: file paths and partition-relevant config stay fixed + // across Open() calls; producer count selects the cached plan. std::shared_ptr partition_plan_cache_; }; diff --git a/src/execution/execute/ops/batch/batch_insert_edge.cc b/src/execution/execute/ops/batch/batch_insert_edge.cc index 826dde914..5cfa347ed 100644 --- a/src/execution/execute/ops/batch/batch_insert_edge.cc +++ b/src/execution/execute/ops/batch/batch_insert_edge.cc @@ -69,22 +69,6 @@ bool resolve_edge_triplet(const Schema& schema, return true; } -std::vector> build_total_edge_mappings( - const std::vector>& source_mappings, - const std::vector>& destination_mappings, - const std::vector>& property_mappings) { - std::vector> mappings; - mappings.reserve(source_mappings.size() + destination_mappings.size() + - property_mappings.size()); - mappings.insert(mappings.end(), source_mappings.begin(), - source_mappings.end()); - mappings.insert(mappings.end(), destination_mappings.begin(), - destination_mappings.end()); - mappings.insert(mappings.end(), property_mappings.begin(), - property_mappings.end()); - return mappings; -} - } // namespace class BatchInsertEdgeOpr : public IOperator { @@ -110,9 +94,8 @@ class BatchInsertEdgeOpr : public IOperator { private: physical::EdgeType edge_type_; - std::vector> property_mappings_; - std::vector> source_mappings_; - std::vector> destination_mappings_; + std::vector> property_mappings_, + source_mappings_, destination_mappings_; std::optional source_; }; @@ -132,8 +115,15 @@ neug::result BatchInsertEdgeOpr::Eval( "BatchInsertEdge"); } - auto mappings = build_total_edge_mappings( - source_mappings_, destination_mappings_, property_mappings_); + std::vector> mappings; + mappings.reserve(source_mappings_.size() + destination_mappings_.size() + + property_mappings_.size()); + mappings.insert(mappings.end(), source_mappings_.begin(), + source_mappings_.end()); + mappings.insert(mappings.end(), destination_mappings_.begin(), + destination_mappings_.end()); + mappings.insert(mappings.end(), property_mappings_.begin(), + property_mappings_.end()); BatchInsertInput input; if (source_) { input = create_batch_insert_input(source_->state, *source_->read_function, diff --git a/src/storages/csr/mutable_csr.cc b/src/storages/csr/mutable_csr.cc index cb4d8609e..91702bd7c 100644 --- a/src/storages/csr/mutable_csr.cc +++ b/src/storages/csr/mutable_csr.cc @@ -33,6 +33,7 @@ #include #include +#include "neug/config.h" #include "neug/storages/container/container_utils.h" #include "neug/storages/container/file_mmap_container.h" #include "neug/utils/exception/exception.h" @@ -42,6 +43,16 @@ namespace neug { +int mutable_csr_detail::capacity_with_reserve(int degree) { + CHECK_GE(degree, 0); + if (degree == 0) { + return 0; + } + const auto capacity = std::ceil(degree * NeugDBConfig::DEFAULT_RESERVE_RATIO); + CHECK_LE(capacity, static_cast(std::numeric_limits::max())); + return static_cast(capacity); +} + template void MutableCsr::Open(Checkpoint& ckp, const ModuleDescriptor& descriptor, @@ -473,7 +484,7 @@ void MutableCsr::batch_put_edges(const std::vector& src_list, int old_deg = sz_arr[i].load(std::memory_order_relaxed); total_to_move += old_deg; int new_degree = degree[i] + old_deg; - int new_cap = reserved_capacity(new_degree); + int new_cap = mutable_csr_detail::capacity_with_reserve(new_degree); cap_arr[i] = new_cap; total_to_allocate += new_cap; } diff --git a/src/storages/graph/edge_table.cc b/src/storages/graph/edge_table.cc index b9ae0052d..c1e4a4acb 100644 --- a/src/storages/graph/edge_table.cc +++ b/src/storages/graph/edge_table.cc @@ -22,84 +22,38 @@ #include #include #include -#include #include -#include -#include #include -#include +#include #include -#include #include #include #include -#include #include #include #include "neug/common/columns/value_columns.h" -#include "neug/common/types/container_types.h" #include "neug/storages/checkpoint_manager.h" #include "neug/storages/csr/csr_view_utils.h" #include "neug/storages/csr/immutable_csr.h" #include "neug/storages/csr/mutable_csr.h" -#include "neug/storages/loader/chunk_pipeline_utils.h" #include "neug/storages/loader/loader_utils.h" #include "neug/storages/module/type_name.h" #include "neug/storages/module_descriptor.h" #include "neug/utils/io/file/file_utils.h" #include "neug/utils/property/types.h" +#include "../loader/bundled_edge_csr_loader.h" + namespace neug { namespace { -bool should_use_bulk_edge_build(const IDataChunkSource& source) { - const char* configured = std::getenv("NEUG_COPY_BULK_BUILD"); - if (configured != nullptr) { - std::string value(configured); - std::transform(value.begin(), value.end(), value.begin(), - [](unsigned char ch) { return std::tolower(ch); }); - if (value == "0" || value == "false" || value == "off" || value == "no") { - return false; - } - if (value == "1" || value == "true" || value == "on" || value == "yes") { - return true; - } - LOG(WARNING) << "Ignore invalid NEUG_COPY_BULK_BUILD=" << configured; - } - - constexpr int64_t kMinBulkBuildBytes = 256LL * 1024 * 1024; - return source.EstimatedBytes() >= kMinBulkBuildBytes; -} - -ChunkSourceOptions resolve_bulk_edge_source_options(int64_t source_bytes, - bool parallel_enabled, - bool preserve_order) { - constexpr int64_t kMinParallelBytes = 256LL * 1024 * 1024; - constexpr int64_t kMinPartitionBytes = 64LL * 1024 * 1024; - constexpr size_t kMaxQueuedChunks = 64; - - ChunkSourceOptions options; - options.preserve_order = preserve_order; - const auto workers = chunk_pipeline_detail::hardware_worker_count(); - if (!parallel_enabled || preserve_order || workers <= 1 || - source_bytes < kMinParallelBytes) { - return options; - } - - const auto useful_partitions = std::max( - 1, source_bytes / kMinPartitionBytes + - (source_bytes % kMinPartitionBytes == 0 ? 0 : 1)); - const auto balanced_producers = (workers + 1) / 2; - options.producer_count = static_cast(std::min( - balanced_producers, std::min(useful_partitions, workers - 1))); - options.producer_count = std::max(1, options.producer_count); - options.consumer_count = - std::max(1, workers - options.producer_count); - options.queue_capacity = std::clamp( - static_cast(options.producer_count) * 2, 2, kMaxQueuedChunks); - return options; +vid_t indexer_vertex_capacity(const IndexerType& indexer) { + const size_t capacity = indexer.capacity(); + CHECK_LE(capacity, static_cast(std::numeric_limits::max())) + << "CSR vertex capacity exceeds the vertex id range"; + return static_cast(capacity); } } // namespace @@ -448,580 +402,6 @@ void batch_add_bundled_edges_impl( } } -template -class EmptyCsrBulkWriter { - public: - static constexpr bool kStoresEdges = false; - static constexpr bool kNeedsDegreeCount = false; - static constexpr bool kChecksSingleUniqueness = false; - static constexpr bool kTracksInputEdgeCount = false; - - void PrepareBuild(vid_t /*vertex_count*/) {} - void AllocateFromCounts() {} - void Finish() {} -}; - -template -bool with_csr_bulk_writer(CsrBase* csr, F&& callback) { - if (auto* typed = dynamic_cast*>(csr)) { - MutableCsrBulkBuildAccess writer(*typed); - std::forward(callback)(writer); - return true; - } - if (auto* typed = dynamic_cast*>(csr)) { - SingleMutableCsrBulkBuildAccess writer(*typed); - std::forward(callback)(writer); - return true; - } - if (dynamic_cast*>(csr) != nullptr) { - EmptyCsrBulkWriter writer; - std::forward(callback)(writer); - return true; - } - return false; -} - -template -class BulkEdgeDataReader { - public: - explicit BulkEdgeDataReader(const std::shared_ptr& column) - : column_(column.get()) { - CHECK(column_ != nullptr); - auto values = std::dynamic_pointer_cast>(column); - if (values) { - values_ = &values->data(); - } - } - - EDATA_T Get(size_t row) const { - if (values_ != nullptr) { - return (*values_)[row]; - } - return column_->get_elem(row).template GetValue(); - } - - private: - const IContextColumn* column_; - const vector_t* values_ = nullptr; -}; - -template <> -class BulkEdgeDataReader { - public: - explicit BulkEdgeDataReader( - const std::shared_ptr& /*column*/) {} - - EmptyType Get(size_t /*row*/) const { return EmptyType(); } -}; - -constexpr uint32_t kInvalidBulkEdgeIndex = std::numeric_limits::max(); - -struct BulkEdgeGroup { - uint32_t count = 0; - uint32_t head = kInvalidBulkEdgeIndex; - uint32_t tail = kInvalidBulkEdgeIndex; -}; - -struct BulkEdgeWorkerScratch { - std::vector src_lids; - std::vector dst_lids; - flat_hash_map out_groups; - flat_hash_map in_groups; - std::vector out_next; - std::vector in_next; - bool out_single_duplicate = false; - bool in_single_duplicate = false; -}; - -void index_bulk_edge_endpoints(const std::shared_ptr& chunk, - const IndexerType& src_indexer, - const IndexerType& dst_indexer, - BulkEdgeWorkerScratch& scratch) { - CHECK(chunk != nullptr); - CHECK_GE(chunk->col_num(), 2); - auto src_column = chunk->get(0); - auto dst_column = chunk->get(1); - CHECK(src_column != nullptr); - CHECK(dst_column != nullptr); - CHECK_EQ(src_column->size(), dst_column->size()); - - scratch.src_lids.clear(); - scratch.dst_lids.clear(); - src_indexer.get_index(*src_column, scratch.src_lids); - dst_indexer.get_index(*dst_column, scratch.dst_lids); - CHECK_EQ(scratch.src_lids.size(), scratch.dst_lids.size()); -} - -inline bool is_valid_bulk_edge(vid_t src, vid_t dst) { - return src != std::numeric_limits::max() && - dst != std::numeric_limits::max(); -} - -struct BulkEdgeCountSummary { - bool out_single_duplicate = false; - bool in_single_duplicate = false; -}; - -constexpr size_t kBulkEdgeInitialGroupReserve = 4096; - -size_t bulk_edge_group_reserve(size_t rows) { - // Eagerly reserving one hash slot per edge wastes substantial memory for a - // supernode. The map can grow past this cap for high-cardinality chunks, and - // the worker-local scratch reuses that larger backing store on later chunks. - return std::min(rows, kBulkEdgeInitialGroupReserve); -} - -void count_bulk_edge_group(flat_hash_map& groups, - vid_t vertex) { - auto& group = groups[vertex]; - CHECK_LT(group.count, std::numeric_limits::max()); - ++group.count; -} - -template -void count_bulk_edge_chunk(BulkEdgeWorkerScratch& scratch, OutWriter& out, - InWriter& in, bool concurrent) { - CHECK_LE(scratch.src_lids.size(), - static_cast(std::numeric_limits::max())); - - if (concurrent) { - if constexpr (OutWriter::kNeedsDegreeCount) { - scratch.out_groups.clear(); - scratch.out_groups.reserve( - bulk_edge_group_reserve(scratch.src_lids.size())); - } - if constexpr (InWriter::kNeedsDegreeCount) { - scratch.in_groups.clear(); - scratch.in_groups.reserve( - bulk_edge_group_reserve(scratch.dst_lids.size())); - } - } - for (size_t row = 0; row < scratch.src_lids.size(); ++row) { - const auto src = scratch.src_lids[row]; - const auto dst = scratch.dst_lids[row]; - if (!is_valid_bulk_edge(src, dst)) { - continue; - } - if (concurrent) { - if constexpr (OutWriter::kNeedsDegreeCount) { - count_bulk_edge_group(scratch.out_groups, src); - } - if constexpr (InWriter::kNeedsDegreeCount) { - count_bulk_edge_group(scratch.in_groups, dst); - } - if constexpr (OutWriter::kChecksSingleUniqueness) { - if (!scratch.out_single_duplicate) { - scratch.out_single_duplicate = out.CheckUniqueConcurrent(src); - } - } - if constexpr (InWriter::kChecksSingleUniqueness) { - if (!scratch.in_single_duplicate) { - scratch.in_single_duplicate = in.CheckUniqueConcurrent(dst); - } - } - } else { - if constexpr (OutWriter::kNeedsDegreeCount) { - out.CountSerial(src); - } - if constexpr (InWriter::kNeedsDegreeCount) { - in.CountSerial(dst); - } - if constexpr (OutWriter::kChecksSingleUniqueness) { - if (!scratch.out_single_duplicate) { - scratch.out_single_duplicate = out.CheckUniqueSerial(src); - } - } - if constexpr (InWriter::kChecksSingleUniqueness) { - if (!scratch.in_single_duplicate) { - scratch.in_single_duplicate = in.CheckUniqueSerial(dst); - } - } - } - } - if (concurrent) { - if constexpr (OutWriter::kNeedsDegreeCount) { - for (const auto& [src, group] : scratch.out_groups) { - CHECK_LE(group.count, - static_cast(std::numeric_limits::max())); - out.CountConcurrent(src, static_cast(group.count)); - } - } - if constexpr (InWriter::kNeedsDegreeCount) { - for (const auto& [dst, group] : scratch.in_groups) { - CHECK_LE(group.count, - static_cast(std::numeric_limits::max())); - in.CountConcurrent(dst, static_cast(group.count)); - } - } - } -} - -using BulkEdgeCountChunk = std::function; - -BulkEdgeCountSummary count_bulk_edges( - const std::shared_ptr& source, - const IndexerType& src_indexer, const IndexerType& dst_indexer, - ChunkSourceOptions options, std::vector& scratches, - const BulkEdgeCountChunk& count_chunk) { - // Degree accumulation and single-slot uniqueness checks are commutative, so - // this pass never needs input order even when fill may later fall back to the - // ordered last-write-wins path. - // The degree pass only needs endpoint OIDs, so edge properties are not - // parsed, typed, allocated, and discarded during the first pass. - options.projected_columns = {0, 1}; - auto supplier = source->Open(options); - CHECK(supplier != nullptr); - - const auto worker_count = options.consumer_count; - CHECK_EQ(scratches.size(), static_cast(worker_count)); - const bool concurrent = worker_count > 1; - auto count = [&](int32_t worker, const std::shared_ptr& chunk) { - CHECK_GE(worker, 0); - CHECK_LT(worker, worker_count); - auto& scratch = scratches[static_cast(worker)]; - index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); - count_chunk(scratch, concurrent); - }; - consume_supplier_indexed(*supplier, options, count); - BulkEdgeCountSummary summary; - for (const auto& scratch : scratches) { - summary.out_single_duplicate = - summary.out_single_duplicate || scratch.out_single_duplicate; - summary.in_single_duplicate = - summary.in_single_duplicate || scratch.in_single_duplicate; - } - return summary; -} - -void append_bulk_edge_group(flat_hash_map& groups, - std::vector& next, vid_t vertex, - uint32_t edge_index) { - auto [it, inserted] = groups.try_emplace(vertex); - auto& group = it->second; - if (inserted) { - group.head = edge_index; - } else { - CHECK_NE(group.tail, kInvalidBulkEdgeIndex); - next[group.tail] = edge_index; - } - group.tail = edge_index; - CHECK_LT(group.count, std::numeric_limits::max()); - ++group.count; -} - -template -void group_bulk_edge_chunk(BulkEdgeWorkerScratch& scratch) { - CHECK_LE(scratch.src_lids.size(), - static_cast(std::numeric_limits::max())); - if constexpr (OutWriter::kNeedsDegreeCount) { - scratch.out_groups.clear(); - scratch.out_groups.reserve( - bulk_edge_group_reserve(scratch.src_lids.size())); - scratch.out_next.assign(scratch.src_lids.size(), kInvalidBulkEdgeIndex); - } - if constexpr (InWriter::kNeedsDegreeCount) { - scratch.in_groups.clear(); - scratch.in_groups.reserve(bulk_edge_group_reserve(scratch.dst_lids.size())); - scratch.in_next.assign(scratch.dst_lids.size(), kInvalidBulkEdgeIndex); - } - for (size_t row = 0; row < scratch.src_lids.size(); ++row) { - const auto src = scratch.src_lids[row]; - const auto dst = scratch.dst_lids[row]; - if (!is_valid_bulk_edge(src, dst)) { - continue; - } - const auto edge_index = static_cast(row); - if constexpr (OutWriter::kNeedsDegreeCount) { - append_bulk_edge_group(scratch.out_groups, scratch.out_next, src, - edge_index); - } - if constexpr (InWriter::kNeedsDegreeCount) { - append_bulk_edge_group(scratch.in_groups, scratch.in_next, dst, - edge_index); - } - } -} - -template -void fill_bulk_edge_chunk_serial(const BulkEdgeWorkerScratch& scratch, - const BulkEdgeDataReader& data_reader, - OutWriter& out, InWriter& in) { - size_t valid_edges = 0; - for (size_t row = 0; row < scratch.src_lids.size(); ++row) { - const auto src = scratch.src_lids[row]; - const auto dst = scratch.dst_lids[row]; - if (!is_valid_bulk_edge(src, dst)) { - continue; - } - ++valid_edges; - const auto data = data_reader.Get(row); - if constexpr (OutWriter::kStoresEdges) { - out.PutSerial(src, dst, data, 0); - } - if constexpr (InWriter::kStoresEdges) { - in.PutSerial(dst, src, data, 0); - } - } - if constexpr (OutWriter::kTracksInputEdgeCount) { - out.RecordFilledEdges(valid_edges); - } - if constexpr (InWriter::kTracksInputEdgeCount) { - in.RecordFilledEdges(valid_edges); - } -} - -template -void fill_bulk_edge_chunk_concurrent( - BulkEdgeWorkerScratch& scratch, - const BulkEdgeDataReader& data_reader, OutWriter& out, - InWriter& in) { - group_bulk_edge_chunk(scratch); - constexpr bool kDirectOut = - OutWriter::kStoresEdges && !OutWriter::kNeedsDegreeCount; - constexpr bool kDirectIn = - InWriter::kStoresEdges && !InWriter::kNeedsDegreeCount; - if constexpr (kDirectOut || kDirectIn) { - size_t valid_edges = 0; - for (size_t row = 0; row < scratch.src_lids.size(); ++row) { - const auto src = scratch.src_lids[row]; - const auto dst = scratch.dst_lids[row]; - if (!is_valid_bulk_edge(src, dst)) { - continue; - } - ++valid_edges; - const auto data = data_reader.Get(row); - if constexpr (kDirectOut) { - out.PutConcurrent(src, dst, data, 0); - } - if constexpr (kDirectIn) { - in.PutConcurrent(dst, src, data, 0); - } - } - if constexpr (OutWriter::kTracksInputEdgeCount) { - out.RecordFilledEdges(valid_edges); - } - if constexpr (InWriter::kTracksInputEdgeCount) { - in.RecordFilledEdges(valid_edges); - } - } - if constexpr (OutWriter::kNeedsDegreeCount) { - for (const auto& [src, group] : scratch.out_groups) { - CHECK_LE(group.count, - static_cast(std::numeric_limits::max())); - auto slot = out.ReserveConcurrent(src, static_cast(group.count)); - auto edge_index = group.head; - for (uint32_t i = 0; i < group.count; ++i) { - CHECK_NE(edge_index, kInvalidBulkEdgeIndex); - const auto data = data_reader.Get(edge_index); - out.PutAt(src, slot++, scratch.dst_lids[edge_index], data, 0); - edge_index = scratch.out_next[edge_index]; - } - CHECK_EQ(edge_index, kInvalidBulkEdgeIndex); - } - } - if constexpr (InWriter::kNeedsDegreeCount) { - for (const auto& [dst, group] : scratch.in_groups) { - CHECK_LE(group.count, - static_cast(std::numeric_limits::max())); - auto slot = in.ReserveConcurrent(dst, static_cast(group.count)); - auto edge_index = group.head; - for (uint32_t i = 0; i < group.count; ++i) { - CHECK_NE(edge_index, kInvalidBulkEdgeIndex); - const auto data = data_reader.Get(edge_index); - in.PutAt(dst, slot++, scratch.src_lids[edge_index], data, 0); - edge_index = scratch.in_next[edge_index]; - } - CHECK_EQ(edge_index, kInvalidBulkEdgeIndex); - } - } -} - -template -using BulkEdgeFillChunk = std::function&, - BulkEdgeWorkerScratch&, bool)>; - -template -struct BulkEdgeBuildOps { - bool needs_degree_count = false; - bool checks_out_single = false; - bool checks_in_single = false; - BulkEdgeCountChunk count_chunk; - std::function allocate_from_counts; - BulkEdgeFillChunk fill_chunk; -}; - -template -BulkEdgeBuildOps make_bulk_edge_build_ops(OutWriter& out, - InWriter& in) { - BulkEdgeBuildOps ops; - constexpr bool kStoresAnyDirection = - OutWriter::kStoresEdges || InWriter::kStoresEdges; - constexpr bool kNeedsAnyDegreeCount = - OutWriter::kNeedsDegreeCount || InWriter::kNeedsDegreeCount; - ops.needs_degree_count = kNeedsAnyDegreeCount; - ops.checks_out_single = OutWriter::kChecksSingleUniqueness; - ops.checks_in_single = InWriter::kChecksSingleUniqueness; - ops.allocate_from_counts = [&out, &in]() { - out.AllocateFromCounts(); - in.AllocateFromCounts(); - }; - if constexpr (kNeedsAnyDegreeCount) { - ops.count_chunk = [&out, &in](BulkEdgeWorkerScratch& scratch, - bool concurrent) { - count_bulk_edge_chunk(scratch, out, in, concurrent); - }; - } - if constexpr (kStoresAnyDirection) { - ops.fill_chunk = [&out, &in](const BulkEdgeDataReader& data_reader, - BulkEdgeWorkerScratch& scratch, - bool concurrent) { - if (concurrent) { - fill_bulk_edge_chunk_concurrent(scratch, data_reader, out, in); - } else { - fill_bulk_edge_chunk_serial(scratch, data_reader, out, in); - } - }; - } - return ops; -} - -template -void fill_bulk_edges(const std::shared_ptr& source, - const IndexerType& src_indexer, - const IndexerType& dst_indexer, - const ChunkSourceOptions& options, - bool allow_concurrent_fill, - std::vector& scratches, - const BulkEdgeBuildOps& ops) { - CHECK(!allow_concurrent_fill || !options.preserve_order); - auto supplier = source->Open(options); - CHECK(supplier != nullptr); - - const auto worker_count = options.consumer_count; - CHECK_GE(scratches.size(), static_cast(worker_count)); - if (allow_concurrent_fill && worker_count > 1) { - auto fill = [&](int32_t worker, const std::shared_ptr& chunk) { - CHECK_GE(worker, 0); - CHECK_LT(worker, worker_count); - auto& scratch = scratches[static_cast(worker)]; - index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); - const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; - BulkEdgeDataReader data_reader(data_column); - ops.fill_chunk(data_reader, scratch, true); - }; - - consume_supplier_indexed(*supplier, options, fill); - return; - } - - auto& scratch = scratches.front(); - while (auto chunk = supplier->GetNextChunk()) { - index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); - const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; - BulkEdgeDataReader data_reader(data_column); - ops.fill_chunk(data_reader, scratch, false); - } -} - -vid_t csr_vertex_capacity(const IndexerType& indexer) { - const size_t capacity = indexer.capacity(); - CHECK_LE(capacity, static_cast(std::numeric_limits::max())) - << "CSR vertex capacity exceeds the vertex id range"; - return static_cast(capacity); -} - -template -void build_bundled_edges_with_ops( - const BulkEdgeBuildOps& ops, const IndexerType& src_indexer, - const IndexerType& dst_indexer, - const std::shared_ptr& source) { - // VertexTable owns its reserve policy. Edge storage consumes the resulting - // indexer capacity instead of duplicating PropertyGraph::Dump's policy. - if (!ops.fill_chunk) { - ops.allocate_from_counts(); - return; - } - const auto estimated_bytes = source->EstimatedBytes(); - const auto count_options = resolve_bulk_edge_source_options( - estimated_bytes, source->ParallelEnabled(), false); - const bool needs_degree_count = ops.needs_degree_count; - std::vector scratches; - if (needs_degree_count) { - scratches.resize(static_cast(count_options.consumer_count)); - } - const bool has_single_direction = - ops.checks_out_single || ops.checks_in_single; - BulkEdgeCountSummary count_summary; - if (needs_degree_count) { - CHECK(static_cast(ops.count_chunk)); - count_summary = count_bulk_edges(source, src_indexer, dst_indexer, - count_options, scratches, ops.count_chunk); - } - const bool single_duplicate = - (ops.checks_out_single && count_summary.out_single_duplicate) || - (ops.checks_in_single && count_summary.in_single_duplicate); - const bool preserve_fill_order = - has_single_direction && (!needs_degree_count || single_duplicate); - const bool allow_concurrent_fill = !preserve_fill_order; - const auto fill_options = - preserve_fill_order - ? resolve_bulk_edge_source_options(estimated_bytes, - source->ParallelEnabled(), true) - : count_options; - if (scratches.size() < static_cast(fill_options.consumer_count)) { - scratches.resize(static_cast(fill_options.consumer_count)); - } - ops.allocate_from_counts(); - fill_bulk_edges(source, src_indexer, dst_indexer, fill_options, - allow_concurrent_fill, scratches, ops); -} - -template -bool build_bundled_edges_typed( - CsrBase* out_csr, CsrBase* in_csr, const IndexerType& src_indexer, - const IndexerType& dst_indexer, - const std::shared_ptr& source) { - bool built = false; - const bool out_supported = - with_csr_bulk_writer(out_csr, [&](auto& out) { - const bool in_supported = - with_csr_bulk_writer(in_csr, [&](auto& in) { - out.PrepareBuild(csr_vertex_capacity(src_indexer)); - in.PrepareBuild(csr_vertex_capacity(dst_indexer)); - auto ops = make_bulk_edge_build_ops(out, in); - build_bundled_edges_with_ops(ops, src_indexer, dst_indexer, - source); - out.Finish(); - in.Finish(); - }); - built = in_supported; - }); - return out_supported && built; -} - -bool build_bundled_edges(CsrBase* out_csr, CsrBase* in_csr, - const std::shared_ptr& schema, - const IndexerType& src_indexer, - const IndexerType& dst_indexer, - const std::shared_ptr& source) { - const auto property_type = schema->properties.empty() - ? DataTypeId::kEmpty - : schema->properties[0].id(); - switch (property_type) { -#define TYPE_DISPATCHER(enum_val, cpp_type) \ - case DataTypeId::enum_val: \ - return build_bundled_edges_typed(out_csr, in_csr, src_indexer, \ - dst_indexer, source); - FOR_EACH_DATA_TYPE_NO_STRING(TYPE_DISPATCHER) -#undef TYPE_DISPATCHER - case DataTypeId::kEmpty: - return build_bundled_edges_typed(out_csr, in_csr, src_indexer, - dst_indexer, source); - default: - return false; - } -} - void EdgeTable::Init(std::shared_ptr ckp, MemoryLevel level) { CHECK(meta_ != nullptr) << "EdgeTable::Init requires schema"; @@ -1405,16 +785,18 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, const IndexerType& dst_indexer, std::shared_ptr supplier) { CHECK(supplier != nullptr); + const auto src_vertex_capacity = indexer_vertex_capacity(src_indexer); + const auto dst_vertex_capacity = indexer_vertex_capacity(dst_indexer); auto source = supplier->RepeatableSource(); - if (source && should_use_bulk_edge_build(*source) && - TryBatchBuildEdges(src_indexer, dst_indexer, source)) { + if (source && TryBatchBuildEdges(src_indexer, dst_indexer, source, + src_vertex_capacity, dst_vertex_capacity)) { return; } // Keep fallback COPY paths aligned with the vertex table's actual capacity, // while leaving completely unloaded edge tables lazy until persistence. - in_csr_->resize(csr_vertex_capacity(dst_indexer)); - out_csr_->resize(csr_vertex_capacity(src_indexer)); + in_csr_->resize(dst_vertex_capacity); + out_csr_->resize(src_vertex_capacity); std::vector src_lid, dst_lid; // Pre-reserve capacity to reduce vector reallocation on large graphs. auto total_rows = supplier->RowNum(); @@ -1477,18 +859,24 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, bool EdgeTable::TryBatchBuildEdges( const IndexerType& src_indexer, const IndexerType& dst_indexer, - const std::shared_ptr& source) { + const std::shared_ptr& source, vid_t src_vertex_capacity, + vid_t dst_vertex_capacity) { if (!source || !meta_ || !meta_->is_bundled() || !out_csr_ || !in_csr_ || out_csr_->edge_num() != 0 || in_csr_->edge_num() != 0 || (meta_->oe_strategy != EdgeStrategy::kNone && !meta_->oe_mutable) || (meta_->ie_strategy != EdgeStrategy::kNone && !meta_->ie_mutable)) { return false; } + const auto source_bytes = source->EstimatedBytes(); + if (!internal::BundledEdgeCsrLoader::ShouldBuild(source_bytes)) { + return false; + } EdgeTable staged(meta_); staged.Init(ckp_, memory_level_); - if (!build_bundled_edges(staged.out_csr_.get(), staged.in_csr_.get(), meta_, - src_indexer, dst_indexer, source)) { + if (!internal::BundledEdgeCsrLoader::TryBuild( + *staged.out_csr_, *staged.in_csr_, *meta_, src_indexer, dst_indexer, + source, source_bytes, src_vertex_capacity, dst_vertex_capacity)) { return false; } Swap(staged); diff --git a/src/storages/loader/bundled_edge_csr_loader.cc b/src/storages/loader/bundled_edge_csr_loader.cc new file mode 100644 index 000000000..4e7a0856a --- /dev/null +++ b/src/storages/loader/bundled_edge_csr_loader.cc @@ -0,0 +1,909 @@ +/** Copyright 2020 Alibaba Group Holding Limited. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "bundled_edge_csr_loader.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "neug/common/columns/value_columns.h" +#include "neug/common/types/container_types.h" +#include "neug/storages/csr/mutable_csr.h" +#include "neug/storages/loader/chunk_pipeline_utils.h" +#include "neug/storages/loader/loader_utils.h" +#include "neug/utils/property/types.h" + +namespace neug { + +namespace internal { + +/// Internal append-only builder used while a fresh CSR is staged. The CSR is +/// not published until both passes complete, so this bypasses transactional +/// incremental growth. +template +class BundledEdgeCsrLoader::MutableWriter { + public: + using csr_t = MutableCsr; + using nbr_t = typename csr_t::nbr_t; + static constexpr EdgeStrategy kStrategy = EdgeStrategy::kMultiple; + + explicit MutableWriter(csr_t& csr) : csr_(csr) {} + + void PrepareBuild(vid_t vertex_count) { + CHECK(csr_.adj_list_buffer_ != nullptr); + CHECK(csr_.degree_list_ != nullptr); + CHECK(csr_.cap_list_ != nullptr); + CHECK(csr_.nbr_list_ != nullptr); + csr_.adj_list_buffer_->Resize(static_cast(vertex_count) * + sizeof(nbr_t*)); + csr_.degree_list_->Resize(static_cast(vertex_count) * sizeof(int)); + csr_.cap_list_->Resize(static_cast(vertex_count) * sizeof(int)); + csr_.locks_ = std::make_unique(vertex_count); + refresh_metadata_ptrs(); + CHECK(adj_lists_ != nullptr || vertex_count == 0); + CHECK(degrees_ != nullptr || vertex_count == 0); + CHECK(capacities_ != nullptr || vertex_count == 0); + for (vid_t i = 0; i < vertex_count; ++i) { + adj_lists_[i] = nullptr; + degrees_[i].store(0, std::memory_order_relaxed); + capacities_[i] = 0; + } + csr_.edge_num_.store(0, std::memory_order_relaxed); + csr_.unsorted_since_ = 0; + } + + void CountSerial(vid_t src, int count = 1) { + DCHECK_LT(src, vertex_capacity_); + CHECK_GT(count, 0); + auto degree = degrees_[src].load(std::memory_order_relaxed); + CHECK_LE(degree, std::numeric_limits::max() - count); + degrees_[src].store(degree + count, std::memory_order_relaxed); + } + + void CountConcurrent(vid_t src, int count = 1) { + DCHECK_LT(src, vertex_capacity_); + CHECK_GT(count, 0); + const auto degree = + degrees_[src].fetch_add(count, std::memory_order_relaxed); + CHECK_LE(degree, std::numeric_limits::max() - count); + } + + void AllocateFromCounts() { + refresh_metadata_ptrs(); + size_t total_capacity = 0; + for (size_t i = 0; i < vertex_capacity_; ++i) { + const auto degree = degrees_[i].load(std::memory_order_relaxed); + const auto storage_capacity = + mutable_csr_detail::capacity_with_reserve(degree); + // During fill cap_list_ stores the exact expected degree, allowing both + // passes to be validated without another O(V) metadata array. + capacities_[i] = degree; + CHECK_LE(static_cast(storage_capacity), + std::numeric_limits::max() - total_capacity); + total_capacity += static_cast(storage_capacity); + } + CHECK_LE(total_capacity, + std::numeric_limits::max() / sizeof(nbr_t)); + csr_.nbr_list_->Resize(total_capacity * sizeof(nbr_t)); + nbrs_ = reinterpret_cast(csr_.nbr_list_->GetData()); + CHECK(nbrs_ != nullptr || total_capacity == 0); + + size_t offset = 0; + for (size_t i = 0; i < vertex_capacity_; ++i) { + const auto storage_capacity = + mutable_csr_detail::capacity_with_reserve(capacities_[i]); + adj_lists_[i] = storage_capacity == 0 ? nullptr : nbrs_ + offset; + offset += static_cast(storage_capacity); + degrees_[i].store(0, std::memory_order_relaxed); + } + } + + int ReserveSerial(vid_t src, int count) { + CHECK_LT(src, vertex_capacity_); + CHECK_GT(count, 0); + const auto slot = degrees_[src].load(std::memory_order_relaxed); + CHECK_LE(slot, capacities_[src] - count); + degrees_[src].store(slot + count, std::memory_order_relaxed); + return slot; + } + + int ReserveConcurrent(vid_t src, int count) { + CHECK_LT(src, vertex_capacity_); + CHECK_GT(count, 0); + CHECK_LE(count, capacities_[src]); + const auto slot = degrees_[src].fetch_add(count, std::memory_order_relaxed); + CHECK_LE(slot, capacities_[src] - count); + return slot; + } + + void PutAt(vid_t src, int slot, vid_t dst, const EDATA_T& data, + timestamp_t ts) { + DCHECK_LT(src, vertex_capacity_); + DCHECK_GE(slot, 0); + DCHECK_LT(slot, capacities_[src]); + auto& nbr = adj_lists_[src][slot]; + nbr.neighbor = dst; + nbr.data = data; + nbr.timestamp.store(ts, std::memory_order_relaxed); + } + + void PutSerial(vid_t src, vid_t dst, const EDATA_T& data, timestamp_t ts) { + PutAt(src, ReserveSerial(src, 1), dst, data, ts); + } + + void Finish(uint64_t filled_edge_count) { + uint64_t edge_num = 0; + for (size_t i = 0; i < vertex_capacity_; ++i) { + const auto degree = degrees_[i].load(std::memory_order_relaxed); + CHECK_EQ(degree, capacities_[i]) + << "Bulk edge count/fill mismatch for vertex " << i; + edge_num += static_cast(degree); + capacities_[i] = mutable_csr_detail::capacity_with_reserve(degree); + } + CHECK_EQ(edge_num, filled_edge_count); + csr_.edge_num_.store(edge_num, std::memory_order_relaxed); + csr_.refresh_prefetch_policy(); + } + + private: + void refresh_metadata_ptrs() { + vertex_capacity_ = csr_.vertex_capacity(); + adj_lists_ = reinterpret_cast( + csr_.adj_list_buffer_ == nullptr ? nullptr + : csr_.adj_list_buffer_->GetData()); + degrees_ = reinterpret_cast*>( + csr_.degree_list_ == nullptr ? nullptr : csr_.degree_list_->GetData()); + capacities_ = reinterpret_cast( + csr_.cap_list_ == nullptr ? nullptr : csr_.cap_list_->GetData()); + } + + csr_t& csr_; + size_t vertex_capacity_ = 0; + nbr_t** adj_lists_ = nullptr; + std::atomic* degrees_ = nullptr; + int* capacities_ = nullptr; + nbr_t* nbrs_ = nullptr; +}; + +/// Builder for single-edge CSR layouts. A degree pass can reuse timestamp as a +/// private seen marker. Unique endpoints are filled concurrently; duplicates +/// fall back to ordered last-write-wins fill. +template +class BundledEdgeCsrLoader::SingleMutableWriter { + public: + using csr_t = SingleMutableCsr; + using nbr_t = typename csr_t::nbr_t; + static constexpr EdgeStrategy kStrategy = EdgeStrategy::kSingle; + + explicit SingleMutableWriter(csr_t& csr) : csr_(csr) {} + + void PrepareBuild(vid_t vertex_count) { + CHECK(csr_.nbr_list_ != nullptr); + csr_.nbr_list_->Resize(static_cast(vertex_count) * sizeof(nbr_t)); + refresh_ptrs(); + CHECK(nbrs_ != nullptr || vertex_count == 0); + for (vid_t i = 0; i < vertex_count; ++i) { + nbrs_[i].timestamp.store(INVALID_TIMESTAMP, std::memory_order_relaxed); + } + csr_.edge_num_.store(0, std::memory_order_relaxed); + } + + bool MarkSeenAndCheckDuplicateSerial(vid_t src) { + DCHECK_LT(src, vertex_capacity_); + auto& timestamp = nbrs_[src].timestamp; + const auto previous = timestamp.load(std::memory_order_relaxed); + timestamp.store(0, std::memory_order_relaxed); + return previous != INVALID_TIMESTAMP; + } + + bool MarkSeenAndCheckDuplicateConcurrent(vid_t src) { + DCHECK_LT(src, vertex_capacity_); + const auto previous = + nbrs_[src].timestamp.exchange(0, std::memory_order_relaxed); + return previous != INVALID_TIMESTAMP; + } + + void PutSerial(vid_t src, vid_t dst, const EDATA_T& data, timestamp_t ts) { + PutAtVertex(src, dst, data, ts); + } + + void PutUnique(vid_t src, vid_t dst, const EDATA_T& data, timestamp_t ts) { + PutAtVertex(src, dst, data, ts); + } + + void Finish(uint64_t filled_edge_count) { + csr_.edge_num_.store(filled_edge_count, std::memory_order_relaxed); + csr_.refresh_prefetch_policy(); + } + + private: + void PutAtVertex(vid_t src, vid_t dst, const EDATA_T& data, timestamp_t ts) { + DCHECK_LT(src, vertex_capacity_); + auto& nbr = nbrs_[src]; + nbr.neighbor = dst; + nbr.data = data; + nbr.timestamp.store(ts, std::memory_order_relaxed); + } + + void refresh_ptrs() { + vertex_capacity_ = csr_.vertex_capacity(); + nbrs_ = reinterpret_cast( + csr_.nbr_list_ == nullptr ? nullptr : csr_.nbr_list_->GetData()); + } + + csr_t& csr_; + size_t vertex_capacity_ = 0; + nbr_t* nbrs_ = nullptr; +}; + +} // namespace internal + +namespace { + +constexpr int64_t kMinBulkEdgeBuildBytes = 256LL * 1024 * 1024; + +ChunkSourceOptions resolve_bulk_edge_source_options(int64_t source_bytes, + bool parallel_enabled, + bool preserve_order) { + constexpr int64_t kMinPartitionBytes = 64LL * 1024 * 1024; + constexpr size_t kMaxQueuedChunks = 64; + + ChunkSourceOptions options; + options.preserve_order = preserve_order; + const auto workers = chunk_pipeline_detail::hardware_worker_count(); + if (!parallel_enabled || preserve_order || workers <= 1 || + source_bytes < kMinBulkEdgeBuildBytes) { + return options; + } + + const auto useful_partitions = std::max( + 1, source_bytes / kMinPartitionBytes + + (source_bytes % kMinPartitionBytes == 0 ? 0 : 1)); + const auto balanced_producers = (workers + 1) / 2; + options.producer_count = static_cast(std::min( + balanced_producers, std::min(useful_partitions, workers - 1))); + options.producer_count = std::max(1, options.producer_count); + options.consumer_count = + std::max(1, workers - options.producer_count); + options.queue_capacity = std::clamp( + static_cast(options.producer_count) * 2, 2, kMaxQueuedChunks); + return options; +} + +class EmptyCsrBulkWriter { + public: + static constexpr EdgeStrategy kStrategy = EdgeStrategy::kNone; + + void PrepareBuild(vid_t /*vertex_count*/) {} + void Finish(uint64_t /*filled_edge_count*/) {} +}; + +template +bool with_csr_bulk_writer(CsrBase* csr, F&& callback) { + if (auto* typed = dynamic_cast*>(csr)) { + internal::BundledEdgeCsrLoader::MutableWriter writer(*typed); + std::forward(callback)(writer); + return true; + } + if (auto* typed = dynamic_cast*>(csr)) { + internal::BundledEdgeCsrLoader::SingleMutableWriter writer(*typed); + std::forward(callback)(writer); + return true; + } + if (dynamic_cast*>(csr) != nullptr) { + EmptyCsrBulkWriter writer; + std::forward(callback)(writer); + return true; + } + return false; +} + +template +class BulkEdgeDataReader { + public: + explicit BulkEdgeDataReader(const std::shared_ptr& column) + : column_(column.get()) { + CHECK(column_ != nullptr); + auto values = std::dynamic_pointer_cast>(column); + if (values) { + values_ = &values->data(); + } + } + + EDATA_T Get(size_t row) const { + if (values_ != nullptr) { + return (*values_)[row]; + } + return column_->get_elem(row).template GetValue(); + } + + private: + const IContextColumn* column_; + const vector_t* values_ = nullptr; +}; + +template <> +class BulkEdgeDataReader { + public: + explicit BulkEdgeDataReader( + const std::shared_ptr& /*column*/) {} + + EmptyType Get(size_t /*row*/) const { return EmptyType(); } +}; + +constexpr uint32_t kInvalidBulkEdgeIndex = std::numeric_limits::max(); + +struct BulkEdgeGroup { + uint32_t count = 0; + uint32_t head = kInvalidBulkEdgeIndex; + uint32_t tail = kInvalidBulkEdgeIndex; +}; + +struct BulkEdgeWorkerScratch { + std::vector src_lids; + std::vector dst_lids; + flat_hash_map out_groups; + flat_hash_map in_groups; + std::vector out_next; + std::vector in_next; + bool out_single_duplicate = false; + bool in_single_duplicate = false; + uint64_t filled_edge_count = 0; +}; + +void index_bulk_edge_endpoints(const std::shared_ptr& chunk, + const IndexerType& src_indexer, + const IndexerType& dst_indexer, + BulkEdgeWorkerScratch& scratch) { + CHECK(chunk != nullptr); + CHECK_GE(chunk->col_num(), 2); + auto src_column = chunk->get(0); + auto dst_column = chunk->get(1); + CHECK(src_column != nullptr); + CHECK(dst_column != nullptr); + CHECK_EQ(src_column->size(), dst_column->size()); + + scratch.src_lids.clear(); + scratch.dst_lids.clear(); + src_indexer.get_index(*src_column, scratch.src_lids); + dst_indexer.get_index(*dst_column, scratch.dst_lids); + CHECK_EQ(scratch.src_lids.size(), scratch.dst_lids.size()); +} + +inline bool is_valid_bulk_edge(vid_t src, vid_t dst) { + return src != std::numeric_limits::max() && + dst != std::numeric_limits::max(); +} + +struct BulkEdgeCountSummary { + bool out_single_duplicate = false; + bool in_single_duplicate = false; +}; + +constexpr size_t kBulkEdgeInitialGroupReserve = 4096; + +size_t bulk_edge_group_reserve(size_t rows) { + // Eagerly reserving one hash slot per edge wastes substantial memory for a + // supernode. The map can grow past this cap for high-cardinality chunks, and + // the worker-local scratch reuses that larger backing store on later chunks. + return std::min(rows, kBulkEdgeInitialGroupReserve); +} + +void count_bulk_edge_group(flat_hash_map& groups, + vid_t vertex) { + auto& group = groups[vertex]; + CHECK_LT(group.count, std::numeric_limits::max()); + ++group.count; +} + +template +void count_bulk_edge_chunk(BulkEdgeWorkerScratch& scratch, OutWriter& out, + InWriter& in, bool concurrent) { + CHECK_LE(scratch.src_lids.size(), + static_cast(std::numeric_limits::max())); + + if (concurrent) { + if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { + scratch.out_groups.clear(); + scratch.out_groups.reserve( + bulk_edge_group_reserve(scratch.src_lids.size())); + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { + scratch.in_groups.clear(); + scratch.in_groups.reserve( + bulk_edge_group_reserve(scratch.dst_lids.size())); + } + } + for (size_t row = 0; row < scratch.src_lids.size(); ++row) { + const auto src = scratch.src_lids[row]; + const auto dst = scratch.dst_lids[row]; + if (!is_valid_bulk_edge(src, dst)) { + continue; + } + if (concurrent) { + if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { + count_bulk_edge_group(scratch.out_groups, src); + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { + count_bulk_edge_group(scratch.in_groups, dst); + } + if constexpr (OutWriter::kStrategy == EdgeStrategy::kSingle) { + if (!scratch.out_single_duplicate) { + scratch.out_single_duplicate = + out.MarkSeenAndCheckDuplicateConcurrent(src); + } + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kSingle) { + if (!scratch.in_single_duplicate) { + scratch.in_single_duplicate = + in.MarkSeenAndCheckDuplicateConcurrent(dst); + } + } + } else { + if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { + out.CountSerial(src); + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { + in.CountSerial(dst); + } + if constexpr (OutWriter::kStrategy == EdgeStrategy::kSingle) { + if (!scratch.out_single_duplicate) { + scratch.out_single_duplicate = + out.MarkSeenAndCheckDuplicateSerial(src); + } + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kSingle) { + if (!scratch.in_single_duplicate) { + scratch.in_single_duplicate = in.MarkSeenAndCheckDuplicateSerial(dst); + } + } + } + } + if (concurrent) { + if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { + for (const auto& [src, group] : scratch.out_groups) { + CHECK_LE(group.count, + static_cast(std::numeric_limits::max())); + out.CountConcurrent(src, static_cast(group.count)); + } + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { + for (const auto& [dst, group] : scratch.in_groups) { + CHECK_LE(group.count, + static_cast(std::numeric_limits::max())); + in.CountConcurrent(dst, static_cast(group.count)); + } + } + } +} + +using BulkEdgeCountChunk = std::function; + +BulkEdgeCountSummary count_bulk_edges( + const std::shared_ptr& source, + const IndexerType& src_indexer, const IndexerType& dst_indexer, + ChunkSourceOptions options, std::vector& scratches, + const BulkEdgeCountChunk& count_chunk) { + // Degree accumulation and single-slot uniqueness checks are commutative, so + // this pass never needs input order even when fill may later fall back to the + // ordered last-write-wins path. + // The degree pass only needs endpoint OIDs, so edge properties are not + // parsed, typed, allocated, and discarded during the first pass. + options.projected_columns = {0, 1}; + auto supplier = source->Open(options); + CHECK(supplier != nullptr); + + const auto worker_count = options.consumer_count; + CHECK_EQ(scratches.size(), static_cast(worker_count)); + const bool concurrent = worker_count > 1; + auto count = [&](int32_t worker, const std::shared_ptr& chunk) { + CHECK_GE(worker, 0); + CHECK_LT(worker, worker_count); + auto& scratch = scratches[static_cast(worker)]; + index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); + count_chunk(scratch, concurrent); + }; + consume_supplier_indexed(*supplier, options, count); + BulkEdgeCountSummary summary; + for (const auto& scratch : scratches) { + summary.out_single_duplicate = + summary.out_single_duplicate || scratch.out_single_duplicate; + summary.in_single_duplicate = + summary.in_single_duplicate || scratch.in_single_duplicate; + } + return summary; +} + +void append_bulk_edge_group(flat_hash_map& groups, + std::vector& next, vid_t vertex, + uint32_t edge_index) { + auto [it, inserted] = groups.try_emplace(vertex); + auto& group = it->second; + if (inserted) { + group.head = edge_index; + } else { + CHECK_NE(group.tail, kInvalidBulkEdgeIndex); + next[group.tail] = edge_index; + } + group.tail = edge_index; + CHECK_LT(group.count, std::numeric_limits::max()); + ++group.count; +} + +template +size_t group_bulk_edge_chunk(BulkEdgeWorkerScratch& scratch) { + CHECK_LE(scratch.src_lids.size(), + static_cast(std::numeric_limits::max())); + if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { + scratch.out_groups.clear(); + scratch.out_groups.reserve( + bulk_edge_group_reserve(scratch.src_lids.size())); + scratch.out_next.assign(scratch.src_lids.size(), kInvalidBulkEdgeIndex); + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { + scratch.in_groups.clear(); + scratch.in_groups.reserve(bulk_edge_group_reserve(scratch.dst_lids.size())); + scratch.in_next.assign(scratch.dst_lids.size(), kInvalidBulkEdgeIndex); + } + size_t valid_edges = 0; + for (size_t row = 0; row < scratch.src_lids.size(); ++row) { + const auto src = scratch.src_lids[row]; + const auto dst = scratch.dst_lids[row]; + if (!is_valid_bulk_edge(src, dst)) { + continue; + } + ++valid_edges; + const auto edge_index = static_cast(row); + if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { + append_bulk_edge_group(scratch.out_groups, scratch.out_next, src, + edge_index); + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { + append_bulk_edge_group(scratch.in_groups, scratch.in_next, dst, + edge_index); + } + } + return valid_edges; +} + +template +void fill_bulk_edge_chunk_serial(BulkEdgeWorkerScratch& scratch, + const BulkEdgeDataReader& data_reader, + OutWriter& out, InWriter& in) { + size_t valid_edges = 0; + for (size_t row = 0; row < scratch.src_lids.size(); ++row) { + const auto src = scratch.src_lids[row]; + const auto dst = scratch.dst_lids[row]; + if (!is_valid_bulk_edge(src, dst)) { + continue; + } + ++valid_edges; + const auto data = data_reader.Get(row); + if constexpr (OutWriter::kStrategy != EdgeStrategy::kNone) { + out.PutSerial(src, dst, data, 0); + } + if constexpr (InWriter::kStrategy != EdgeStrategy::kNone) { + in.PutSerial(dst, src, data, 0); + } + } + CHECK_LE(valid_edges, + std::numeric_limits::max() - scratch.filled_edge_count); + scratch.filled_edge_count += static_cast(valid_edges); +} + +template +void fill_bulk_edge_chunk_concurrent( + BulkEdgeWorkerScratch& scratch, + const BulkEdgeDataReader& data_reader, OutWriter& out, + InWriter& in) { + const auto valid_edges = group_bulk_edge_chunk(scratch); + constexpr bool kDirectOut = OutWriter::kStrategy == EdgeStrategy::kSingle; + constexpr bool kDirectIn = InWriter::kStrategy == EdgeStrategy::kSingle; + if constexpr (kDirectOut || kDirectIn) { + for (size_t row = 0; row < scratch.src_lids.size(); ++row) { + const auto src = scratch.src_lids[row]; + const auto dst = scratch.dst_lids[row]; + if (!is_valid_bulk_edge(src, dst)) { + continue; + } + const auto data = data_reader.Get(row); + if constexpr (kDirectOut) { + out.PutUnique(src, dst, data, 0); + } + if constexpr (kDirectIn) { + in.PutUnique(dst, src, data, 0); + } + } + } + if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { + for (const auto& [src, group] : scratch.out_groups) { + CHECK_LE(group.count, + static_cast(std::numeric_limits::max())); + auto slot = out.ReserveConcurrent(src, static_cast(group.count)); + auto edge_index = group.head; + for (uint32_t i = 0; i < group.count; ++i) { + CHECK_NE(edge_index, kInvalidBulkEdgeIndex); + const auto data = data_reader.Get(edge_index); + out.PutAt(src, slot++, scratch.dst_lids[edge_index], data, 0); + edge_index = scratch.out_next[edge_index]; + } + CHECK_EQ(edge_index, kInvalidBulkEdgeIndex); + } + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { + for (const auto& [dst, group] : scratch.in_groups) { + CHECK_LE(group.count, + static_cast(std::numeric_limits::max())); + auto slot = in.ReserveConcurrent(dst, static_cast(group.count)); + auto edge_index = group.head; + for (uint32_t i = 0; i < group.count; ++i) { + CHECK_NE(edge_index, kInvalidBulkEdgeIndex); + const auto data = data_reader.Get(edge_index); + in.PutAt(dst, slot++, scratch.src_lids[edge_index], data, 0); + edge_index = scratch.in_next[edge_index]; + } + CHECK_EQ(edge_index, kInvalidBulkEdgeIndex); + } + } + CHECK_LE(valid_edges, + std::numeric_limits::max() - scratch.filled_edge_count); + scratch.filled_edge_count += static_cast(valid_edges); +} + +template +using BulkEdgeFillChunk = std::function&, + BulkEdgeWorkerScratch&, bool)>; + +template +struct BulkEdgeBuildOps { + bool needs_degree_count = false; + bool checks_out_single = false; + bool checks_in_single = false; + BulkEdgeCountChunk count_chunk; + std::function allocate_from_counts; + BulkEdgeFillChunk fill_chunk; +}; + +template +BulkEdgeBuildOps make_bulk_edge_build_ops(OutWriter& out, + InWriter& in) { + BulkEdgeBuildOps ops; + constexpr bool kStoresAnyDirection = + OutWriter::kStrategy != EdgeStrategy::kNone || + InWriter::kStrategy != EdgeStrategy::kNone; + constexpr bool kNeedsAnyDegreeCount = + OutWriter::kStrategy == EdgeStrategy::kMultiple || + InWriter::kStrategy == EdgeStrategy::kMultiple; + ops.needs_degree_count = kNeedsAnyDegreeCount; + ops.checks_out_single = OutWriter::kStrategy == EdgeStrategy::kSingle; + ops.checks_in_single = InWriter::kStrategy == EdgeStrategy::kSingle; + ops.allocate_from_counts = [&]() { + if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { + out.AllocateFromCounts(); + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { + in.AllocateFromCounts(); + } + }; + if constexpr (kNeedsAnyDegreeCount) { + ops.count_chunk = [&out, &in](BulkEdgeWorkerScratch& scratch, + bool concurrent) { + count_bulk_edge_chunk(scratch, out, in, concurrent); + }; + } + if constexpr (kStoresAnyDirection) { + ops.fill_chunk = [&out, &in](const BulkEdgeDataReader& data_reader, + BulkEdgeWorkerScratch& scratch, + bool concurrent) { + if (concurrent) { + fill_bulk_edge_chunk_concurrent(scratch, data_reader, out, in); + } else { + fill_bulk_edge_chunk_serial(scratch, data_reader, out, in); + } + }; + } + return ops; +} + +template +uint64_t fill_bulk_edges(const std::shared_ptr& source, + const IndexerType& src_indexer, + const IndexerType& dst_indexer, + const ChunkSourceOptions& options, + bool allow_concurrent_fill, + std::vector& scratches, + const BulkEdgeBuildOps& ops) { + CHECK(!allow_concurrent_fill || !options.preserve_order); + auto supplier = source->Open(options); + CHECK(supplier != nullptr); + + const auto worker_count = options.consumer_count; + CHECK_GE(scratches.size(), static_cast(worker_count)); + for (auto& scratch : scratches) { + scratch.filled_edge_count = 0; + } + if (allow_concurrent_fill && worker_count > 1) { + auto fill = [&](int32_t worker, const std::shared_ptr& chunk) { + CHECK_GE(worker, 0); + CHECK_LT(worker, worker_count); + auto& scratch = scratches[static_cast(worker)]; + index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); + const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; + BulkEdgeDataReader data_reader(data_column); + ops.fill_chunk(data_reader, scratch, true); + }; + + consume_supplier_indexed(*supplier, options, fill); + } else { + auto& scratch = scratches.front(); + while (auto chunk = supplier->GetNextChunk()) { + index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); + const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; + BulkEdgeDataReader data_reader(data_column); + ops.fill_chunk(data_reader, scratch, false); + } + } + + uint64_t filled_edge_count = 0; + for (const auto& scratch : scratches) { + CHECK_LE(scratch.filled_edge_count, + std::numeric_limits::max() - filled_edge_count); + filled_edge_count += scratch.filled_edge_count; + } + return filled_edge_count; +} + +template +uint64_t build_bundled_edges_with_ops( + const BulkEdgeBuildOps& ops, const IndexerType& src_indexer, + const IndexerType& dst_indexer, + const std::shared_ptr& source, int64_t source_bytes) { + // VertexTable owns its reserve policy. Edge storage consumes the resulting + // indexer capacity instead of duplicating PropertyGraph::Dump's policy. + if (!ops.fill_chunk) { + ops.allocate_from_counts(); + return 0; + } + const auto count_options = resolve_bulk_edge_source_options( + source_bytes, source->ParallelEnabled(), false); + const bool needs_degree_count = ops.needs_degree_count; + std::vector scratches; + if (needs_degree_count) { + scratches.resize(static_cast(count_options.consumer_count)); + } + const bool has_single_direction = + ops.checks_out_single || ops.checks_in_single; + BulkEdgeCountSummary count_summary; + if (needs_degree_count) { + CHECK(static_cast(ops.count_chunk)); + count_summary = count_bulk_edges(source, src_indexer, dst_indexer, + count_options, scratches, ops.count_chunk); + } + const bool single_duplicate = + (ops.checks_out_single && count_summary.out_single_duplicate) || + (ops.checks_in_single && count_summary.in_single_duplicate); + const bool preserve_fill_order = + has_single_direction && (!needs_degree_count || single_duplicate); + const bool allow_concurrent_fill = !preserve_fill_order; + const auto fill_options = + preserve_fill_order ? resolve_bulk_edge_source_options( + source_bytes, source->ParallelEnabled(), true) + : count_options; + if (scratches.size() < static_cast(fill_options.consumer_count)) { + scratches.resize(static_cast(fill_options.consumer_count)); + } + ops.allocate_from_counts(); + return fill_bulk_edges(source, src_indexer, dst_indexer, + fill_options, allow_concurrent_fill, + scratches, ops); +} + +template +bool build_bundled_edges_typed(CsrBase* out_csr, CsrBase* in_csr, + const IndexerType& src_indexer, + const IndexerType& dst_indexer, + const std::shared_ptr& source, + int64_t source_bytes, vid_t src_vertex_capacity, + vid_t dst_vertex_capacity) { + bool built = false; + const bool out_supported = + with_csr_bulk_writer(out_csr, [&](auto& out) { + const bool in_supported = + with_csr_bulk_writer(in_csr, [&](auto& in) { + out.PrepareBuild(src_vertex_capacity); + in.PrepareBuild(dst_vertex_capacity); + auto ops = make_bulk_edge_build_ops(out, in); + const auto filled_edge_count = build_bundled_edges_with_ops( + ops, src_indexer, dst_indexer, source, source_bytes); + out.Finish(filled_edge_count); + in.Finish(filled_edge_count); + }); + built = in_supported; + }); + return out_supported && built; +} + +bool build_bundled_edges(CsrBase* out_csr, CsrBase* in_csr, + const EdgeSchema& schema, + const IndexerType& src_indexer, + const IndexerType& dst_indexer, + const std::shared_ptr& source, + int64_t source_bytes, vid_t src_vertex_capacity, + vid_t dst_vertex_capacity) { + const auto property_type = schema.properties.empty() + ? DataTypeId::kEmpty + : schema.properties[0].id(); + switch (property_type) { +#define TYPE_DISPATCHER(enum_val, cpp_type) \ + case DataTypeId::enum_val: \ + return build_bundled_edges_typed( \ + out_csr, in_csr, src_indexer, dst_indexer, source, source_bytes, \ + src_vertex_capacity, dst_vertex_capacity); + FOR_EACH_DATA_TYPE_NO_STRING(TYPE_DISPATCHER) +#undef TYPE_DISPATCHER + case DataTypeId::kEmpty: + return build_bundled_edges_typed( + out_csr, in_csr, src_indexer, dst_indexer, source, source_bytes, + src_vertex_capacity, dst_vertex_capacity); + default: + return false; + } +} + +} // namespace + +bool internal::BundledEdgeCsrLoader::ShouldBuild(int64_t source_bytes) { + const char* configured = std::getenv("NEUG_COPY_BULK_BUILD"); + if (configured != nullptr) { + std::string value(configured); + std::transform(value.begin(), value.end(), value.begin(), + [](unsigned char ch) { return std::tolower(ch); }); + if (value == "0" || value == "false" || value == "off" || value == "no") { + return false; + } + if (value == "1" || value == "true" || value == "on" || value == "yes") { + return true; + } + LOG(WARNING) << "Ignore invalid NEUG_COPY_BULK_BUILD=" << configured; + } + + return source_bytes >= kMinBulkEdgeBuildBytes; +} + +bool internal::BundledEdgeCsrLoader::TryBuild( + CsrBase& out_csr, CsrBase& in_csr, const EdgeSchema& schema, + const IndexerType& src_indexer, const IndexerType& dst_indexer, + const std::shared_ptr& source, int64_t source_bytes, + vid_t src_vertex_capacity, vid_t dst_vertex_capacity) { + if (!source) { + return false; + } + return build_bundled_edges(&out_csr, &in_csr, schema, src_indexer, + dst_indexer, source, source_bytes, + src_vertex_capacity, dst_vertex_capacity); +} + +} // namespace neug diff --git a/src/storages/loader/bundled_edge_csr_loader.h b/src/storages/loader/bundled_edge_csr_loader.h new file mode 100644 index 000000000..ac382d413 --- /dev/null +++ b/src/storages/loader/bundled_edge_csr_loader.h @@ -0,0 +1,55 @@ +/** Copyright 2020 Alibaba Group Holding Limited. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include "neug/storages/graph/schema.h" +#include "neug/utils/indexers.h" + +namespace neug { + +class CsrBase; +class IDataChunkSource; + +namespace internal { + +/// Builds the outgoing and incoming CSR pair for a bundled edge table directly +/// from a repeatable chunk source. EdgeTable owns staging and publication; this +/// class owns source planning and the CSR bulk-build protocol. +class BundledEdgeCsrLoader { + public: + template + class MutableWriter; + + template + class SingleMutableWriter; + + static bool ShouldBuild(int64_t source_bytes); + + /// Returns false when either CSR layout or the property type is unsupported. + /// Callers must pass fresh, unpublished CSR instances. + static bool TryBuild(CsrBase& out_csr, CsrBase& in_csr, + const EdgeSchema& schema, const IndexerType& src_indexer, + const IndexerType& dst_indexer, + const std::shared_ptr& source, + int64_t source_bytes, vid_t src_vertex_capacity, + vid_t dst_vertex_capacity); +}; + +} // namespace internal +} // namespace neug diff --git a/src/storages/loader/loader_utils.cc b/src/storages/loader/loader_utils.cc index 9f6f24479..decb17050 100644 --- a/src/storages/loader/loader_utils.cc +++ b/src/storages/loader/loader_utils.cc @@ -86,8 +86,8 @@ struct CsvPartitionPlanCache { std::shared_ptr plan; }; - std::mutex mutex; - std::unordered_map> entries; + std::mutex mutex_; + std::unordered_map> entries_; }; namespace { @@ -454,40 +454,31 @@ struct RowCounterState { } }; -/// Fast CSV row counter: scans raw bytes to count row boundaries -/// without parsing fields. Uses a quote-tracking state machine; -/// parallelizes via speculative dual-state-machine scan for large files. +/// Result of scanning a CSV file without materializing its fields. struct CsvScanResult { int64_t row_count = 0; std::vector ranges; }; -class CsvRowCountCounter { +/// Scans CSV bytes to count rows and optionally produce record-aligned ranges. +/// csv-parser always recognizes RFC-style doubled quotes, so the scanner uses +/// the same effective dialect instead of the legacy double_quote option. +class CsvFileScanner { public: - // rows_to_skip is intentionally NOT a parameter: the counter counts + // rows_to_skip is intentionally NOT a parameter: the scanner counts // all non-empty rows. The reader's skip_rows() handles skipping // separately, and RowNum() is only a pre-allocation hint, so a slight // overcount (by at most skip_rows, typically 1 for header) is safe. - CsvRowCountCounter(std::string file_path, bool quoting, char quote_char, - bool /*double_quote*/, char delimiter, - bool use_threads = true) + CsvFileScanner(std::string file_path, bool quoting, char quote_char, + char delimiter, bool use_threads = true) : file_path_(std::move(file_path)), quoting_(quoting), quote_char_(quote_char), - // csv::CSVReader always uses RFC-style doubled quotes and exposes no - // switch for disabling that behavior. Partition planning must follow - // the parser's actual DFA even when the legacy DOUBLE_QUOTE option is - // false, otherwise a range can start in the middle of a record. - double_quote_(true), delimiter_(delimiter), use_threads_(use_threads) {} - int64_t count() const { - struct stat st; - if (stat(file_path_.c_str(), &st) != 0) { - THROW_IO_EXCEPTION("Failed to get file size: " + file_path_); - } - auto file_size = static_cast(st.st_size); + int64_t count_rows() const { + const auto file_size = get_file_size(); if (file_size == 0) return 0; @@ -502,25 +493,20 @@ class CsvRowCountCounter { } if (num_threads <= 1) return count_single(file_size); - return count_parallel(file_size, num_threads); + return scan_parallel(file_size, num_threads, false).row_count; } - CsvScanResult ScanWithRanges(int32_t requested_partitions, - int32_t scan_threads) const { - struct stat st; - if (stat(file_path_.c_str(), &st) != 0) { - THROW_IO_EXCEPTION("Failed to get file size: " + file_path_); - } - const auto file_size = static_cast(st.st_size); + CsvScanResult scan_with_ranges(size_t file_size, int32_t requested_partitions, + int32_t scan_threads) const { CsvScanResult result; if (file_size == 0) { return result; } - const auto target_partitions = static_cast(std::max( - 1, std::min(requested_partitions, - static_cast(file_size)))); - const auto workers = static_cast(std::max( - 1, std::min(scan_threads, static_cast(file_size)))); + const auto target_partitions = std::min( + file_size, + static_cast(std::max(1, requested_partitions))); + const auto workers = static_cast(std::min( + file_size, static_cast(std::max(1, scan_threads)))); if (workers == 1) { result.row_count = count_single(file_size); result.ranges.push_back({0, file_size, 0}); @@ -554,6 +540,18 @@ class CsvRowCountCounter { RowCounterState inside; // assumed start inside quotes }; + size_t get_file_size() const { + struct stat st; + if (stat(file_path_.c_str(), &st) != 0 || st.st_size < 0) { + THROW_IO_EXCEPTION("Failed to get file size: " + file_path_); + } + const auto file_size = static_cast(st.st_size); + if (file_size > std::numeric_limits::max()) { + THROW_IO_EXCEPTION("CSV file is too large to address: " + file_path_); + } + return static_cast(file_size); + } + /// Read [start, end) from the file in 1 MB buffers, calling \p fn /// for each buffer. Shared by count_single() and scan_chunk(). template @@ -570,6 +568,9 @@ class CsvRowCountCounter { size_t to_read = std::min(kBufSize, end - pos); file.read(buffer.data(), to_read); auto bytes_read = static_cast(file.gcount()); + if (file.bad()) { + THROW_IO_EXCEPTION("Failed to scan CSV file: " + file_path_); + } if (bytes_read == 0) break; fn(buffer.data(), bytes_read); @@ -597,6 +598,9 @@ class CsvRowCountCounter { size_t to_read = std::min(kScanBuf, end - pos); file.read(buf, to_read); auto n = static_cast(file.gcount()); + if (file.bad()) { + THROW_IO_EXCEPTION("Failed to scan CSV file: " + file_path_); + } if (n == 0) break; for (size_t i = 0; i < n; ++i) { @@ -627,8 +631,8 @@ class CsvRowCountCounter { /// other is copied — both produce identical results. ChunkResult scan_chunk(size_t start, size_t end) const { ChunkResult res; - res.outside.init(false, quoting_, quote_char_, double_quote_, delimiter_); - res.inside.init(true, quoting_, quote_char_, double_quote_, delimiter_); + res.outside.init(false, quoting_, quote_char_, true, delimiter_); + res.inside.init(true, quoting_, quote_char_, true, delimiter_); if (quoting_) { scan_range(start, end, [&](const char* data, size_t n) { for (size_t i = 0; i < n; ++i) { @@ -653,7 +657,7 @@ class CsvRowCountCounter { /// Single-threaded scan. int64_t count_single(size_t file_size) const { RowCounterState state; - state.init(false, quoting_, quote_char_, double_quote_, delimiter_); + state.init(false, quoting_, quote_char_, true, delimiter_); scan_range(0, file_size, [&](const char* data, size_t n) { for (size_t i = 0; i < n; ++i) { state.step(data[i]); @@ -666,44 +670,6 @@ class CsvRowCountCounter { } /// Parallel scan with speculative dual state machines. - int64_t count_parallel(size_t file_size, unsigned num_threads) const { - return scan_parallel(file_size, num_threads, false).row_count; - } - - CsvPartitionRange find_record_boundary(size_t start, size_t end, - int64_t start_row) const { - RowCounterState state; - state.init(true, quoting_, quote_char_, double_quote_, delimiter_); - size_t boundary = end; - bool found = false; - size_t position = start; - scan_range(start, end, [&](const char* data, size_t n) { - if (found) { - position += n; - return; - } - for (size_t i = 0; i < n; ++i) { - const auto before = state.count; - const char c = data[i]; - state.step(c); - if (state.count != before) { - boundary = position + i + 1; - found = true; - break; - } - } - position += n; - }); - if (found && boundary < end) { - std::ifstream file(file_path_, std::ios::binary); - file.seekg(static_cast(boundary)); - if (file && file.peek() == '\n') { - ++boundary; - } - } - return {boundary, end, start_row + state.count}; - } - CsvScanResult scan_parallel(size_t file_size, unsigned num_threads, bool build_ranges) const { // Compute newline-aligned chunk boundaries. @@ -773,12 +739,13 @@ class CsvRowCountCounter { std::vector> safe_bounds; safe_bounds.emplace_back(0, 0); for (unsigned i = 1; i < actual_threads; ++i) { - CsvPartitionRange safe{bounds[i], file_size, rows_before[i]}; - if (starts_inside[i]) { - safe = find_record_boundary(bounds[i], file_size, rows_before[i]); - } - if (safe.start > safe_bounds.back().first && safe.start < file_size) { - safe_bounds.emplace_back(safe.start, safe.start_row); + // A nominal split is already positioned after a physical newline. It is + // a valid record boundary exactly when that newline was outside quotes. + // Unsafe splits are merged into the preceding range instead of rescanning + // the file tail to manufacture another boundary. + if (!starts_inside[i] && bounds[i] > safe_bounds.back().first && + bounds[i] < file_size) { + safe_bounds.emplace_back(bounds[i], rows_before[i]); } } safe_bounds.emplace_back(file_size, total); @@ -793,18 +760,30 @@ class CsvRowCountCounter { std::string file_path_; bool quoting_; char quote_char_; - bool double_quote_; char delimiter_; bool use_threads_; }; -class BoundedFileStreamBuf final : public std::streambuf { +class CsvRangeStreamBuf final : public std::streambuf { public: - BoundedFileStreamBuf(const std::string& file_path, size_t start, size_t end) - : file_(file_path, std::ios::binary), remaining_(end - start) { + CsvRangeStreamBuf(const std::string& file_path, size_t start, size_t end) + : file_(file_path, std::ios::binary), file_path_(file_path) { if (!file_.is_open()) { THROW_IO_EXCEPTION("Failed to open CSV range: " + file_path); } + if (end < start) { + THROW_INVALID_ARGUMENT_EXCEPTION("Invalid CSV byte range: [" + + std::to_string(start) + ", " + + std::to_string(end) + ")"); + } + const auto max_offset = + static_cast(std::numeric_limits::max()); + if (static_cast(start) > max_offset || + static_cast(end) > max_offset) { + THROW_IO_EXCEPTION("CSV byte range exceeds stream offset limit: " + + file_path); + } + remaining_ = end - start; file_.seekg(static_cast(start)); if (!file_) { THROW_IO_EXCEPTION("Failed to seek CSV range: " + file_path); @@ -823,236 +802,290 @@ class BoundedFileStreamBuf final : public std::streambuf { const auto requested = std::min(remaining_, buffer_.size()); file_.read(buffer_.data(), static_cast(requested)); const auto read = static_cast(file_.gcount()); + if (file_.bad()) { + THROW_IO_EXCEPTION("Failed to read CSV range: " + file_path_); + } if (read == 0) { remaining_ = 0; return traits_type::eof(); } remaining_ -= read; + if (read < requested) { + remaining_ = 0; + } setg(buffer_.data(), buffer_.data(), buffer_.data() + read); return traits_type::to_int_type(*gptr()); } + std::streamsize xsgetn(char_type* destination, + std::streamsize count) override { + if (count <= 0) { + return 0; + } + + std::streamsize total = 0; + const auto buffered = static_cast(egptr() - gptr()); + const auto from_buffer = + std::min(buffered, static_cast(count - total)); + if (from_buffer > 0) { + std::memcpy(destination, gptr(), from_buffer); + gbump(static_cast(from_buffer)); + total += static_cast(from_buffer); + } + + if (total == count || remaining_ == 0) { + return total; + } + const auto requested = + std::min(remaining_, static_cast(count - total)); + file_.read(destination + total, static_cast(requested)); + const auto read = static_cast(file_.gcount()); + if (file_.bad()) { + THROW_IO_EXCEPTION("Failed to read CSV range: " + file_path_); + } + remaining_ -= read; + if (read < requested) { + remaining_ = 0; + } + return total + static_cast(read); + } + private: std::ifstream file_; - size_t remaining_; - std::array buffer_{}; + std::string file_path_; + size_t remaining_ = 0; + std::array buffer_{}; }; -class BoundedFileStream final : public std::istream { +class CsvRangeStream final : public std::istream { public: - BoundedFileStream(const std::string& file_path, size_t start, size_t end) + CsvRangeStream(const std::string& file_path, size_t start, size_t end) : std::istream(nullptr), buffer_(file_path, start, end) { rdbuf(&buffer_); } private: - BoundedFileStreamBuf buffer_; + CsvRangeStreamBuf buffer_; }; } // namespace -std::shared_ptr CsvPartitionPlanCache::GetOrCreate( +namespace { + +std::shared_ptr build_csv_partition_plan( const std::vector& file_paths, const CsvReadConfig& config, int32_t producer_count) { - producer_count = std::clamp( - producer_count, 1, chunk_pipeline_detail::hardware_worker_count()); - std::shared_ptr entry; - { - std::lock_guard lock(mutex); - auto& cached = entries[producer_count]; - if (!cached) { - cached = std::make_shared(); + CHECK_GE(producer_count, 1); + auto plan = std::make_shared(); + + // producer_count is a budget for the complete COPY, not for every input + // file. Give every non-empty file one range, then split the currently + // largest range until the global producer budget is reached. + std::vector file_sizes(file_paths.size(), 0); + std::vector range_counts(file_paths.size(), 0); + size_t non_empty_files = 0; + for (size_t i = 0; i < file_paths.size(); ++i) { + std::error_code error; + file_sizes[i] = std::filesystem::file_size(file_paths[i], error); + if (error) { + THROW_IO_EXCEPTION("Failed to get file size: " + file_paths[i]); + } + if (file_sizes[i] > 0) { + range_counts[i] = 1; + ++non_empty_files; } - entry = cached; } - std::call_once(entry->once, [&] { - auto plan = std::make_shared(); - - // producer_count is a budget for the complete COPY, not for every input - // file. Give every non-empty file one range, then split the currently - // largest range until the global producer budget is reached. - std::vector file_sizes(file_paths.size(), 0); - std::vector range_counts(file_paths.size(), 0); - size_t non_empty_files = 0; + size_t assigned_ranges = non_empty_files; + const size_t target_ranges = + std::max(non_empty_files, static_cast(producer_count)); + while (assigned_ranges < target_ranges) { + size_t best = file_paths.size(); + long double best_range_bytes = -1; for (size_t i = 0; i < file_paths.size(); ++i) { - std::error_code error; - file_sizes[i] = std::filesystem::file_size(file_paths[i], error); - if (error) { - THROW_IO_EXCEPTION("Failed to get file size: " + file_paths[i]); + if (range_counts[i] <= 0 || + static_cast(range_counts[i]) >= file_sizes[i]) { + continue; } - if (file_sizes[i] > 0) { - range_counts[i] = 1; - ++non_empty_files; + const auto range_bytes = + static_cast(file_sizes[i]) / range_counts[i]; + if (range_bytes > best_range_bytes) { + best = i; + best_range_bytes = range_bytes; } } + if (best == file_paths.size()) { + break; + } + ++range_counts[best]; + ++assigned_ranges; + } - size_t assigned_ranges = non_empty_files; - const size_t target_ranges = - std::max(non_empty_files, static_cast(producer_count)); - while (assigned_ranges < target_ranges) { + // Row counting/planning is a separate phase, so it can use the complete + // hardware budget. When there are many input files, H workers scan files + // concurrently. With fewer files, the same H budget is divided among + // intra-file speculative scans. scan_parallel() uses its caller for one + // byte range, therefore the number of active scanners never exceeds H. + const auto scan_budget = chunk_pipeline_detail::hardware_worker_count(); + std::vector scan_threads(file_paths.size(), 0); + if (non_empty_files < static_cast(scan_budget)) { + size_t assigned_scanners = 0; + for (size_t i = 0; i < file_paths.size(); ++i) { + if (file_sizes[i] == 0) { + continue; + } + scan_threads[i] = std::max(1, range_counts[i]); + assigned_scanners += static_cast(scan_threads[i]); + } + CHECK_LE(assigned_scanners, static_cast(scan_budget)); + while (assigned_scanners < static_cast(scan_budget)) { size_t best = file_paths.size(); - long double best_range_bytes = -1; + long double best_scanner_bytes = -1; for (size_t i = 0; i < file_paths.size(); ++i) { - if (range_counts[i] <= 0 || - static_cast(range_counts[i]) >= file_sizes[i]) { + if (scan_threads[i] <= 0 || + static_cast(scan_threads[i]) >= file_sizes[i]) { continue; } - const auto range_bytes = - static_cast(file_sizes[i]) / range_counts[i]; - if (range_bytes > best_range_bytes) { + const auto scanner_bytes = + static_cast(file_sizes[i]) / scan_threads[i]; + if (scanner_bytes > best_scanner_bytes) { best = i; - best_range_bytes = range_bytes; + best_scanner_bytes = scanner_bytes; } } if (best == file_paths.size()) { break; } - ++range_counts[best]; - ++assigned_ranges; - } - - // Row counting/planning is a separate phase, so it can use the complete - // hardware budget. When there are many input files, H workers scan files - // concurrently. With fewer files, the same H budget is divided among - // intra-file speculative scans. scan_parallel() uses its caller for one - // byte range, therefore the number of active scanners never exceeds H. - const auto scan_budget = chunk_pipeline_detail::hardware_worker_count(); - std::vector scan_threads(file_paths.size(), 0); - if (non_empty_files < static_cast(scan_budget)) { - size_t assigned_scanners = 0; - for (size_t i = 0; i < file_paths.size(); ++i) { - if (file_sizes[i] == 0) { - continue; - } - scan_threads[i] = std::max(1, range_counts[i]); - assigned_scanners += static_cast(scan_threads[i]); - } - CHECK_LE(assigned_scanners, static_cast(scan_budget)); - while (assigned_scanners < static_cast(scan_budget)) { - size_t best = file_paths.size(); - long double best_scanner_bytes = -1; - for (size_t i = 0; i < file_paths.size(); ++i) { - if (scan_threads[i] <= 0 || - static_cast(scan_threads[i]) >= file_sizes[i]) { - continue; - } - const auto scanner_bytes = - static_cast(file_sizes[i]) / scan_threads[i]; - if (scanner_bytes > best_scanner_bytes) { - best = i; - best_scanner_bytes = scanner_bytes; - } - } - if (best == file_paths.size()) { - break; - } - ++scan_threads[best]; - ++assigned_scanners; - } - } else { - for (size_t i = 0; i < file_paths.size(); ++i) { - scan_threads[i] = file_sizes[i] == 0 ? 0 : 1; - } + ++scan_threads[best]; + ++assigned_scanners; } + } else { + for (size_t i = 0; i < file_paths.size(); ++i) { + scan_threads[i] = file_sizes[i] == 0 ? 0 : 1; + } + } - std::vector scans(file_paths.size()); - std::atomic next_file{0}; - std::atomic scan_cancelled{false}; - std::mutex scan_error_mutex; - std::exception_ptr scan_error; - auto capture_scan_error = [&](std::exception_ptr error) { - bool expected = false; - if (scan_cancelled.compare_exchange_strong(expected, true, - std::memory_order_acq_rel)) { - std::lock_guard lock(scan_error_mutex); - scan_error = std::move(error); - } - }; - auto scan_file = [&](size_t file_index) { - if (file_sizes[file_index] == 0) { - return; - } - scans[file_index] = - CsvRowCountCounter(file_paths[file_index], config.quoting, - config.quote_char, config.double_quote, - config.delimiter, config.use_threads) - .ScanWithRanges(std::max(1, range_counts[file_index]), - std::max(1, scan_threads[file_index])); - }; - - std::vector planners; - if (non_empty_files >= static_cast(scan_budget)) { - planners.reserve(static_cast(scan_budget)); - for (int32_t worker = 0; worker < scan_budget; ++worker) { - planners.emplace_back([&] { - try { - while (!scan_cancelled.load(std::memory_order_acquire)) { - const auto file_index = - next_file.fetch_add(1, std::memory_order_relaxed); - if (file_index >= file_paths.size()) { - break; - } - scan_file(file_index); - } - } catch (...) { capture_scan_error(std::current_exception()); } - }); - } - } else { - planners.reserve(non_empty_files); - for (size_t i = 0; i < file_paths.size(); ++i) { - if (file_sizes[i] == 0) { - continue; - } - planners.emplace_back([&, i] { - try { - if (!scan_cancelled.load(std::memory_order_acquire)) { - scan_file(i); - } - } catch (...) { capture_scan_error(std::current_exception()); } - }); - } + std::vector scans(file_paths.size()); + std::atomic next_file{0}; + std::atomic scan_cancelled{false}; + std::mutex scan_error_mutex; + std::exception_ptr scan_error; + auto capture_scan_error = [&](std::exception_ptr error) { + bool expected = false; + if (scan_cancelled.compare_exchange_strong(expected, true, + std::memory_order_acq_rel)) { + std::lock_guard lock(scan_error_mutex); + scan_error = std::move(error); } - for (auto& planner : planners) { - planner.join(); + }; + auto scan_file = [&](size_t file_index) { + if (file_sizes[file_index] == 0) { + return; } - if (scan_error) { - std::rethrow_exception(scan_error); + if (file_sizes[file_index] > std::numeric_limits::max()) { + THROW_IO_EXCEPTION("CSV file is too large to address: " + + file_paths[file_index]); } + scans[file_index] = + CsvFileScanner(file_paths[file_index], config.quoting, + config.quote_char, config.delimiter, config.use_threads) + .scan_with_ranges(static_cast(file_sizes[file_index]), + std::max(1, range_counts[file_index]), + std::max(1, scan_threads[file_index])); + }; - int64_t total = 0; + std::vector planners; + if (non_empty_files >= static_cast(scan_budget)) { + planners.reserve(static_cast(scan_budget)); + for (int32_t worker = 0; worker < scan_budget; ++worker) { + planners.emplace_back([&] { + try { + while (!scan_cancelled.load(std::memory_order_acquire)) { + const auto file_index = + next_file.fetch_add(1, std::memory_order_relaxed); + if (file_index >= file_paths.size()) { + break; + } + scan_file(file_index); + } + } catch (...) { capture_scan_error(std::current_exception()); } + }); + } + } else { + planners.reserve(non_empty_files); for (size_t i = 0; i < file_paths.size(); ++i) { - const auto& scan = scans[i]; - if (scan.row_count < 0 || - scan.row_count > std::numeric_limits::max() - total) { - total = kUnknownRowNum; - } else if (total != kUnknownRowNum) { - total += scan.row_count; + if (file_sizes[i] == 0) { + continue; } + planners.emplace_back([&, i] { + try { + if (!scan_cancelled.load(std::memory_order_acquire)) { + scan_file(i); + } + } catch (...) { capture_scan_error(std::current_exception()); } + }); + } + } + for (auto& planner : planners) { + planner.join(); + } + if (scan_error) { + std::rethrow_exception(scan_error); + } - // SKIP is scoped per input file. Distribute it over record-aligned - // ranges so a short first range cannot leak rows that still need to be - // skipped into a later producer. - int64_t remaining_skip = std::max(0, config.skip_rows); - for (size_t range_index = 0; range_index < scan.ranges.size(); - ++range_index) { - const auto& range = scan.ranges[range_index]; - const auto range_end_row = range_index + 1 < scan.ranges.size() - ? scan.ranges[range_index + 1].start_row - : scan.row_count; - CHECK_GE(range_end_row, range.start_row); - const auto range_rows = range_end_row - range.start_row; - const auto range_skip = std::min(remaining_skip, range_rows); - remaining_skip -= range_skip; - plan->tasks.push_back({file_paths[i], range, range_skip}); - } + int64_t total = 0; + for (size_t i = 0; i < file_paths.size(); ++i) { + const auto& scan = scans[i]; + if (scan.row_count < 0 || + scan.row_count > std::numeric_limits::max() - total) { + total = kUnknownRowNum; + } else if (total != kUnknownRowNum) { + total += scan.row_count; + } + + // SKIP is scoped per input file. Distribute it over record-aligned ranges + // so a short first range cannot leak rows that still need to be skipped + // into a later producer. + int64_t remaining_skip = std::max(0, config.skip_rows); + for (size_t range_index = 0; range_index < scan.ranges.size(); + ++range_index) { + const auto& range = scan.ranges[range_index]; + const auto range_end_row = range_index + 1 < scan.ranges.size() + ? scan.ranges[range_index + 1].start_row + : scan.row_count; + CHECK_GE(range_end_row, range.start_row); + const auto range_rows = range_end_row - range.start_row; + const auto range_skip = std::min(remaining_skip, range_rows); + remaining_skip -= range_skip; + plan->tasks.push_back({file_paths[i], range, range_skip}); + } + } + plan->row_count = total; + VLOG(1) << "CSV partition plan: files=" << file_paths.size() + << ", ranges=" << plan->tasks.size() + << ", producers=" << producer_count << ", scan_budget=" << scan_budget + << ", rows=" << total; + return plan; +} + +} // namespace + +std::shared_ptr CsvPartitionPlanCache::GetOrCreate( + const std::vector& file_paths, const CsvReadConfig& config, + int32_t producer_count) { + producer_count = std::clamp( + producer_count, 1, chunk_pipeline_detail::hardware_worker_count()); + Entry* entry; + { + std::lock_guard lock(mutex_); + auto& cached = entries_[producer_count]; + if (!cached) { + cached = std::make_unique(); } - plan->row_count = total; - VLOG(1) << "CSV partition plan: files=" << file_paths.size() - << ", ranges=" << plan->tasks.size() - << ", producers=" << producer_count - << ", scan_budget=" << scan_budget << ", rows=" << total; - entry->plan = std::move(plan); + entry = cached.get(); + } + std::call_once(entry->once, [&] { + entry->plan = build_csv_partition_plan(file_paths, config, producer_count); }); return entry->plan; } @@ -1078,7 +1111,6 @@ struct CsvSupplierRuntime { escape_char_(config.escape_char), quoting_(config.quoting), quote_char_(config.quote_char), - double_quote_(config.double_quote), delimiter_(config.delimiter), use_threads_(config.use_threads), range_(std::move(range)) { @@ -1086,9 +1118,9 @@ struct CsvSupplierRuntime { THROW_SCHEMA_MISMATCH("No columns selected for CSV file: " + file_path_); } if (row_count_mode == CsvRowCountMode::kCountOnOpen) { - row_num_ = CsvRowCountCounter(file_path, quoting_, quote_char_, - double_quote_, delimiter_, use_threads_) - .count(); + row_num_ = CsvFileScanner(file_path, quoting_, quote_char_, delimiter_, + use_threads_) + .count_rows(); } if (range_) { csv_format_.threading(false); @@ -1157,9 +1189,12 @@ struct CsvSupplierRuntime { int64_t row_num() const { if (row_num_ == kUnknownRowNum) { - row_num_ = CsvRowCountCounter(file_path_, quoting_, quote_char_, - double_quote_, delimiter_, use_threads_) - .count(); + if (range_) { + return kUnknownRowNum; + } + row_num_ = CsvFileScanner(file_path_, quoting_, quote_char_, delimiter_, + use_threads_) + .count_rows(); } return row_num_; } @@ -1168,7 +1203,7 @@ struct CsvSupplierRuntime { void reset_reader() { try { if (range_) { - range_stream_ = std::make_unique( + range_stream_ = std::make_unique( file_path_, range_->start, range_->end); reader_ = std::make_unique(*range_stream_, csv_format_); } else { @@ -1206,13 +1241,12 @@ struct CsvSupplierRuntime { char escape_char_ = '\\'; bool quoting_ = false; char quote_char_ = '"'; - bool double_quote_ = true; char delimiter_ = ','; bool use_threads_ = true; std::optional range_; mutable int64_t row_num_ = kUnknownRowNum; int64_t current_row_number_ = 0; - std::unique_ptr range_stream_; + std::unique_ptr range_stream_; std::unique_ptr reader_; }; diff --git a/tests/storage/test_edge_table.cc b/tests/storage/test_edge_table.cc index b18d4e7ce..80627167d 100644 --- a/tests/storage/test_edge_table.cc +++ b/tests/storage/test_edge_table.cc @@ -155,6 +155,42 @@ class EdgeTableTest : public ::testing::Test { edge_table->EnsureCapacity(src_v_cap, dst_v_cap); } + void InitEdgeTable(const std::shared_ptr& ckp, + neug::vid_t src_num, neug::vid_t dst_num, + neug::label_t edge_label) { + InitIndexers(*ckp, src_num, dst_num); + ConstructEdgeTable(src_label_, dst_label_, edge_label); + OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), src_num, dst_num); + } + + template + std::vector> MakeIntEdgeChunks( + const std::vector& srcs, const std::vector& dsts, + const std::vector& data, size_t chunk_num) { + return convert_to_data_chunks({split_column_to_chunks(srcs, chunk_num), + split_column_to_chunks(dsts, chunk_num), + split_column_to_chunks(data, chunk_num)}); + } + + template + std::shared_ptr MakeIntEdgeChunk(SrcT src = 0, DstT dst = 0, + int32_t data = 42) { + auto chunks = + MakeIntEdgeChunks(std::vector{src}, std::vector{dst}, + std::vector{data}, 1); + CHECK_EQ(chunks.size(), 1); + return chunks.front(); + } + + template + std::shared_ptr MakeIntEdgeSource( + const std::vector& srcs, const std::vector& dsts, + const std::vector& data, size_t chunk_num, + int64_t estimated_bytes = kForceBulkBuildBytes) { + return std::make_shared( + MakeIntEdgeChunks(srcs, dsts, data, chunk_num), estimated_bytes); + } + void BatchInsert(std::vector>&& chunks) { auto supplier = std::make_shared(std::move(chunks)); edge_table->BatchAddEdges(src_indexer, dst_indexer, supplier); @@ -257,6 +293,48 @@ class EdgeTableTest : public ::testing::Test { } } + template + void ExpectEdges(std::vector> expected, + bool outgoing, bool unordered = false) { + std::vector srcs, dsts; + std::vector data; + if (outgoing) { + OutputOutgoingEndpoints(srcs, dsts, MAX_TIMESTAMP); + OutputOutgoingEdgeData(data, MAX_TIMESTAMP, 0); + } else { + OutputIncomingEndpoints(srcs, dsts, MAX_TIMESTAMP); + OutputIncomingEdgeData(data, MAX_TIMESTAMP, 0); + } + ASSERT_EQ(srcs.size(), expected.size()); + ASSERT_EQ(dsts.size(), expected.size()); + ASSERT_EQ(data.size(), expected.size()); + + std::vector> actual; + actual.reserve(expected.size()); + for (size_t row = 0; row < expected.size(); ++row) { + actual.emplace_back(srcs[row], dsts[row], data[row]); + } + if (unordered) { + std::sort(expected.begin(), expected.end()); + std::sort(actual.begin(), actual.end()); + } + EXPECT_EQ(actual, expected); + } + + template + void ExpectOutgoingEdges( + std::vector> expected, + bool unordered = false) { + ExpectEdges(std::move(expected), true, unordered); + } + + template + void ExpectIncomingEdges( + std::vector> expected, + bool unordered = false) { + ExpectEdges(std::move(expected), false, unordered); + } + neug::vid_t GetSrcLid(const neug::Value& src_oid) { neug::vid_t src_lid; if (!src_indexer.get_index(src_oid, src_lid)) { @@ -739,131 +817,122 @@ TEST_F(EdgeTableTest, TestBatchAddEdgesBundled) { ASSERT_EQ(dsts.size(), edge_num + more_edge_num); } -TEST_F(EdgeTableTest, BatchBuildEdgesBundledWithParallelEncoding) { - auto ckp = make_checkpoint(workspace()); - constexpr int64_t kSrcNum = 100; - constexpr int64_t kDstNum = 100; - constexpr size_t kEdgeNum = 4000; - - auto src_list = generate_random_vertices(kSrcNum, kEdgeNum); - auto dst_list = generate_random_vertices(kDstNum, kEdgeNum); - auto data_list = generate_random_data(kEdgeNum); - auto batches = - convert_to_data_chunks({split_column_to_chunks(src_list, 16), - split_column_to_chunks(dst_list, 16), - split_column_to_chunks(data_list, 16)}); - - InitIndexers(*ckp, kSrcNum, kDstNum); - ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); - OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), kSrcNum, kDstNum); - auto source = std::make_shared(std::move(batches), - 256LL * 1024 * 1024); - BatchBuild(source); - - EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); - EXPECT_EQ(source->OpenCount(), 2); - ASSERT_EQ(source->OpenedProjections().size(), 2); - EXPECT_EQ(source->OpenedProjections()[0], (std::vector{0, 1})); - EXPECT_TRUE(source->OpenedProjections()[1].empty()); - std::vector output_srcs, output_dsts; - OutputOutgoingEndpoints(output_srcs, output_dsts, neug::MAX_TIMESTAMP); - ASSERT_EQ(output_srcs.size(), kEdgeNum); - ASSERT_EQ(output_dsts.size(), kEdgeNum); - - std::vector> expected; - std::vector> actual; - expected.reserve(kEdgeNum); - actual.reserve(kEdgeNum); - for (size_t i = 0; i < kEdgeNum; ++i) { - expected.emplace_back(src_list[i], dst_list[i], data_list[i]); - } - std::vector output_data; - OutputOutgoingEdgeData(output_data, neug::MAX_TIMESTAMP, 0); - ASSERT_EQ(output_data.size(), kEdgeNum); - for (size_t i = 0; i < kEdgeNum; ++i) { - actual.emplace_back(output_srcs[i], output_dsts[i], output_data[i]); - } - std::sort(expected.begin(), expected.end()); - std::sort(actual.begin(), actual.end()); - EXPECT_EQ(actual, expected); -} - -TEST_F(EdgeTableTest, BatchBuildEdgesParallelFillWithoutProperties) { - auto ckp = make_checkpoint(workspace()); - constexpr int64_t kVertexNum = 32; - constexpr size_t kEdgeNum = 2048; - std::vector srcs(kEdgeNum); - std::vector dsts(kEdgeNum); - for (size_t row = 0; row < kEdgeNum; ++row) { - srcs[row] = static_cast(row % kVertexNum); - dsts[row] = static_cast((row * 7) % kVertexNum); - } - auto chunks = convert_to_data_chunks( - {split_column_to_chunks(srcs, 32), split_column_to_chunks(dsts, 32)}); - - InitIndexers(*ckp, kVertexNum, kVertexNum); - ConstructEdgeTable(src_label_, dst_label_, edge_label_empty_); - OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); - auto source = std::make_shared(std::move(chunks), - kForceBulkBuildBytes); - BatchBuild(source); - - EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); - std::vector actual_srcs, actual_dsts; - OutputOutgoingEndpoints(actual_srcs, actual_dsts, MAX_TIMESTAMP); - ASSERT_EQ(actual_srcs.size(), kEdgeNum); - ASSERT_EQ(actual_dsts.size(), kEdgeNum); - std::vector> expected; - std::vector> actual; - expected.reserve(kEdgeNum); - actual.reserve(kEdgeNum); - for (size_t row = 0; row < kEdgeNum; ++row) { - expected.emplace_back(srcs[row], dsts[row]); - actual.emplace_back(actual_srcs[row], actual_dsts[row]); - } - std::sort(expected.begin(), expected.end()); - std::sort(actual.begin(), actual.end()); - EXPECT_EQ(actual, expected); -} +TEST_F(EdgeTableTest, BatchBuildEdgesHandlesParallelVariants) { + auto with_property = [&] { + SCOPED_TRACE("bundled int property"); + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kSrcNum = 100; + constexpr int64_t kDstNum = 100; + constexpr size_t kEdgeNum = 4000; + auto srcs = generate_random_vertices(kSrcNum, kEdgeNum); + auto dsts = generate_random_vertices(kDstNum, kEdgeNum); + auto data = generate_random_data(kEdgeNum); + InitEdgeTable(ckp, kSrcNum, kDstNum, edge_label_int_); + auto source = MakeIntEdgeSource(srcs, dsts, data, 16, 256LL * 1024 * 1024); + BatchBuild(source); + + EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); + EXPECT_EQ(source->OpenCount(), 2); + ASSERT_EQ(source->OpenedProjections().size(), 2); + EXPECT_EQ(source->OpenedProjections()[0], (std::vector{0, 1})); + EXPECT_TRUE(source->OpenedProjections()[1].empty()); + std::vector> expected; + expected.reserve(kEdgeNum); + for (size_t row = 0; row < kEdgeNum; ++row) { + expected.emplace_back(srcs[row], dsts[row], data[row]); + } + ExpectOutgoingEdges(std::move(expected), true); + }; -TEST_F(EdgeTableTest, BatchBuildEdgesHandlesEmptyInput) { - auto ckp = make_checkpoint(workspace()); - constexpr int64_t kVertexNum = 16; - InitIndexers(*ckp, kVertexNum, kVertexNum); - ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); - OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); - auto source = std::make_shared( - std::vector>{}, kForceBulkBuildBytes); + auto without_properties = [&] { + SCOPED_TRACE("no properties"); + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kVertexNum = 32; + constexpr size_t kEdgeNum = 2048; + std::vector srcs(kEdgeNum); + std::vector dsts(kEdgeNum); + for (size_t row = 0; row < kEdgeNum; ++row) { + srcs[row] = static_cast(row % kVertexNum); + dsts[row] = static_cast((row * 7) % kVertexNum); + } + auto chunks = convert_to_data_chunks( + {split_column_to_chunks(srcs, 32), split_column_to_chunks(dsts, 32)}); + InitEdgeTable(ckp, kVertexNum, kVertexNum, edge_label_empty_); + auto source = std::make_shared(std::move(chunks), + kForceBulkBuildBytes); + BatchBuild(source); + + EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); + std::vector actual_srcs, actual_dsts; + OutputOutgoingEndpoints(actual_srcs, actual_dsts, MAX_TIMESTAMP); + ASSERT_EQ(actual_srcs.size(), kEdgeNum); + ASSERT_EQ(actual_dsts.size(), kEdgeNum); + std::vector> expected; + std::vector> actual; + expected.reserve(kEdgeNum); + actual.reserve(kEdgeNum); + for (size_t row = 0; row < kEdgeNum; ++row) { + expected.emplace_back(srcs[row], dsts[row]); + actual.emplace_back(actual_srcs[row], actual_dsts[row]); + } + std::sort(expected.begin(), expected.end()); + std::sort(actual.begin(), actual.end()); + EXPECT_EQ(actual, expected); + }; - BatchBuild(source); + auto supernode = [&] { + SCOPED_TRACE("supernode"); + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kVertexNum = 4; + constexpr size_t kEdgeNum = 8192; + std::vector srcs(kEdgeNum, 0); + std::vector dsts(kEdgeNum, 1); + std::vector data(kEdgeNum); + for (size_t row = 0; row < kEdgeNum; ++row) { + data[row] = static_cast(row); + } + InitEdgeTable(ckp, kVertexNum, kVertexNum, edge_label_int_); + auto source = MakeIntEdgeSource(srcs, dsts, data, 64); + BatchBuild(source); + + EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); + std::vector> expected; + expected.reserve(kEdgeNum); + for (const auto value : data) { + expected.emplace_back(0, 1, value); + } + ExpectOutgoingEdges(expected, true); + ExpectIncomingEdges(std::move(expected), true); + }; - EXPECT_EQ(source->OpenCount(), 2); - EXPECT_EQ(edge_table->EdgeNum(), 0); - std::vector output_srcs, output_dsts; - OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); - EXPECT_TRUE(output_srcs.empty()); - EXPECT_TRUE(output_dsts.empty()); + with_property(); + without_properties(); + supernode(); } -TEST_F(EdgeTableTest, BatchBuildEdgesSkipsSourceForNoAdjacency) { - auto ckp = make_checkpoint(workspace()); - constexpr int64_t kVertexNum = 4; - auto chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{0, 1, 2}, 2), - split_column_to_chunks(std::vector{1, 2, 3}, 2), - split_column_to_chunks(std::vector{10, 20, 30}, 2)}); - - InitIndexers(*ckp, kVertexNum, kVertexNum); - ConstructEdgeTable(src_label_, dst_label_, edge_label_none_); - OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); - auto source = std::make_shared(std::move(chunks), - kForceBulkBuildBytes); - - BatchBuild(source); +TEST_F(EdgeTableTest, BatchBuildEdgesHandlesInputsWithNoStoredEdges) { + auto run = [&](neug::label_t edge_label, + std::vector> chunks, + size_t expected_open_count) { + SCOPED_TRACE(edge_label); + auto ckp = make_checkpoint(workspace()); + constexpr neug::vid_t kVertexNum = 4; + InitEdgeTable(ckp, kVertexNum, kVertexNum, edge_label); + auto source = std::make_shared(std::move(chunks), + kForceBulkBuildBytes); + + BatchBuild(source); + + EXPECT_EQ(source->OpenCount(), expected_open_count); + EXPECT_EQ(edge_table->EdgeNum(), 0); + }; - EXPECT_EQ(source->OpenCount(), 0); - EXPECT_TRUE(source->OpenedProjections().empty()); - EXPECT_EQ(edge_table->EdgeNum(), 0); + run(edge_label_int_, {}, 2); + run(edge_label_none_, + MakeIntEdgeChunks(std::vector{0, 1, 2}, + std::vector{1, 2, 3}, + std::vector{10, 20, 30}, 2), + 0); } TEST_F(EdgeTableTest, SecondBatchFallsBackWhenIncomingCsrIsNotEmpty) { @@ -873,77 +942,19 @@ TEST_F(EdgeTableTest, SecondBatchFallsBackWhenIncomingCsrIsNotEmpty) { ConstructEdgeTable(src_label_, dst_label_, edge_label_out_none_); OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); - auto first_chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{10}, 1)}); - auto first_source = std::make_shared( - std::move(first_chunks), kForceBulkBuildBytes); + auto first_source = + MakeIntEdgeSource(std::vector{0}, std::vector{0}, + std::vector{10}, 1); BatchBuild(first_source); EXPECT_EQ(first_source->OpenCount(), 2); - auto second_chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{1}, 1), - split_column_to_chunks(std::vector{1}, 1), - split_column_to_chunks(std::vector{20}, 1)}); - auto second_source = std::make_shared( - std::move(second_chunks), kForceBulkBuildBytes); + auto second_source = + MakeIntEdgeSource(std::vector{1}, std::vector{1}, + std::vector{20}, 1); BatchBuild(second_source); EXPECT_EQ(second_source->OpenCount(), 1); - std::vector incoming_srcs, incoming_dsts; - std::vector incoming_data; - OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); - OutputIncomingEdgeData(incoming_data, MAX_TIMESTAMP, 0); - EXPECT_EQ(incoming_srcs, (std::vector{0, 1})); - EXPECT_EQ(incoming_dsts, (std::vector{0, 1})); - EXPECT_EQ(incoming_data, (std::vector{10, 20})); -} - -TEST_F(EdgeTableTest, BatchBuildEdgesParallelFillHandlesSupernodes) { - auto ckp = make_checkpoint(workspace()); - constexpr int64_t kVertexNum = 4; - constexpr size_t kEdgeNum = 8192; - std::vector srcs(kEdgeNum, 0); - std::vector dsts(kEdgeNum, 1); - std::vector data(kEdgeNum); - for (size_t row = 0; row < kEdgeNum; ++row) { - data[row] = static_cast(row); - } - auto chunks = convert_to_data_chunks({split_column_to_chunks(srcs, 64), - split_column_to_chunks(dsts, 64), - split_column_to_chunks(data, 64)}); - - InitIndexers(*ckp, kVertexNum, kVertexNum); - ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); - OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); - auto source = std::make_shared(std::move(chunks), - kForceBulkBuildBytes); - BatchBuild(source); - - EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); - std::vector output_srcs, output_dsts; - std::vector output_data; - OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); - OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); - ASSERT_EQ(output_srcs.size(), kEdgeNum); - ASSERT_EQ(output_dsts.size(), kEdgeNum); - ASSERT_EQ(output_data.size(), kEdgeNum); - EXPECT_TRUE(std::all_of(output_srcs.begin(), output_srcs.end(), - [](int64_t src) { return src == 0; })); - EXPECT_TRUE(std::all_of(output_dsts.begin(), output_dsts.end(), - [](int64_t dst) { return dst == 1; })); - std::sort(output_data.begin(), output_data.end()); - EXPECT_EQ(output_data, data); - - std::vector incoming_srcs, incoming_dsts; - OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); - ASSERT_EQ(incoming_srcs.size(), kEdgeNum); - ASSERT_EQ(incoming_dsts.size(), kEdgeNum); - EXPECT_TRUE(std::all_of(incoming_srcs.begin(), incoming_srcs.end(), - [](int64_t src) { return src == 0; })); - EXPECT_TRUE(std::all_of(incoming_dsts.begin(), incoming_dsts.end(), - [](int64_t dst) { return dst == 1; })); + ExpectIncomingEdges({{0, 0, 10}, {1, 1, 20}}); } TEST_F(EdgeTableTest, SingleEdgeBulkBuildFillsUniqueSlotsAcrossChunks) { @@ -967,14 +978,10 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildFillsUniqueSlotsAcrossChunks) { dsts[row] = row % kDstNum; data[row] = static_cast(row * 3); } - auto chunks = convert_to_data_chunks({split_column_to_chunks(srcs, 128), - split_column_to_chunks(dsts, 128), - split_column_to_chunks(data, 128)}); + auto chunks = MakeIntEdgeChunks(srcs, dsts, data, 128); ASSERT_EQ(chunks.size(), 128); - InitIndexers(*ckp, kSrcNum, kDstNum); - ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); - OpenEdgeTableInMemory(ckp, CheckpointManifest(), kSrcNum, kDstNum); + InitEdgeTable(ckp, kSrcNum, kDstNum, edge_label_single_); auto shared_chunks = std::make_shared>>( std::move(chunks)); @@ -1015,68 +1022,53 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildFillsUniqueSlotsAcrossChunks) { EXPECT_GT(activities[1]->max_active.load(), 1); EXPECT_EQ(edge_table->EdgeNum(), kSrcNum); - std::vector output_srcs, output_dsts; - std::vector output_data; - OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); - OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); - EXPECT_EQ(output_srcs, srcs); - EXPECT_EQ(output_dsts, dsts); - EXPECT_EQ(output_data, data); - - std::vector incoming_srcs, incoming_dsts; - std::vector incoming_data; - OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); - OutputIncomingEdgeData(incoming_data, MAX_TIMESTAMP, 0); - ASSERT_EQ(incoming_srcs.size(), srcs.size()); - ASSERT_EQ(incoming_dsts.size(), dsts.size()); - ASSERT_EQ(incoming_data.size(), data.size()); std::vector> expected; - std::vector> actual; expected.reserve(srcs.size()); - actual.reserve(srcs.size()); for (size_t row = 0; row < srcs.size(); ++row) { expected.emplace_back(srcs[row], dsts[row], data[row]); - actual.emplace_back(incoming_srcs[row], incoming_dsts[row], - incoming_data[row]); } - std::sort(expected.begin(), expected.end()); - std::sort(actual.begin(), actual.end()); - EXPECT_EQ(actual, expected); + ExpectOutgoingEdges(expected); + ExpectIncomingEdges(std::move(expected), true); } -TEST_F(EdgeTableTest, SingleEdgeBulkBuildSkipsInvalidEndpoints) { - auto ckp = make_checkpoint(workspace()); - InitIndexers(*ckp, 3, 3); - ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); - OpenEdgeTableInMemory(ckp, CheckpointManifest(), 3, 3); - - const std::vector srcs = {0, 99, 1, 2}; - const std::vector dsts = {0, 1, 99, 2}; - const std::vector data = {10, 20, 30, 40}; - auto chunks = convert_to_data_chunks({split_column_to_chunks(srcs, 4), - split_column_to_chunks(dsts, 4), - split_column_to_chunks(data, 4)}); - auto source = std::make_shared(std::move(chunks), - kForceBulkBuildBytes); - BatchBuild(source); +TEST_F(EdgeTableTest, SingleEdgeBulkBuildHandlesEndpointAndDirectionVariants) { + auto skips_invalid_endpoints = [&] { + SCOPED_TRACE("invalid endpoints"); + auto ckp = make_checkpoint(workspace()); + InitEdgeTable(ckp, 3, 3, edge_label_single_); + auto source = MakeIntEdgeSource(std::vector{0, 99, 1, 2}, + std::vector{0, 1, 99, 2}, + std::vector{10, 20, 30, 40}, 4); + BatchBuild(source); + + EXPECT_EQ(source->OpenCount(), 2); + EXPECT_EQ(edge_table->EdgeNum(), 2); + const std::vector> expected = { + {0, 0, 10}, {2, 2, 40}}; + ExpectOutgoingEdges(expected); + ExpectIncomingEdges(expected); + }; - EXPECT_EQ(source->OpenCount(), 2); - EXPECT_EQ(edge_table->EdgeNum(), 2); - std::vector output_srcs, output_dsts; - std::vector output_data; - OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); - OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); - EXPECT_EQ(output_srcs, (std::vector{0, 2})); - EXPECT_EQ(output_dsts, (std::vector{0, 2})); - EXPECT_EQ(output_data, (std::vector{10, 40})); - - std::vector incoming_srcs, incoming_dsts; - std::vector incoming_data; - OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); - OutputIncomingEdgeData(incoming_data, MAX_TIMESTAMP, 0); - EXPECT_EQ(incoming_srcs, (std::vector{0, 2})); - EXPECT_EQ(incoming_dsts, (std::vector{0, 2})); - EXPECT_EQ(incoming_data, (std::vector{10, 40})); + auto preserves_both_single_directions = [&] { + SCOPED_TRACE("single in both directions"); + auto ckp = make_checkpoint(workspace()); + InitEdgeTable(ckp, 2, 2, edge_label_single_both_); + auto source = MakeIntEdgeSource(std::vector{0, 1, 0}, + std::vector{0, 0, 1}, + std::vector{10, 20, 30}, 1); + BatchBuild(source); + + EXPECT_EQ(source->OpenCount(), 1); + ASSERT_EQ(source->OpenedProjections().size(), 1); + EXPECT_TRUE(source->OpenedProjections()[0].empty()); + EXPECT_EQ(edge_table->EdgeNum(), 3); + + ExpectOutgoingEdges({{0, 1, 30}, {1, 0, 20}}); + ExpectIncomingEdges({{1, 0, 20}, {0, 1, 30}}); + }; + + skips_invalid_endpoints(); + preserves_both_single_directions(); } TEST_F(EdgeTableTest, SingleEdgeBulkBuildDetectsCrossWorkerDuplicate) { @@ -1094,14 +1086,10 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildDetectsCrossWorkerDuplicate) { }; auto ckp = make_checkpoint(workspace()); - InitIndexers(*ckp, 1, 2); - ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); - OpenEdgeTableInMemory(ckp, CheckpointManifest(), 1, 2); - - auto chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{0, 0}, 2), - split_column_to_chunks(std::vector{0, 1}, 2), - split_column_to_chunks(std::vector{10, 20}, 2)}); + InitEdgeTable(ckp, 1, 2, edge_label_single_); + auto chunks = + MakeIntEdgeChunks(std::vector{0, 0}, std::vector{0, 1}, + std::vector{10, 20}, 2); ASSERT_EQ(chunks.size(), 2); auto shared_chunks = std::make_shared>>( @@ -1150,89 +1138,7 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildDetectsCrossWorkerDuplicate) { EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); EXPECT_TRUE(source->OpenedOptions()[1].preserve_order); EXPECT_EQ(edge_table->EdgeNum(), 2); - std::vector output_srcs, output_dsts; - std::vector output_data; - OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); - OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); - EXPECT_EQ(output_srcs, (std::vector{0})); - EXPECT_EQ(output_dsts, (std::vector{1})); - EXPECT_EQ(output_data, (std::vector{20})); -} - -TEST_F(EdgeTableTest, SingleEdgeBulkBuildPreservesLastWriteOrderAcrossChunks) { - auto ckp = make_checkpoint(workspace()); - InitIndexers(*ckp, 1, 3); - ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); - OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), 1, 3); - - const std::vector srcs = {0, 0, 0}; - const std::vector dsts = {0, 1, 2}; - const std::vector data = {10, 20, 30}; - auto chunks = convert_to_data_chunks({split_column_to_chunks(srcs, 3), - split_column_to_chunks(dsts, 3), - split_column_to_chunks(data, 3)}); - ASSERT_EQ(chunks.size(), 3); - auto shared_chunks = - std::make_shared>>( - std::move(chunks)); - auto source = std::make_shared( - [shared_chunks](const ChunkSourceOptions&, size_t) { - return std::make_shared(*shared_chunks, true); - }, - kForceBulkBuildBytes); - BatchBuild(source); - - std::vector output_srcs, output_dsts; - std::vector output_data; - OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); - OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); - ASSERT_EQ(output_srcs, (std::vector{0})); - ASSERT_EQ(output_dsts, (std::vector{2})); - ASSERT_EQ(output_data, (std::vector{30})); - EXPECT_EQ(source->OpenCount(), 2); - ASSERT_EQ(source->OpenedOptions().size(), 2); - EXPECT_EQ(source->OpenedOptions()[0].projected_columns, - (std::vector{0, 1})); - EXPECT_TRUE(source->OpenedOptions()[1].projected_columns.empty()); - EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); - EXPECT_TRUE(source->OpenedOptions()[1].preserve_order); - EXPECT_EQ(edge_table->EdgeNum(), 3); -} - -TEST_F(EdgeTableTest, SingleOnlyBulkBuildPreservesBothDirections) { - auto ckp = make_checkpoint(workspace()); - InitIndexers(*ckp, 2, 2); - ConstructEdgeTable(src_label_, dst_label_, edge_label_single_both_); - OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), 2, 2); - - auto chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{0, 1, 0}, 1), - split_column_to_chunks(std::vector{0, 0, 1}, 1), - split_column_to_chunks(std::vector{10, 20, 30}, 1)}); - auto source = std::make_shared(std::move(chunks), - kForceBulkBuildBytes); - BatchBuild(source); - - EXPECT_EQ(source->OpenCount(), 1); - ASSERT_EQ(source->OpenedProjections().size(), 1); - EXPECT_TRUE(source->OpenedProjections()[0].empty()); - EXPECT_EQ(edge_table->EdgeNum(), 3); - - std::vector output_srcs, output_dsts; - std::vector output_data; - OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); - OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); - ASSERT_EQ(output_srcs, (std::vector{0, 1})); - ASSERT_EQ(output_dsts, (std::vector{1, 0})); - ASSERT_EQ(output_data, (std::vector{30, 20})); - - std::vector incoming_srcs, incoming_dsts; - std::vector incoming_data; - OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); - OutputIncomingEdgeData(incoming_data, MAX_TIMESTAMP, 0); - ASSERT_EQ(incoming_srcs, (std::vector{1, 0})); - ASSERT_EQ(incoming_dsts, (std::vector{0, 1})); - ASSERT_EQ(incoming_data, (std::vector{20, 30})); + ExpectOutgoingEdges({{0, 1, 20}}); } TEST_F(EdgeTableTest, BatchBuildEdgesFromPartitionedCsvInTwoPasses) { @@ -1287,192 +1193,102 @@ TEST_F(EdgeTableTest, BatchBuildEdgesFromPartitionedCsvInTwoPasses) { EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); EXPECT_EQ(source->OpenCount(), 2); - std::vector output_srcs, output_dsts; - std::vector output_data; - OutputOutgoingEndpoints(output_srcs, output_dsts, MAX_TIMESTAMP); - OutputOutgoingEdgeData(output_data, MAX_TIMESTAMP, 0); - ASSERT_EQ(output_srcs.size(), kEdgeNum); - ASSERT_EQ(output_dsts.size(), kEdgeNum); - ASSERT_EQ(output_data.size(), kEdgeNum); - std::vector> actual; - actual.reserve(kEdgeNum); - for (size_t row = 0; row < kEdgeNum; ++row) { - actual.emplace_back(output_srcs[row], output_dsts[row], output_data[row]); - } - std::sort(expected.begin(), expected.end()); - std::sort(actual.begin(), actual.end()); - EXPECT_EQ(actual, expected); - - std::vector incoming_srcs, incoming_dsts; - OutputIncomingEndpoints(incoming_srcs, incoming_dsts, MAX_TIMESTAMP); - EXPECT_EQ(incoming_srcs.size(), kEdgeNum); - EXPECT_EQ(incoming_dsts.size(), kEdgeNum); -} - -TEST_F(EdgeTableTest, FailedBatchBuildDoesNotPublishStagedCsr) { - auto ckp = make_checkpoint(workspace()); - InitIndexers(*ckp, 1, 1); - ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); - OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), 1, 1); - auto chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{42}, 1)}); - ASSERT_EQ(chunks.size(), 1); - - auto source = std::make_shared( - [chunk = chunks.front()](const ChunkSourceOptions&, size_t) { - auto remaining = std::make_shared>(chunk); - return std::make_shared( - [remaining] { - if (*remaining) { - return std::exchange(*remaining, std::shared_ptr{}); - } - throw std::runtime_error("injected bulk edge source failure"); - }, - 1); - }, - kForceBulkBuildBytes); - EXPECT_THROW(BatchBuild(source), std::runtime_error); - EXPECT_EQ(edge_table->EdgeNum(), 0); + ExpectOutgoingEdges(expected, true); + ExpectIncomingEdges(std::move(expected), true); } -TEST_F(EdgeTableTest, - ConcurrentSingleCheckFailureCancelsSupplierAndDoesNotPublish) { - if (std::thread::hardware_concurrency() < 4) { - GTEST_SKIP() - << "Concurrent uniqueness cancellation needs two consumer workers"; - } - - auto ckp = make_checkpoint(workspace()); - InitIndexers(*ckp, 1, 1); - ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); - OpenEdgeTableInMemory(ckp, CheckpointManifest(), 1, 1); - auto chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{42}, 1)}); - ASSERT_EQ(chunks.size(), 1); - std::atomic cancel_count{0}; - auto source = std::make_shared( - [chunk = chunks.front(), &cancel_count](const ChunkSourceOptions&, - size_t) { - auto next_call = std::make_shared>(0); - return std::make_shared( - [chunk, next_call]() -> std::shared_ptr { - const auto call = - next_call->fetch_add(1, std::memory_order_relaxed); - if (call == 0) { - return chunk; - } - if (call == 1) { - throw std::runtime_error( - "injected concurrent uniqueness failure"); - } - return nullptr; - }, - 1, true, - [&cancel_count] { - cancel_count.fetch_add(1, std::memory_order_relaxed); - }); - }, - kForceBulkBuildBytes); - - EXPECT_THROW(BatchBuild(source), std::runtime_error); - EXPECT_EQ(source->OpenCount(), 1); - EXPECT_EQ(cancel_count.load(std::memory_order_relaxed), 1); - ASSERT_EQ(source->OpenedOptions().size(), 1); - EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); - EXPECT_EQ(source->OpenedOptions()[0].projected_columns, - (std::vector{0, 1})); - EXPECT_EQ(edge_table->EdgeNum(), 0); -} - -TEST_F(EdgeTableTest, SecondPassFailureDoesNotPublishAllocatedCsr) { - auto ckp = make_checkpoint(workspace()); - InitIndexers(*ckp, 1, 1); - ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); - OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), 1, 1); - auto chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{42}, 1)}); - ASSERT_EQ(chunks.size(), 1); - auto source = std::make_shared( - [chunk = chunks.front()](const ChunkSourceOptions&, size_t open_index) { - if (open_index == 0) { +TEST_F(EdgeTableTest, BatchBuildFailuresDoNotPublishPartialCsr) { + for (size_t failure_pass : {size_t{0}, size_t{1}}) { + SCOPED_TRACE(::testing::Message() << "failure pass " << failure_pass); + auto ckp = make_checkpoint(workspace()); + InitEdgeTable(ckp, 1, 1, edge_label_int_); + auto chunk = MakeIntEdgeChunk(); + + auto source = std::make_shared( + [chunk, failure_pass](const ChunkSourceOptions&, size_t open_index) { + if (open_index != failure_pass) { + return std::shared_ptr( + std::make_shared( + std::vector>{chunk})); + } + auto remaining = std::make_shared>(chunk); return std::shared_ptr( - std::make_shared( - std::vector>{chunk})); - } - auto remaining = std::make_shared>(chunk); - return std::shared_ptr( - std::make_shared( - [remaining]() -> std::shared_ptr { - if (*remaining) { - return std::exchange(*remaining, nullptr); - } - throw std::runtime_error("injected second-pass failure"); - }, - 1)); - }, - kForceBulkBuildBytes); + std::make_shared( + [remaining] { + if (*remaining) { + return std::exchange(*remaining, + std::shared_ptr{}); + } + throw std::runtime_error("injected bulk build failure"); + }, + 1)); + }, + kForceBulkBuildBytes); - EXPECT_THROW(BatchBuild(source), std::runtime_error); - EXPECT_EQ(source->OpenCount(), 2); - EXPECT_EQ(edge_table->EdgeNum(), 0); + EXPECT_THROW(BatchBuild(source), std::runtime_error); + EXPECT_EQ(source->OpenCount(), failure_pass + 1); + EXPECT_EQ(edge_table->EdgeNum(), 0); + } } TEST_F(EdgeTableTest, - ConcurrentSecondPassFailureCancelsSupplierAndDoesNotPublish) { + ConcurrentBatchBuildFailuresCancelSupplierAndDoNotPublish) { if (std::thread::hardware_concurrency() < 2) { GTEST_SKIP() << "Concurrent supplier cancellation requires two workers"; } - auto ckp = make_checkpoint(workspace()); - InitIndexers(*ckp, 1, 1); - ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); - OpenEdgeTableInMemory(ckp, CheckpointManifest(), 1, 1); - auto chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{42}, 1)}); - ASSERT_EQ(chunks.size(), 1); - std::atomic cancel_count{0}; - auto source = std::make_shared( - [chunk = chunks.front(), &cancel_count](const ChunkSourceOptions&, - size_t open_index) { - if (open_index == 0) { + auto run = [&](neug::label_t edge_label, size_t failure_pass) { + SCOPED_TRACE(::testing::Message() << "failure pass " << failure_pass); + auto ckp = make_checkpoint(workspace()); + InitEdgeTable(ckp, 1, 1, edge_label); + auto chunk = MakeIntEdgeChunk(); + std::atomic cancel_count{0}; + auto source = std::make_shared( + [chunk, failure_pass, &cancel_count](const ChunkSourceOptions&, + size_t open_index) { + if (open_index != failure_pass) { + return std::shared_ptr( + std::make_shared( + std::vector>{chunk})); + } + auto next_call = std::make_shared>(0); return std::shared_ptr( - std::make_shared( - std::vector>{chunk})); - } - auto next_call = std::make_shared>(0); - return std::shared_ptr( - std::make_shared( - [chunk, next_call]() -> std::shared_ptr { - const auto call = - next_call->fetch_add(1, std::memory_order_relaxed); - if (call == 0) { - return chunk; - } - if (call == 1) { - throw std::runtime_error( - "injected concurrent fill failure"); - } - return nullptr; - }, - 1, true, - [&cancel_count] { - cancel_count.fetch_add(1, std::memory_order_relaxed); - })); - }, - kForceBulkBuildBytes); + std::make_shared( + [chunk, next_call]() -> std::shared_ptr { + const auto call = + next_call->fetch_add(1, std::memory_order_relaxed); + if (call == 0) { + return chunk; + } + if (call == 1) { + throw std::runtime_error( + "injected concurrent bulk build failure"); + } + return nullptr; + }, + 1, true, + [&cancel_count] { + cancel_count.fetch_add(1, std::memory_order_relaxed); + })); + }, + kForceBulkBuildBytes); + + EXPECT_THROW(BatchBuild(source), std::runtime_error); + EXPECT_EQ(source->OpenCount(), failure_pass + 1); + EXPECT_EQ(cancel_count.load(std::memory_order_relaxed), 1); + EXPECT_EQ(edge_table->EdgeNum(), 0); + if (failure_pass == 0) { + ASSERT_EQ(source->OpenedOptions().size(), 1); + EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); + EXPECT_EQ(source->OpenedOptions()[0].projected_columns, + (std::vector{0, 1})); + } + }; - EXPECT_THROW(BatchBuild(source), std::runtime_error); - EXPECT_EQ(source->OpenCount(), 2); - EXPECT_EQ(cancel_count.load(std::memory_order_relaxed), 1); - EXPECT_EQ(edge_table->EdgeNum(), 0); + run(edge_label_int_, 1); + if (std::thread::hardware_concurrency() >= 4) { + run(edge_label_single_, 0); + } } TEST_F(EdgeTableTest, @@ -1491,27 +1307,15 @@ TEST_F(EdgeTableTest, }; auto ckp = make_checkpoint(workspace()); - InitIndexers(*ckp, 1, 1); - ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); - OpenEdgeTableInMemory(ckp, CheckpointManifest(), 1, 1); - - auto valid_chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{42}, 1)}); + InitEdgeTable(ckp, 1, 1, edge_label_int_); + auto valid_chunk = MakeIntEdgeChunk(); // The source indexer expects int64 endpoints. The uint64 source column makes // a fill consumer throw after the producer has requested its next chunk. - auto invalid_chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{42}, 1)}); - ASSERT_EQ(valid_chunks.size(), 1); - ASSERT_EQ(invalid_chunks.size(), 1); + auto invalid_chunk = MakeIntEdgeChunk(uint64_t{0}); std::shared_ptr second_pass_state; auto source = std::make_shared( - [valid_chunk = valid_chunks.front(), - invalid_chunk = invalid_chunks.front(), - &second_pass_state](const ChunkSourceOptions&, size_t open_index) { + [valid_chunk, invalid_chunk, &second_pass_state]( + const ChunkSourceOptions&, size_t open_index) { if (open_index == 0) { return std::shared_ptr( std::make_shared( @@ -1553,57 +1357,39 @@ TEST_F(EdgeTableTest, EXPECT_EQ(edge_table->EdgeNum(), 0); } -TEST_F(EdgeTableTest, BatchBuildEdgesUsesIndexerVertexCapacity) { - auto ckp = make_checkpoint(workspace()); - constexpr neug::vid_t kVertexNum = 4097; - constexpr neug::vid_t kCheckpointCapacity = kVertexNum + kVertexNum / 4; - std::vector endpoints = {0}; - std::vector edge_data = {42}; - auto batches = - convert_to_data_chunks({split_column_to_chunks(endpoints, 16), - split_column_to_chunks(endpoints, 16), - split_column_to_chunks(edge_data, 16)}); +TEST_F(EdgeTableTest, BatchInsertPathsUseIndexerVertexCapacity) { + auto run = [&](bool bulk_build, neug::vid_t vertex_num, + neug::vid_t vertex_capacity) { + SCOPED_TRACE(bulk_build ? "bulk build" : "incremental insert"); + auto ckp = make_checkpoint(workspace()); + InitIndexers(*ckp, vertex_num, vertex_num); + src_indexer.reserve(vertex_capacity); + dst_indexer.reserve(vertex_capacity); + ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); + OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), vertex_num, + vertex_num); + auto chunks = + MakeIntEdgeChunks(std::vector{0}, std::vector{0}, + std::vector{42}, 1); + + if (bulk_build) { + BatchBuild(std::move(chunks)); + } else { + BatchInsert(std::move(chunks)); + } - InitIndexers(*ckp, kVertexNum, kVertexNum); - src_indexer.reserve(kCheckpointCapacity); - dst_indexer.reserve(kCheckpointCapacity); - ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); - OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), kVertexNum, - kVertexNum); - BatchBuild(std::move(batches)); - - EXPECT_EQ(edge_table->EdgeNum(), 1); - auto out_csr = edge_table->TakeOutCsr(); - auto in_csr = edge_table->TakeInCsr(); - ASSERT_NE(out_csr, nullptr); - ASSERT_NE(in_csr, nullptr); - EXPECT_EQ(out_csr->size(), kCheckpointCapacity); - EXPECT_EQ(in_csr->size(), kCheckpointCapacity); -} + EXPECT_EQ(edge_table->EdgeNum(), 1); + auto out_csr = edge_table->TakeOutCsr(); + auto in_csr = edge_table->TakeInCsr(); + ASSERT_NE(out_csr, nullptr); + ASSERT_NE(in_csr, nullptr); + EXPECT_EQ(out_csr->size(), vertex_capacity); + EXPECT_EQ(in_csr->size(), vertex_capacity); + }; -TEST_F(EdgeTableTest, BatchAddEdgesUsesIndexerVertexCapacity) { - auto ckp = make_checkpoint(workspace()); - constexpr neug::vid_t kVertexNum = 16; - constexpr neug::vid_t kVertexCapacity = 128; - InitIndexers(*ckp, kVertexNum, kVertexNum); - src_indexer.reserve(kVertexCapacity); - dst_indexer.reserve(kVertexCapacity); - ConstructEdgeTable(src_label_, dst_label_, edge_label_int_); - OpenEdgeTableInMemory(ckp, neug::CheckpointManifest(), kVertexNum, - kVertexNum); - - auto chunks = convert_to_data_chunks( - {split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{0}, 1), - split_column_to_chunks(std::vector{42}, 1)}); - BatchInsert(std::move(chunks)); - - auto out_csr = edge_table->TakeOutCsr(); - auto in_csr = edge_table->TakeInCsr(); - ASSERT_NE(out_csr, nullptr); - ASSERT_NE(in_csr, nullptr); - EXPECT_EQ(out_csr->size(), kVertexCapacity); - EXPECT_EQ(in_csr->size(), kVertexCapacity); + constexpr neug::vid_t kBulkVertexNum = 4097; + run(true, kBulkVertexNum, kBulkVertexNum + kBulkVertexNum / 4); + run(false, 16, 128); } TEST_F(EdgeTableTest, TestBatchAddEdgesUnbundled) { diff --git a/tests/storage/test_mutable_csr.cc b/tests/storage/test_mutable_csr.cc index c5a3b456b..8ab438da3 100644 --- a/tests/storage/test_mutable_csr.cc +++ b/tests/storage/test_mutable_csr.cc @@ -15,7 +15,6 @@ #include #include -#include #include #include #include @@ -552,67 +551,6 @@ class MutableCsrTest : public ::testing::Test { }; TYPED_TEST_SUITE(MutableCsrTest, Datatypes); -TEST(MutableCsrBulkBuildAccessTest, ReservesTwentyPercentPerVertex) { - auto test_dir = make_unique_test_dir("mutable_csr_bulk_build"); - CheckpointManager workspace; - workspace.Open(test_dir.string()); - auto ckp = make_checkpoint(workspace); - MutableCsr csr; - csr.Open(*ckp, ModuleDescriptor(), MemoryLevel::kInMemory); - - MutableCsrBulkBuildAccess writer(csr); - writer.PrepareBuild(3); - std::vector counters; - for (int thread = 0; thread < 4; ++thread) { - counters.emplace_back([&writer] { writer.CountConcurrent(0, 5); }); - } - for (auto& counter : counters) { - counter.join(); - } - writer.CountConcurrent(1, 2); - writer.AllocateFromCounts(); - - std::vector fillers; - for (int thread = 0; thread < 4; ++thread) { - fillers.emplace_back([&writer] { - const auto begin = writer.ReserveConcurrent(0, 5); - for (int slot = begin; slot < begin + 5; ++slot) { - writer.PutAt(0, slot, static_cast(slot), slot, 0); - } - }); - } - for (auto& filler : fillers) { - filler.join(); - } - for (int i = 0; i < 2; ++i) { - writer.PutSerial(1, static_cast(i), i, 0); - } - writer.Finish(); - EXPECT_EQ(csr.edge_num(), 22); - auto built_edges = csr.get_generic_view(0).get_edges(0); - std::vector built_neighbors; - for (auto it = built_edges.begin(); it != built_edges.end(); ++it) { - built_neighbors.push_back(*it); - } - ASSERT_EQ(built_neighbors.size(), 20); - std::sort(built_neighbors.begin(), built_neighbors.end()); - for (int i = 0; i < 20; ++i) { - EXPECT_EQ(built_neighbors[static_cast(i)], static_cast(i)); - } - - auto before = csr.get_generic_view(0).get_edges(0).start_ptr; - Allocator allocator(MemoryLevel::kInMemory, ""); - for (int i = 20; i < 24; ++i) { - csr.put_edge(0, static_cast(i), i, 0, allocator); - } - auto after = csr.get_generic_view(0).get_edges(0).start_ptr; - EXPECT_EQ(after, before); - EXPECT_EQ(csr.edge_num(), 26); - - workspace.Close(); - std::filesystem::remove_all(test_dir); -} - TYPED_TEST(MutableCsrTest, TestCsrType) { MutableCsr mutable_csr; EXPECT_EQ(mutable_csr.csr_type(), CsrType::kMutable); diff --git a/tests/utils/test_reader.cc b/tests/utils/test_reader.cc index 8e43a3794..73d142815 100644 --- a/tests/utils/test_reader.cc +++ b/tests/utils/test_reader.cc @@ -16,6 +16,7 @@ #include "test_reader.h" #include +#include #include #include "neug/storages/loader/chunk_pipeline_utils.h" @@ -37,6 +38,33 @@ ChunkSourceOptions parallel_source_options( }; } +CsvReadConfig csv_config( + std::initializer_list> columns, + size_t skip_rows = 0) { + CsvReadConfig config; + config.delimiter = '|'; + config.skip_rows = skip_rows; + for (const auto& [name, type] : columns) { + config.column_names.push_back(name); + config.column_types.emplace(name, DataType(type)); + } + config.include_columns = config.column_names; + return config; +} + +template +std::vector read_sorted_column( + const std::shared_ptr& supplier, size_t column = 0) { + std::vector values; + while (auto chunk = supplier->GetNextChunk()) { + for (size_t row = 0; row < chunk->row_num(); ++row) { + values.push_back(chunk->get(column)->get_elem(row).GetValue()); + } + } + std::sort(values.begin(), values.end()); + return values; +} + } // namespace // Test 1: Basic CSV reading with default options @@ -66,204 +94,192 @@ TEST_F(ReaderTest, TestBasicCsvRead) { EXPECT_EQ(ctx.row_num(), 3); } -TEST_F(ReaderTest, CsvChunkSourceCanBeReopened) { - createCsvFile("repeatable.csv", "id|name\n1|Alice\n2|Bob\n3|Carol\n"); - std::vector column_names = {"id", "name"}; - std::vector> column_types = { - createInt32Type(), createStringType()}; - auto shared_state = - createSharedState("repeatable.csv", column_names, column_types, - {{"skip_rows", "1"}, {"batch_read", "true"}}); - auto reader = createCsvReader(shared_state); - auto source = reader->createChunkSource(); - - ASSERT_NE(source, nullptr); - for (int pass = 0; pass < 2; ++pass) { - auto supplier = source->Open(); - ASSERT_NE(supplier, nullptr); - size_t rows = 0; - while (auto chunk = supplier->GetNextChunk()) { - EXPECT_EQ(chunk->col_num(), 2); - rows += chunk->row_num(); - } - EXPECT_EQ(rows, 3); - } +TEST_F(ReaderTest, CsvChunkSourceReusesPartitionPlanAcrossParallelOpens) { + createCsvFile("cached-plan.csv", "id|name\n1|Alice\n2|Bob\n3|Carol\n"); + auto config = csv_config( + {{"id", DataTypeId::kInt32}, {"name", DataTypeId::kVarchar}}, 1); + const auto csv_path = + std::filesystem::path(ARROW_READER_TEST_DIR) / "cached-plan.csv"; + const auto moved_path = + std::filesystem::path(ARROW_READER_TEST_DIR) / "cached-plan.moved"; + CSVChunkSource source({csv_path.string()}, std::move(config)); + const auto options = parallel_source_options(4, 4); + + auto first = source.Open(options); + ASSERT_NE(first, nullptr); + EXPECT_EQ(first->RowNum(), 4); + + std::filesystem::rename(csv_path, moved_path); + auto second = source.Open(options); + ASSERT_NE(second, nullptr); + EXPECT_EQ(second->RowNum(), 4); + std::filesystem::rename(moved_path, csv_path); } -TEST_F(ReaderTest, CsvChunkSourcePartitionsQuotedRecords) { - createCsvFile( - "partitioned.csv", - "id|name\r\n\r\n1|\"Alice|Smith\"\r\n2|\"line one\nline two\"\r\n" - "3|Carol\r\n4|\"D\"\"Angelo\"\r\n5|Last"); - std::vector column_names = {"id", "name"}; - std::vector> column_types = { - createInt32Type(), createStringType()}; - auto shared_state = createSharedState( - "partitioned.csv", column_names, column_types, - {{"skip_rows", "1"}, {"batch_read", "true"}, {"batch_size", "1"}}); - auto reader = createCsvReader(shared_state); - auto source = reader->createChunkSource(); - ASSERT_NE(source, nullptr); - - auto read_rows = [](const std::shared_ptr& supplier) { - std::vector> rows; - while (auto chunk = supplier->GetNextChunk()) { - EXPECT_EQ(chunk->col_num(), 2); - for (size_t row = 0; row < chunk->row_num(); ++row) { - rows.emplace_back(chunk->get(0)->get_elem(row).GetValue(), - chunk->get(1)->get_elem(row).GetValue()); +TEST_F(ReaderTest, PartitionedCsvHandlesQuotedRecordBoundaries) { + { + SCOPED_TRACE("quoted records"); + createCsvFile( + "partitioned.csv", + "id|name\r\n\r\n1|\"Alice|Smith\"\r\n2|\"line one\nline two\"\r\n" + "3|Carol\r\n4|\"D\"\"Angelo\"\r\n5|Last"); + std::vector column_names = {"id", "name"}; + std::vector> column_types = { + createInt32Type(), createStringType()}; + auto shared_state = createSharedState( + "partitioned.csv", column_names, column_types, + {{"skip_rows", "1"}, {"batch_read", "true"}, {"batch_size", "1"}}); + auto reader = createCsvReader(shared_state); + auto source = reader->createChunkSource(); + ASSERT_NE(source, nullptr); + + auto read_rows = [](const std::shared_ptr& supplier) { + std::vector> rows; + while (auto chunk = supplier->GetNextChunk()) { + EXPECT_EQ(chunk->col_num(), 2); + for (size_t row = 0; row < chunk->row_num(); ++row) { + rows.emplace_back( + chunk->get(0)->get_elem(row).GetValue(), + chunk->get(1)->get_elem(row).GetValue()); + } } - } - std::sort(rows.begin(), rows.end()); - return rows; - }; - - auto expected = read_rows(source->Open()); - auto partitioned = source->Open(parallel_source_options()); - ASSERT_NE(partitioned, nullptr); - EXPECT_TRUE(partitioned->SupportsConcurrentGetNext()); - EXPECT_EQ(partitioned->RowNum(), 6); // Header is an intentional overcount. - auto actual = read_rows(partitioned); - - ASSERT_EQ(actual, expected); - ASSERT_EQ(actual.size(), 5); - EXPECT_EQ(actual[0], std::make_pair(1, std::string("Alice|Smith"))); - EXPECT_EQ(actual[1], std::make_pair(2, std::string("line one\nline two"))); - EXPECT_EQ(actual[3], std::make_pair(4, std::string("D\"Angelo"))); - EXPECT_EQ(actual[4], std::make_pair(5, std::string("Last"))); -} - -TEST_F(ReaderTest, PartitionedCsvCarriesSkipAcrossRanges) { - std::string long_quoted_field(64 * 1024, 'x'); - for (size_t i = 32; i < long_quoted_field.size(); i += 64) { - long_quoted_field[i] = '\n'; + std::sort(rows.begin(), rows.end()); + return rows; + }; + + auto expected = read_rows(source->Open()); + auto partitioned = source->Open(parallel_source_options()); + ASSERT_NE(partitioned, nullptr); + EXPECT_TRUE(partitioned->SupportsConcurrentGetNext()); + EXPECT_EQ(partitioned->RowNum(), 6); // Header is an intentional overcount. + auto actual = read_rows(partitioned); + + ASSERT_EQ(actual, expected); + ASSERT_EQ(actual.size(), 5); + EXPECT_EQ(actual[0], std::make_pair(1, std::string("Alice|Smith"))); + EXPECT_EQ(actual[1], std::make_pair(2, std::string("line one\nline two"))); + EXPECT_EQ(actual[3], std::make_pair(4, std::string("D\"Angelo"))); + EXPECT_EQ(actual[4], std::make_pair(5, std::string("Last"))); } - createCsvFile("partition-skip.csv", - "0|\"" + long_quoted_field + - "\"\r\n1|also-skipped\r\n2|kept\r\n3|last\r\n"); - CsvReadConfig config; - config.delimiter = '|'; - config.quoting = true; - config.skip_rows = 2; - config.chunk_size = 1; - config.column_names = {"id", "name"}; - config.include_columns = config.column_names; - config.column_types.emplace("id", DataType(DataTypeId::kInt32)); - config.column_types.emplace("name", DataType(DataTypeId::kVarchar)); - CSVChunkSource source( - {std::string(ARROW_READER_TEST_DIR) + "/partition-skip.csv"}, config); + { + SCOPED_TRACE("skip rows across ranges"); + std::string long_quoted_field(64 * 1024, 'x'); + for (size_t i = 32; i < long_quoted_field.size(); i += 64) { + long_quoted_field[i] = '\n'; + } + createCsvFile("partition-skip.csv", + "0|\"" + long_quoted_field + + "\"\r\n1|also-skipped\r\n2|kept\r\n3|last\r\n"); + + auto config = csv_config( + {{"id", DataTypeId::kInt32}, {"name", DataTypeId::kVarchar}}, 2); + config.quoting = true; + config.chunk_size = 1; + CSVChunkSource source( + {std::string(ARROW_READER_TEST_DIR) + "/partition-skip.csv"}, config); + + auto supplier = source.Open(parallel_source_options(4, 4)); + ASSERT_NE(supplier, nullptr); - auto supplier = source.Open(parallel_source_options(4, 4)); - ASSERT_NE(supplier, nullptr); + EXPECT_EQ(read_sorted_column(supplier), + (std::vector{2, 3})); + } - std::vector ids; - while (auto chunk = supplier->GetNextChunk()) { - for (size_t row = 0; row < chunk->row_num(); ++row) { - ids.push_back(chunk->get(0)->get_elem(row).GetValue()); - } + { + SCOPED_TRACE("double quote disabled"); + std::string multiline = "first \"\"quoted\n"; + multiline.append(64 * 1024, 'z'); + multiline += "\ncontinued\"\" tail"; + createCsvFile("partition-double-quote.csv", + "0|\"" + multiline + "\"\n1|one\n2|two\n"); + + auto config = csv_config( + {{"id", DataTypeId::kInt32}, {"name", DataTypeId::kVarchar}}); + config.quoting = true; + config.double_quote = false; + config.chunk_size = 1; + CSVChunkSource source( + {std::string(ARROW_READER_TEST_DIR) + "/partition-double-quote.csv"}, + config); + + const auto expected = read_sorted_column(source.Open()); + + EXPECT_EQ( + read_sorted_column(source.Open(parallel_source_options(4, 4))), + expected); + EXPECT_EQ(expected, (std::vector{0, 1, 2})); } - std::sort(ids.begin(), ids.end()); - EXPECT_EQ(ids, (std::vector{2, 3})); } -TEST_F(ReaderTest, PartitionPlannerMatchesParserWhenDoubleQuoteIsFalse) { - std::string multiline = "first \"\"quoted\n"; - multiline.append(64 * 1024, 'z'); - multiline += "\ncontinued\"\" tail"; - createCsvFile("partition-double-quote.csv", - "0|\"" + multiline + "\"\n1|one\n2|two\n"); - - CsvReadConfig config; - config.delimiter = '|'; - config.quoting = true; - config.double_quote = false; - config.chunk_size = 1; - config.column_names = {"id", "name"}; - config.include_columns = config.column_names; - config.column_types.emplace("id", DataType(DataTypeId::kInt32)); - config.column_types.emplace("name", DataType(DataTypeId::kVarchar)); - CSVChunkSource source( - {std::string(ARROW_READER_TEST_DIR) + "/partition-double-quote.csv"}, - config); +TEST_F(ReaderTest, CsvChunkSourceHonorsProjectionAndParallelOptions) { + { + SCOPED_TRACE("partitioned projection"); + createCsvFile("partition-projection.csv", + "id|ignored|score\n1|Alice|10\n2|Bob|20\n"); + auto config = csv_config({{"id", DataTypeId::kInt32}, + {"ignored", DataTypeId::kVarchar}, + {"score", DataTypeId::kInt32}}, + 1); + config.chunk_size = 1; + CSVChunkSource source( + {std::string(ARROW_READER_TEST_DIR) + "/partition-projection.csv"}, + config); + auto supplier = source.Open(parallel_source_options(2, 2, {2, 0})); + ASSERT_NE(supplier, nullptr); - auto read_ids = [](std::shared_ptr supplier) { - std::vector ids; + std::vector> rows; while (auto chunk = supplier->GetNextChunk()) { + ASSERT_EQ(chunk->col_num(), 2); for (size_t row = 0; row < chunk->row_num(); ++row) { - ids.push_back(chunk->get(0)->get_elem(row).GetValue()); + rows.emplace_back(chunk->get(0)->get_elem(row).GetValue(), + chunk->get(1)->get_elem(row).GetValue()); } } - std::sort(ids.begin(), ids.end()); - return ids; - }; - const auto expected = read_ids(source.Open()); - - EXPECT_EQ(read_ids(source.Open(parallel_source_options(4, 4))), expected); - EXPECT_EQ(expected, (std::vector{0, 1, 2})); -} - -TEST_F(ReaderTest, CsvChunkSourcePushesProjectionIntoPartitionedParsing) { - createCsvFile("partition-projection.csv", - "id|ignored|score\n1|Alice|10\n2|Bob|20\n"); - CsvReadConfig config; - config.delimiter = '|'; - config.skip_rows = 1; - config.chunk_size = 1; - config.column_names = {"id", "ignored", "score"}; - config.include_columns = config.column_names; - config.column_types.emplace("id", DataType(DataTypeId::kInt32)); - config.column_types.emplace("ignored", DataType(DataTypeId::kVarchar)); - config.column_types.emplace("score", DataType(DataTypeId::kInt32)); - CSVChunkSource source( - {std::string(ARROW_READER_TEST_DIR) + "/partition-projection.csv"}, - config); - - auto supplier = source.Open(parallel_source_options(2, 2, {2, 0})); - ASSERT_NE(supplier, nullptr); + std::sort(rows.begin(), rows.end()); + EXPECT_EQ(rows, + (std::vector>{{10, 1}, {20, 2}})); + } - std::vector> rows; - while (auto chunk = supplier->GetNextChunk()) { - ASSERT_EQ(chunk->col_num(), 2); - for (size_t row = 0; row < chunk->row_num(); ++row) { - rows.emplace_back(chunk->get(0)->get_elem(row).GetValue(), - chunk->get(1)->get_elem(row).GetValue()); - } + { + SCOPED_TRACE("parallel disabled"); + createCsvFile("serial.csv", "id|name\n1|Alice\n2|Bob\n"); + std::vector column_names = {"id", "name"}; + std::vector> column_types = { + createInt32Type(), createStringType()}; + auto shared_state = createSharedState( + "serial.csv", column_names, column_types, + {{"skip_rows", "1"}, {"batch_read", "true"}, {"parallel", "false"}}); + auto source = createCsvReader(shared_state)->createChunkSource(); + ASSERT_NE(source, nullptr); + EXPECT_FALSE(source->ParallelEnabled()); + + auto supplier = source->Open(parallel_source_options()); + ASSERT_NE(supplier, nullptr); + EXPECT_FALSE(supplier->SupportsConcurrentGetNext()); } - std::sort(rows.begin(), rows.end()); - EXPECT_EQ(rows, (std::vector>{{10, 1}, {20, 2}})); -} -TEST_F(ReaderTest, CsvChunkSourceHonorsParallelFalse) { - createCsvFile("serial.csv", "id|name\n1|Alice\n2|Bob\n"); - std::vector column_names = {"id", "name"}; - std::vector> column_types = { - createInt32Type(), createStringType()}; - auto shared_state = createSharedState( - "serial.csv", column_names, column_types, - {{"skip_rows", "1"}, {"batch_read", "true"}, {"parallel", "false"}}); - auto source = createCsvReader(shared_state)->createChunkSource(); - ASSERT_NE(source, nullptr); - EXPECT_FALSE(source->ParallelEnabled()); - - auto supplier = source->Open(parallel_source_options()); - ASSERT_NE(supplier, nullptr); - EXPECT_FALSE(supplier->SupportsConcurrentGetNext()); + { + SCOPED_TRACE("post-read projection rejected"); + createCsvFile("projected.csv", "id|name\n1|Alice\n"); + std::vector column_names = {"id", "name"}; + std::vector> column_types = { + createInt32Type(), createStringType()}; + auto shared_state = + createSharedState("projected.csv", column_names, column_types, {}); + shared_state->projectColumns = {"id"}; + EXPECT_EQ(createCsvReader(shared_state)->createChunkSource(), nullptr); + } } TEST_F(ReaderTest, PartitionedCsvSkipsHeaderForEveryFile) { createCsvFile("part-a.csv", "id|name\n1|Alice\n2|Bob\n"); createCsvFile("part-b.csv", "id|name\n3|Carol\n4|Dave"); - CsvReadConfig config; - config.delimiter = '|'; + auto config = csv_config( + {{"id", DataTypeId::kInt32}, {"name", DataTypeId::kVarchar}}, 1); config.quoting = true; config.escaping = false; - config.skip_rows = 1; config.chunk_size = 1; - config.column_names = {"id", "name"}; - config.include_columns = config.column_names; - config.column_types.emplace("id", DataType(DataTypeId::kInt32)); - config.column_types.emplace("name", DataType(DataTypeId::kVarchar)); auto path = [](const char* file) { return std::string(ARROW_READER_TEST_DIR) + "/" + file; }; @@ -274,29 +290,18 @@ TEST_F(ReaderTest, PartitionedCsvSkipsHeaderForEveryFile) { EXPECT_EQ(supplier->RowNum(), 6); // Two headers are counted as reserve hints. - std::vector ids; - while (auto chunk = supplier->GetNextChunk()) { - for (size_t row = 0; row < chunk->row_num(); ++row) { - ids.push_back(chunk->get(0)->get_elem(row).GetValue()); - } - } - std::sort(ids.begin(), ids.end()); - EXPECT_EQ(ids, (std::vector{1, 2, 3, 4})); + EXPECT_EQ(read_sorted_column(supplier), + (std::vector{1, 2, 3, 4})); } TEST_F(ReaderTest, PartitionedCsvPropagatesProducerErrors) { createCsvFile("partition-error.csv", "id|name\n1|Alice\nnot-an-int|Broken\n3|Carol\n4|Dave"); - CsvReadConfig config; - config.delimiter = '|'; + auto config = csv_config( + {{"id", DataTypeId::kInt32}, {"name", DataTypeId::kVarchar}}, 1); config.quoting = true; config.escaping = false; - config.skip_rows = 1; config.chunk_size = 1; - config.column_names = {"id", "name"}; - config.include_columns = config.column_names; - config.column_types.emplace("id", DataType(DataTypeId::kInt32)); - config.column_types.emplace("name", DataType(DataTypeId::kVarchar)); CSVChunkSource source( {std::string(ARROW_READER_TEST_DIR) + "/partition-error.csv"}, config); @@ -307,19 +312,6 @@ TEST_F(ReaderTest, PartitionedCsvPropagatesProducerErrors) { }); } -TEST_F(ReaderTest, CsvChunkSourceRejectsPostReadProjection) { - createCsvFile("projected.csv", "id|name\n1|Alice\n"); - std::vector column_names = {"id", "name"}; - std::vector> column_types = { - createInt32Type(), createStringType()}; - auto shared_state = - createSharedState("projected.csv", column_names, column_types, {}); - shared_state->projectColumns = {"id"}; - - auto reader = createCsvReader(shared_state); - EXPECT_EQ(reader->createChunkSource(), nullptr); -} - // Test 2: CSV with different delimiter (tab) TEST_F(ReaderTest, TestCsvWithTabDelimiter) { createCsvFile("test2.csv", "id\tname\tage\n1\tAlice\t95.5\n2\tBob\t87.0\n"); diff --git a/tests/utils/test_reader.h b/tests/utils/test_reader.h index 0ef410c37..6d5432fbd 100644 --- a/tests/utils/test_reader.h +++ b/tests/utils/test_reader.h @@ -21,11 +21,6 @@ #include #include -#include -#include -#include -#include - #include "neug/common/types/data_chunk.h" #include "neug/compiler/common/case_insensitive_map.h" #include "neug/execution/common/context.h" From fa98077a48d57dfab2fb7ff417228c035b1126f4 Mon Sep 17 00:00:00 2001 From: "xiaolei.zl" Date: Mon, 20 Jul 2026 10:23:48 +0800 Subject: [PATCH 4/8] fix vertex load issue --- include/neug/storages/graph/vertex_table.h | 8 ++- include/neug/storages/loader/loader_utils.h | 37 +++++++++++ src/storages/graph/edge_table.cc | 2 +- src/storages/graph/property_graph.cc | 2 +- src/storages/graph/vertex_table.cc | 50 ++++++++++++++- .../loader/bundled_edge_csr_loader.cc | 63 +++---------------- src/storages/loader/bundled_edge_csr_loader.h | 2 - src/storages/loader/loader_utils.cc | 38 +++++++++++ tests/storage/test_vertex_table.cc | 4 +- 9 files changed, 143 insertions(+), 63 deletions(-) diff --git a/include/neug/storages/graph/vertex_table.h b/include/neug/storages/graph/vertex_table.h index 1b3d111ef..d8f8d56d8 100644 --- a/include/neug/storages/graph/vertex_table.h +++ b/include/neug/storages/graph/vertex_table.h @@ -28,6 +28,7 @@ class ModuleBroker; class CheckpointManifest; class Checkpoint; class IDataChunkSupplier; +class IDataChunkSource; class VertexTableView; class VertexSet { @@ -257,7 +258,7 @@ class VertexTable { void Compact(timestamp_t ts = MAX_TIMESTAMP); - void insert_vertices(std::shared_ptr supplier); + void BatchAddVertices(std::shared_ptr supplier); const VertexTimestamp& get_vertex_timestamp() const { return *v_ts_; } @@ -265,6 +266,11 @@ class VertexTable { Table& get_table() { return *table_; } private: + bool try_batch_build_vertices( + const std::shared_ptr& source); + + void batch_add_vertices_impl(std::shared_ptr supplier); + vid_t insert_vertex_pk(const Value& id, timestamp_t ts, bool insert_safe); std::vector insert_primary_keys( diff --git a/include/neug/storages/loader/loader_utils.h b/include/neug/storages/loader/loader_utils.h index 68f1a9d13..4436d10e5 100644 --- a/include/neug/storages/loader/loader_utils.h +++ b/include/neug/storages/loader/loader_utils.h @@ -17,6 +17,10 @@ #include #include +#include +#include +#include +#include #include #include #include @@ -125,7 +129,40 @@ class IDataChunkSource { std::shared_ptr make_data_chunk_supplier( std::shared_ptr source); +enum class BulkBuildWorkerStrategy { + kMaxProducers, + kBalancedProducerConsumer, +}; + +ChunkSourceOptions ResolveBulkBuildSourceOptions( + int64_t source_bytes, bool parallel_enabled, bool preserve_order, + BulkBuildWorkerStrategy worker_strategy); + inline constexpr int64_t kUnknownRowNum = -1; +inline constexpr int64_t kDefaultBulkBuildMinBytes = 256LL * 1024 * 1024; + +inline bool ShouldUseBulkBuild(int64_t source_bytes, + int64_t min_bytes = kDefaultBulkBuildMinBytes) { + const char* configured = std::getenv("NEUG_COPY_BULK_BUILD"); + if (configured != nullptr) { + std::string value(configured); + std::transform(value.begin(), value.end(), value.begin(), + [](unsigned char ch) { return std::tolower(ch); }); + if (value == "0" || value == "false" || value == "off" || value == "no") { + return false; + } + if (value == "1" || value == "true" || value == "on" || value == "yes") { + return true; + } + LOG(WARNING) << "Ignore invalid NEUG_COPY_BULK_BUILD=" << configured; + } + return source_bytes >= min_bytes; +} + +inline bool ShouldUseBulkBuild(const IDataChunkSource& source, + int64_t min_bytes = kDefaultBulkBuildMinBytes) { + return ShouldUseBulkBuild(source.EstimatedBytes(), min_bytes); +} enum class CsvRowCountMode { kCountOnOpen, diff --git a/src/storages/graph/edge_table.cc b/src/storages/graph/edge_table.cc index c1e4a4acb..36c0e4f88 100644 --- a/src/storages/graph/edge_table.cc +++ b/src/storages/graph/edge_table.cc @@ -868,7 +868,7 @@ bool EdgeTable::TryBatchBuildEdges( return false; } const auto source_bytes = source->EstimatedBytes(); - if (!internal::BundledEdgeCsrLoader::ShouldBuild(source_bytes)) { + if (!ShouldUseBulkBuild(source_bytes)) { return false; } diff --git a/src/storages/graph/property_graph.cc b/src/storages/graph/property_graph.cc index 1eb225923..13a0220a5 100644 --- a/src/storages/graph/property_graph.cc +++ b/src/storages/graph/property_graph.cc @@ -137,7 +137,7 @@ Status PropertyGraph::EnsureCapacity(label_t src_label, label_t dst_label, Status PropertyGraph::BatchAddVertices( label_t v_label, std::shared_ptr supplier) { RETURN_IF_NOT_OK(vertex_label_check(v_label)); - vertex_tables_[v_label].insert_vertices(std::move(supplier)); + vertex_tables_[v_label].BatchAddVertices(std::move(supplier)); return neug::Status::OK(); } diff --git a/src/storages/graph/vertex_table.cc b/src/storages/graph/vertex_table.cc index db8518b43..8c9ae7200 100644 --- a/src/storages/graph/vertex_table.cc +++ b/src/storages/graph/vertex_table.cc @@ -15,6 +15,8 @@ #include "neug/storages/graph/vertex_table.h" +#include + #include "neug/storages/checkpoint_manifest.h" #include "neug/storages/loader/loader_utils.h" #include "neug/storages/module/module_broker.h" @@ -46,7 +48,53 @@ void VertexTable::Init(std::shared_ptr ckp, MemoryLevel level) { v_ts_->Open(*ckp_, ModuleDescriptor{}, level); } -void VertexTable::insert_vertices( +void VertexTable::BatchAddVertices( + std::shared_ptr supplier) { + CHECK(supplier != nullptr); + auto source = supplier->RepeatableSource(); + if (source && ShouldUseBulkBuild(*source) && + try_batch_build_vertices(source)) { + return; + } + batch_add_vertices_impl(std::move(supplier)); +} + +bool VertexTable::try_batch_build_vertices( + const std::shared_ptr& source) { + if (!source || Size() != 0) { + return false; + } + + auto staged = VertexTable(vertex_schema_); + staged.Init(ckp_, memory_level_); + const auto source_bytes = source->EstimatedBytes(); + auto options = ResolveBulkBuildSourceOptions( + source_bytes, source->ParallelEnabled(), false, + BulkBuildWorkerStrategy::kMaxProducers); + auto supplier = source->Open(options); + if (!supplier) { + return false; + } + + auto row_num = supplier->RowNum(); + if (row_num >= 0) { + if (static_cast(row_num) > std::numeric_limits::max()) { + THROW_RUNTIME_ERROR("Vertex row count exceeds addressable capacity"); + } + const auto row_count = static_cast(row_num); + const size_t checkpoint_headroom = row_count / 4; + if (checkpoint_headroom > std::numeric_limits::max() - row_count) { + THROW_RUNTIME_ERROR("Vertex bulk load capacity overflow"); + } + staged.EnsureCapacity(row_count + checkpoint_headroom); + } + + staged.batch_add_vertices_impl(std::move(supplier)); + Swap(staged); + return true; +} + +void VertexTable::batch_add_vertices_impl( std::shared_ptr supplier) { CHECK(supplier != nullptr); auto reserve_checkpoint_headroom = [this](size_t required_size) { diff --git a/src/storages/loader/bundled_edge_csr_loader.cc b/src/storages/loader/bundled_edge_csr_loader.cc index 4e7a0856a..8bfc70c0b 100644 --- a/src/storages/loader/bundled_edge_csr_loader.cc +++ b/src/storages/loader/bundled_edge_csr_loader.cc @@ -19,9 +19,7 @@ #include #include -#include #include -#include #include #include #include @@ -263,36 +261,6 @@ class BundledEdgeCsrLoader::SingleMutableWriter { namespace { -constexpr int64_t kMinBulkEdgeBuildBytes = 256LL * 1024 * 1024; - -ChunkSourceOptions resolve_bulk_edge_source_options(int64_t source_bytes, - bool parallel_enabled, - bool preserve_order) { - constexpr int64_t kMinPartitionBytes = 64LL * 1024 * 1024; - constexpr size_t kMaxQueuedChunks = 64; - - ChunkSourceOptions options; - options.preserve_order = preserve_order; - const auto workers = chunk_pipeline_detail::hardware_worker_count(); - if (!parallel_enabled || preserve_order || workers <= 1 || - source_bytes < kMinBulkEdgeBuildBytes) { - return options; - } - - const auto useful_partitions = std::max( - 1, source_bytes / kMinPartitionBytes + - (source_bytes % kMinPartitionBytes == 0 ? 0 : 1)); - const auto balanced_producers = (workers + 1) / 2; - options.producer_count = static_cast(std::min( - balanced_producers, std::min(useful_partitions, workers - 1))); - options.producer_count = std::max(1, options.producer_count); - options.consumer_count = - std::max(1, workers - options.producer_count); - options.queue_capacity = std::clamp( - static_cast(options.producer_count) * 2, 2, kMaxQueuedChunks); - return options; -} - class EmptyCsrBulkWriter { public: static constexpr EdgeStrategy kStrategy = EdgeStrategy::kNone; @@ -787,8 +755,9 @@ uint64_t build_bundled_edges_with_ops( ops.allocate_from_counts(); return 0; } - const auto count_options = resolve_bulk_edge_source_options( - source_bytes, source->ParallelEnabled(), false); + const auto count_options = ResolveBulkBuildSourceOptions( + source_bytes, source->ParallelEnabled(), false, + BulkBuildWorkerStrategy::kBalancedProducerConsumer); const bool needs_degree_count = ops.needs_degree_count; std::vector scratches; if (needs_degree_count) { @@ -809,9 +778,11 @@ uint64_t build_bundled_edges_with_ops( has_single_direction && (!needs_degree_count || single_duplicate); const bool allow_concurrent_fill = !preserve_fill_order; const auto fill_options = - preserve_fill_order ? resolve_bulk_edge_source_options( - source_bytes, source->ParallelEnabled(), true) - : count_options; + preserve_fill_order + ? ResolveBulkBuildSourceOptions( + source_bytes, source->ParallelEnabled(), true, + BulkBuildWorkerStrategy::kBalancedProducerConsumer) + : count_options; if (scratches.size() < static_cast(fill_options.consumer_count)) { scratches.resize(static_cast(fill_options.consumer_count)); } @@ -875,24 +846,6 @@ bool build_bundled_edges(CsrBase* out_csr, CsrBase* in_csr, } // namespace -bool internal::BundledEdgeCsrLoader::ShouldBuild(int64_t source_bytes) { - const char* configured = std::getenv("NEUG_COPY_BULK_BUILD"); - if (configured != nullptr) { - std::string value(configured); - std::transform(value.begin(), value.end(), value.begin(), - [](unsigned char ch) { return std::tolower(ch); }); - if (value == "0" || value == "false" || value == "off" || value == "no") { - return false; - } - if (value == "1" || value == "true" || value == "on" || value == "yes") { - return true; - } - LOG(WARNING) << "Ignore invalid NEUG_COPY_BULK_BUILD=" << configured; - } - - return source_bytes >= kMinBulkEdgeBuildBytes; -} - bool internal::BundledEdgeCsrLoader::TryBuild( CsrBase& out_csr, CsrBase& in_csr, const EdgeSchema& schema, const IndexerType& src_indexer, const IndexerType& dst_indexer, diff --git a/src/storages/loader/bundled_edge_csr_loader.h b/src/storages/loader/bundled_edge_csr_loader.h index ac382d413..495590a29 100644 --- a/src/storages/loader/bundled_edge_csr_loader.h +++ b/src/storages/loader/bundled_edge_csr_loader.h @@ -39,8 +39,6 @@ class BundledEdgeCsrLoader { template class SingleMutableWriter; - static bool ShouldBuild(int64_t source_bytes); - /// Returns false when either CSR layout or the property type is unsupported. /// Callers must pass fresh, unpublished CSR instances. static bool TryBuild(CsrBase& out_csr, CsrBase& in_csr, diff --git a/src/storages/loader/loader_utils.cc b/src/storages/loader/loader_utils.cc index decb17050..da6dcb87b 100644 --- a/src/storages/loader/loader_utils.cc +++ b/src/storages/loader/loader_utils.cc @@ -1880,6 +1880,44 @@ std::shared_ptr make_data_chunk_supplier( return std::make_shared(std::move(source)); } +ChunkSourceOptions ResolveBulkBuildSourceOptions( + int64_t source_bytes, bool parallel_enabled, bool preserve_order, + BulkBuildWorkerStrategy worker_strategy) { + constexpr int64_t kMinPartitionBytes = 64LL * 1024 * 1024; + constexpr size_t kMaxQueuedChunks = 64; + + ChunkSourceOptions options; + options.preserve_order = preserve_order; + const auto workers = chunk_pipeline_detail::hardware_worker_count(); + if (!parallel_enabled || preserve_order || workers <= 1 || + source_bytes < kDefaultBulkBuildMinBytes) { + return options; + } + + const auto useful_partitions = std::max( + 1, source_bytes / kMinPartitionBytes + + (source_bytes % kMinPartitionBytes == 0 ? 0 : 1)); + switch (worker_strategy) { + case BulkBuildWorkerStrategy::kMaxProducers: { + options.producer_count = static_cast(std::min( + workers - 1, std::min(useful_partitions, workers))); + break; + } + case BulkBuildWorkerStrategy::kBalancedProducerConsumer: { + const auto balanced_producers = (workers + 1) / 2; + options.producer_count = static_cast(std::min( + balanced_producers, std::min(useful_partitions, workers - 1))); + options.consumer_count = + std::max(1, workers - options.producer_count); + break; + } + } + options.producer_count = std::max(1, options.producer_count); + options.queue_capacity = std::clamp( + static_cast(options.producer_count) * 2, 2, kMaxQueuedChunks); + return options; +} + CSVChunkSource::CSVChunkSource(std::vector file_paths, CsvReadConfig config, std::vector projected_columns) diff --git a/tests/storage/test_vertex_table.cc b/tests/storage/test_vertex_table.cc index ae0664aec..3d60870f7 100644 --- a/tests/storage/test_vertex_table.cc +++ b/tests/storage/test_vertex_table.cc @@ -687,7 +687,7 @@ TEST_F(VertexTableTest, VertexTableResizeTest) { auto data_chunks = generate_data_chunks(10000); std::shared_ptr batch_supplier = std::make_shared(std::move(data_chunks)); - table.insert_vertices(batch_supplier); + table.BatchAddVertices(batch_supplier); EXPECT_EQ(table.VertexNum(), 10000); EXPECT_EQ(table.LidNum(), 10000); @@ -715,7 +715,7 @@ TEST_F(VertexTableTest, InsertVerticesFromRepeatableSource) { constexpr size_t kVertexNum = 4097; auto source = std::make_shared( generate_data_chunks(kVertexNum), kForceBulkBuildBytes); - table.insert_vertices(make_data_chunk_supplier(source)); + table.BatchAddVertices(make_data_chunk_supplier(source)); EXPECT_EQ(table.VertexNum(), kVertexNum); EXPECT_EQ(table.LidNum(), kVertexNum); From 9fbec25354ef0dd610ae2dbd60964d469104309a Mon Sep 17 00:00:00 2001 From: "xiaolei.zl" Date: Mon, 20 Jul 2026 13:58:40 +0800 Subject: [PATCH 5/8] refactor --- .../function/import/csv_read_function.h | 61 ++-- .../execute/ops/batch/batch_insert_edge.h | 3 - .../execute/ops/batch/batch_insert_vertex.h | 3 - include/neug/storages/loader/loader_utils.h | 7 +- .../execute/ops/batch/batch_insert_edge.cc | 18 +- .../execute/ops/batch/batch_insert_vertex.cc | 8 +- src/storages/graph/edge_table.cc | 17 +- src/storages/graph/vertex_table.cc | 10 +- .../loader/bundled_edge_csr_loader.cc | 279 ++++++------------ src/storages/loader/loader_utils.cc | 114 +++---- 10 files changed, 181 insertions(+), 339 deletions(-) diff --git a/include/neug/compiler/function/import/csv_read_function.h b/include/neug/compiler/function/import/csv_read_function.h index caaeee5fb..f1a936443 100644 --- a/include/neug/compiler/function/import/csv_read_function.h +++ b/include/neug/compiler/function/import/csv_read_function.h @@ -112,21 +112,31 @@ struct CSVReadFunction { } } - static execution::Context execFunc( - std::shared_ptr state) { - validateAndConvertExecOptions(state); - const auto& vfs = neug::main::MetadataRegistry::getVFS(); - const auto& fs = vfs->Provide(state->schema.file); - auto resolvedPaths = std::vector(); + static void resolvePaths( + const std::shared_ptr& state) { + const auto& fs = + neug::main::MetadataRegistry::getVFS()->Provide(state->schema.file); + std::vector resolvedPaths; for (const auto& path : state->schema.file.paths) { const auto& resolved = fs->glob(path); resolvedPaths.insert(resolvedPaths.end(), resolved.begin(), resolved.end()); } state->schema.file.paths = std::move(resolvedPaths); + } + + static std::shared_ptr createReader( + const std::shared_ptr& state) { auto optionsBuilder = std::make_unique(state); - auto reader = - std::make_unique(state, std::move(optionsBuilder)); + return std::make_shared(state, + std::move(optionsBuilder)); + } + + static execution::Context execFunc( + std::shared_ptr state) { + validateAndConvertExecOptions(state); + resolvePaths(state); + auto reader = createReader(state); execution::Context ctx; auto localState = std::make_shared(); reader->read(localState, ctx); @@ -144,20 +154,9 @@ struct CSVReadFunction { // the original state pristine for that fallback. auto source_state = std::make_shared(*state); validateAndConvertExecOptions(source_state); - const auto& vfs = neug::main::MetadataRegistry::getVFS(); - const auto& fs = vfs->Provide(source_state->schema.file); - std::vector resolved_paths; - for (const auto& path : source_state->schema.file.paths) { - const auto& resolved = fs->glob(path); - resolved_paths.insert(resolved_paths.end(), resolved.begin(), - resolved.end()); - } - source_state->schema.file.paths = std::move(resolved_paths); - auto options_builder = - std::make_unique(source_state); - auto reader = std::make_unique( - source_state, std::move(options_builder)); - return reader->createChunkSource(std::move(projected_columns)); + resolvePaths(source_state); + return createReader(source_state) + ->createChunkSource(std::move(projected_columns)); } static std::shared_ptr sniffFunc( @@ -171,18 +170,8 @@ struct CSVReadFunction { validateAndConvertSniffOptions(externalSchema.file); externalSchema.file.options["BATCH_SIZE"] = std::to_string(reader::kSniffBlockSize); - const auto& vfs = neug::main::MetadataRegistry::getVFS(); - const auto& fs = vfs->Provide(state->schema.file); - auto resolvedPaths = std::vector(); - for (const auto& path : state->schema.file.paths) { - const auto& resolved = fs->glob(path); - resolvedPaths.insert(resolvedPaths.end(), resolved.begin(), - resolved.end()); - } - state->schema.file.paths = std::move(resolvedPaths); - auto optionsBuilder = std::make_unique(state); - auto reader = - std::make_shared(state, std::move(optionsBuilder)); + resolvePaths(state); + auto reader = createReader(state); auto sniffer = std::make_shared(reader); auto sniffResult = sniffer->sniff(); if (sniffResult) { @@ -198,9 +187,7 @@ struct CSVReadFunction { if (hasHeader) { options.insert({"SKIP_ROWS", "1"}); options.insert({"AUTOGENERATE_COLUMN_NAMES", "TRUE"}); - auto optionsBuilder2 = std::make_unique(state); - auto reader2 = std::make_shared( - state, std::move(optionsBuilder2)); + auto reader2 = createReader(state); auto sniffer2 = std::make_shared(reader2); sniffResult = sniffer2->sniff(); if (sniffResult) { diff --git a/include/neug/execution/execute/ops/batch/batch_insert_edge.h b/include/neug/execution/execute/ops/batch/batch_insert_edge.h index 080e0c1ab..d0630f3d2 100644 --- a/include/neug/execution/execute/ops/batch/batch_insert_edge.h +++ b/include/neug/execution/execute/ops/batch/batch_insert_edge.h @@ -44,9 +44,6 @@ class BatchInsertEdgeOprBuilder : public IOperatorBuilder { /// build or normal BatchAdd from the supplied repeatable source. class BatchInsertEdgeFromSourceOprBuilder : public IOperatorBuilder { public: - BatchInsertEdgeFromSourceOprBuilder() = default; - ~BatchInsertEdgeFromSourceOprBuilder() = default; - neug::result Build(const Schema& schema, const ContextMeta& ctx_meta, const physical::PhysicalPlan& plan, diff --git a/include/neug/execution/execute/ops/batch/batch_insert_vertex.h b/include/neug/execution/execute/ops/batch/batch_insert_vertex.h index a96e87a0a..379cef23f 100644 --- a/include/neug/execution/execute/ops/batch/batch_insert_vertex.h +++ b/include/neug/execution/execute/ops/batch/batch_insert_vertex.h @@ -43,9 +43,6 @@ class BatchInsertVertexOprBuilder : public IOperatorBuilder { /// build or normal BatchAdd from the supplied repeatable source. class BatchInsertVertexFromSourceOprBuilder : public IOperatorBuilder { public: - BatchInsertVertexFromSourceOprBuilder() = default; - ~BatchInsertVertexFromSourceOprBuilder() = default; - neug::result Build(const Schema& schema, const ContextMeta& ctx_meta, const physical::PhysicalPlan& plan, diff --git a/include/neug/storages/loader/loader_utils.h b/include/neug/storages/loader/loader_utils.h index 4436d10e5..2425f3cd4 100644 --- a/include/neug/storages/loader/loader_utils.h +++ b/include/neug/storages/loader/loader_utils.h @@ -135,7 +135,7 @@ enum class BulkBuildWorkerStrategy { }; ChunkSourceOptions ResolveBulkBuildSourceOptions( - int64_t source_bytes, bool parallel_enabled, bool preserve_order, + int64_t source_bytes, bool parallel_enabled, BulkBuildWorkerStrategy worker_strategy); inline constexpr int64_t kUnknownRowNum = -1; @@ -159,11 +159,6 @@ inline bool ShouldUseBulkBuild(int64_t source_bytes, return source_bytes >= min_bytes; } -inline bool ShouldUseBulkBuild(const IDataChunkSource& source, - int64_t min_bytes = kDefaultBulkBuildMinBytes) { - return ShouldUseBulkBuild(source.EstimatedBytes(), min_bytes); -} - enum class CsvRowCountMode { kCountOnOpen, kUnknown, diff --git a/src/execution/execute/ops/batch/batch_insert_edge.cc b/src/execution/execute/ops/batch/batch_insert_edge.cc index 5cfa347ed..ba0356586 100644 --- a/src/execution/execute/ops/batch/batch_insert_edge.cc +++ b/src/execution/execute/ops/batch/batch_insert_edge.cc @@ -63,10 +63,7 @@ bool resolve_edge_triplet(const Schema& schema, if (!resolve_vertex_label_id(schema, edge_type.src_type_name(), src_type)) { return false; } - if (!resolve_vertex_label_id(schema, edge_type.dst_type_name(), dst_type)) { - return false; - } - return true; + return resolve_vertex_label_id(schema, edge_type.dst_type_name(), dst_type); } } // namespace @@ -141,7 +138,6 @@ neug::result BatchInsertEdgeOprBuilder::Build( const Schema& schema, const ContextMeta& ctx_meta, const physical::PhysicalPlan& plan, int op_idx) { (void) schema; - ContextMeta ret_meta = ctx_meta; const auto& opr = plan.plan(op_idx).opr().load_edge(); if (!opr.has_edge_type()) { @@ -161,16 +157,15 @@ neug::result BatchInsertEdgeOprBuilder::Build( std::make_unique( std::move(edge_type), std::move(prop_mappings), std::move(src_vertex_bindings), std::move(dst_vertex_binds)), - ret_meta); + ctx_meta); } neug::result BatchInsertEdgeFromSourceOprBuilder::Build( const Schema& schema, const ContextMeta& ctx_meta, const physical::PhysicalPlan& plan, int op_idx) { (void) schema; - ContextMeta result_meta = ctx_meta; if (!is_terminal_batch_insert(plan, op_idx)) { - return std::make_pair(nullptr, result_meta); + return std::make_pair(nullptr, ctx_meta); } const auto& edge_pb = plan.plan(op_idx + 1).opr().load_edge(); if (!edge_pb.has_edge_type()) { @@ -180,9 +175,8 @@ neug::result BatchInsertEdgeFromSourceOprBuilder::Build( auto source = build_batch_insert_source(plan, op_idx); - std::vector> property_mappings; - std::vector> source_mappings; - std::vector> destination_mappings; + std::vector> property_mappings, + source_mappings, destination_mappings; parse_property_mappings(edge_pb.property_mappings(), property_mappings); parse_property_mappings(edge_pb.source_vertex_binding(), source_mappings); parse_property_mappings(edge_pb.destination_vertex_binding(), @@ -193,7 +187,7 @@ neug::result BatchInsertEdgeFromSourceOprBuilder::Build( std::move(edge_type), std::move(property_mappings), std::move(source_mappings), std::move(destination_mappings), std::move(source)), - result_meta); + ctx_meta); } } // namespace ops diff --git a/src/execution/execute/ops/batch/batch_insert_vertex.cc b/src/execution/execute/ops/batch/batch_insert_vertex.cc index d5024dd0c..4f8b1143b 100644 --- a/src/execution/execute/ops/batch/batch_insert_vertex.cc +++ b/src/execution/execute/ops/batch/batch_insert_vertex.cc @@ -83,7 +83,6 @@ neug::result BatchInsertVertexOprBuilder::Build( const Schema& schema, const ContextMeta& ctx_meta, const physical::PhysicalPlan& plan, int op_idx) { (void) schema; - ContextMeta ret_meta = ctx_meta; const auto& opr = plan.plan(op_idx).opr().load_vertex(); if (!opr.has_vertex_type()) { @@ -96,16 +95,15 @@ neug::result BatchInsertVertexOprBuilder::Build( vertex_type.CopyFrom(opr.vertex_type()); return std::make_pair(std::make_unique( std::move(vertex_type), std::move(prop_mappings)), - ret_meta); + ctx_meta); } neug::result BatchInsertVertexFromSourceOprBuilder::Build( const Schema& schema, const ContextMeta& ctx_meta, const physical::PhysicalPlan& plan, int op_idx) { (void) schema; - ContextMeta result_meta = ctx_meta; if (!is_terminal_batch_insert(plan, op_idx)) { - return std::make_pair(nullptr, result_meta); + return std::make_pair(nullptr, ctx_meta); } const auto& vertex_pb = plan.plan(op_idx + 1).opr().load_vertex(); if (!vertex_pb.has_vertex_type()) { @@ -122,7 +120,7 @@ neug::result BatchInsertVertexFromSourceOprBuilder::Build( return std::make_pair(std::make_unique( std::move(vertex_type), std::move(property_mappings), std::move(source)), - result_meta); + ctx_meta); } } // namespace ops diff --git a/src/storages/graph/edge_table.cc b/src/storages/graph/edge_table.cc index 36c0e4f88..999cada08 100644 --- a/src/storages/graph/edge_table.cc +++ b/src/storages/graph/edge_table.cc @@ -47,17 +47,6 @@ namespace neug { -namespace { - -vid_t indexer_vertex_capacity(const IndexerType& indexer) { - const size_t capacity = indexer.capacity(); - CHECK_LE(capacity, static_cast(std::numeric_limits::max())) - << "CSR vertex capacity exceeds the vertex id range"; - return static_cast(capacity); -} - -} // namespace - void filterInvalidEdges(std::vector& src_lid, std::vector& dst_lid, std::vector& valid_flags) { @@ -785,8 +774,8 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, const IndexerType& dst_indexer, std::shared_ptr supplier) { CHECK(supplier != nullptr); - const auto src_vertex_capacity = indexer_vertex_capacity(src_indexer); - const auto dst_vertex_capacity = indexer_vertex_capacity(dst_indexer); + const auto src_vertex_capacity = static_cast(src_indexer.capacity()); + const auto dst_vertex_capacity = static_cast(dst_indexer.capacity()); auto source = supplier->RepeatableSource(); if (source && TryBatchBuildEdges(src_indexer, dst_indexer, source, src_vertex_capacity, dst_vertex_capacity)) { @@ -861,7 +850,7 @@ bool EdgeTable::TryBatchBuildEdges( const IndexerType& src_indexer, const IndexerType& dst_indexer, const std::shared_ptr& source, vid_t src_vertex_capacity, vid_t dst_vertex_capacity) { - if (!source || !meta_ || !meta_->is_bundled() || !out_csr_ || !in_csr_ || + if (!meta_ || !meta_->is_bundled() || !out_csr_ || !in_csr_ || out_csr_->edge_num() != 0 || in_csr_->edge_num() != 0 || (meta_->oe_strategy != EdgeStrategy::kNone && !meta_->oe_mutable) || (meta_->ie_strategy != EdgeStrategy::kNone && !meta_->ie_mutable)) { diff --git a/src/storages/graph/vertex_table.cc b/src/storages/graph/vertex_table.cc index 8c9ae7200..c16334a39 100644 --- a/src/storages/graph/vertex_table.cc +++ b/src/storages/graph/vertex_table.cc @@ -52,7 +52,7 @@ void VertexTable::BatchAddVertices( std::shared_ptr supplier) { CHECK(supplier != nullptr); auto source = supplier->RepeatableSource(); - if (source && ShouldUseBulkBuild(*source) && + if (source && ShouldUseBulkBuild(source->EstimatedBytes()) && try_batch_build_vertices(source)) { return; } @@ -61,16 +61,16 @@ void VertexTable::BatchAddVertices( bool VertexTable::try_batch_build_vertices( const std::shared_ptr& source) { - if (!source || Size() != 0) { + if (Size() != 0) { return false; } auto staged = VertexTable(vertex_schema_); staged.Init(ckp_, memory_level_); const auto source_bytes = source->EstimatedBytes(); - auto options = ResolveBulkBuildSourceOptions( - source_bytes, source->ParallelEnabled(), false, - BulkBuildWorkerStrategy::kMaxProducers); + auto options = + ResolveBulkBuildSourceOptions(source_bytes, source->ParallelEnabled(), + BulkBuildWorkerStrategy::kMaxProducers); auto supplier = source->Open(options); if (!supplier) { return false; diff --git a/src/storages/loader/bundled_edge_csr_loader.cc b/src/storages/loader/bundled_edge_csr_loader.cc index 8bfc70c0b..b9d1793d2 100644 --- a/src/storages/loader/bundled_edge_csr_loader.cc +++ b/src/storages/loader/bundled_edge_csr_loader.cc @@ -20,7 +20,6 @@ #include #include #include -#include #include #include #include @@ -228,10 +227,6 @@ class BundledEdgeCsrLoader::SingleMutableWriter { PutAtVertex(src, dst, data, ts); } - void PutUnique(vid_t src, vid_t dst, const EDATA_T& data, timestamp_t ts) { - PutAtVertex(src, dst, data, ts); - } - void Finish(uint64_t filled_edge_count) { csr_.edge_num_.store(filled_edge_count, std::memory_order_relaxed); csr_.refresh_prefetch_policy(); @@ -295,8 +290,7 @@ class BulkEdgeDataReader { explicit BulkEdgeDataReader(const std::shared_ptr& column) : column_(column.get()) { CHECK(column_ != nullptr); - auto values = std::dynamic_pointer_cast>(column); - if (values) { + if (auto* values = dynamic_cast*>(column_)) { values_ = &values->data(); } } @@ -366,11 +360,6 @@ inline bool is_valid_bulk_edge(vid_t src, vid_t dst) { dst != std::numeric_limits::max(); } -struct BulkEdgeCountSummary { - bool out_single_duplicate = false; - bool in_single_duplicate = false; -}; - constexpr size_t kBulkEdgeInitialGroupReserve = 4096; size_t bulk_edge_group_reserve(size_t rows) { @@ -468,13 +457,13 @@ void count_bulk_edge_chunk(BulkEdgeWorkerScratch& scratch, OutWriter& out, } } -using BulkEdgeCountChunk = std::function; - -BulkEdgeCountSummary count_bulk_edges( - const std::shared_ptr& source, - const IndexerType& src_indexer, const IndexerType& dst_indexer, - ChunkSourceOptions options, std::vector& scratches, - const BulkEdgeCountChunk& count_chunk) { +template +void count_bulk_edges(const std::shared_ptr& source, + const IndexerType& src_indexer, + const IndexerType& dst_indexer, + ChunkSourceOptions options, + std::vector& scratches, + OutWriter& out, InWriter& in) { // Degree accumulation and single-slot uniqueness checks are commutative, so // this pass never needs input order even when fill may later fall back to the // ordered last-write-wins path. @@ -492,17 +481,9 @@ BulkEdgeCountSummary count_bulk_edges( CHECK_LT(worker, worker_count); auto& scratch = scratches[static_cast(worker)]; index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); - count_chunk(scratch, concurrent); + count_bulk_edge_chunk(scratch, out, in, concurrent); }; consume_supplier_indexed(*supplier, options, count); - BulkEdgeCountSummary summary; - for (const auto& scratch : scratches) { - summary.out_single_duplicate = - summary.out_single_duplicate || scratch.out_single_duplicate; - summary.in_single_duplicate = - summary.in_single_duplicate || scratch.in_single_duplicate; - } - return summary; } void append_bulk_edge_group(flat_hash_map& groups, @@ -521,42 +502,6 @@ void append_bulk_edge_group(flat_hash_map& groups, ++group.count; } -template -size_t group_bulk_edge_chunk(BulkEdgeWorkerScratch& scratch) { - CHECK_LE(scratch.src_lids.size(), - static_cast(std::numeric_limits::max())); - if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { - scratch.out_groups.clear(); - scratch.out_groups.reserve( - bulk_edge_group_reserve(scratch.src_lids.size())); - scratch.out_next.assign(scratch.src_lids.size(), kInvalidBulkEdgeIndex); - } - if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { - scratch.in_groups.clear(); - scratch.in_groups.reserve(bulk_edge_group_reserve(scratch.dst_lids.size())); - scratch.in_next.assign(scratch.dst_lids.size(), kInvalidBulkEdgeIndex); - } - size_t valid_edges = 0; - for (size_t row = 0; row < scratch.src_lids.size(); ++row) { - const auto src = scratch.src_lids[row]; - const auto dst = scratch.dst_lids[row]; - if (!is_valid_bulk_edge(src, dst)) { - continue; - } - ++valid_edges; - const auto edge_index = static_cast(row); - if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { - append_bulk_edge_group(scratch.out_groups, scratch.out_next, src, - edge_index); - } - if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { - append_bulk_edge_group(scratch.in_groups, scratch.in_next, dst, - edge_index); - } - } - return valid_edges; -} - template void fill_bulk_edge_chunk_serial(BulkEdgeWorkerScratch& scratch, const BulkEdgeDataReader& data_reader, @@ -587,24 +532,48 @@ void fill_bulk_edge_chunk_concurrent( BulkEdgeWorkerScratch& scratch, const BulkEdgeDataReader& data_reader, OutWriter& out, InWriter& in) { - const auto valid_edges = group_bulk_edge_chunk(scratch); + CHECK_LE(scratch.src_lids.size(), + static_cast(std::numeric_limits::max())); constexpr bool kDirectOut = OutWriter::kStrategy == EdgeStrategy::kSingle; constexpr bool kDirectIn = InWriter::kStrategy == EdgeStrategy::kSingle; - if constexpr (kDirectOut || kDirectIn) { - for (size_t row = 0; row < scratch.src_lids.size(); ++row) { - const auto src = scratch.src_lids[row]; - const auto dst = scratch.dst_lids[row]; - if (!is_valid_bulk_edge(src, dst)) { - continue; - } + if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { + scratch.out_groups.clear(); + scratch.out_groups.reserve( + bulk_edge_group_reserve(scratch.src_lids.size())); + scratch.out_next.assign(scratch.src_lids.size(), kInvalidBulkEdgeIndex); + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { + scratch.in_groups.clear(); + scratch.in_groups.reserve(bulk_edge_group_reserve(scratch.dst_lids.size())); + scratch.in_next.assign(scratch.dst_lids.size(), kInvalidBulkEdgeIndex); + } + + size_t valid_edges = 0; + for (size_t row = 0; row < scratch.src_lids.size(); ++row) { + const auto src = scratch.src_lids[row]; + const auto dst = scratch.dst_lids[row]; + if (!is_valid_bulk_edge(src, dst)) { + continue; + } + ++valid_edges; + if constexpr (kDirectOut || kDirectIn) { const auto data = data_reader.Get(row); if constexpr (kDirectOut) { - out.PutUnique(src, dst, data, 0); + out.PutSerial(src, dst, data, 0); } if constexpr (kDirectIn) { - in.PutUnique(dst, src, data, 0); + in.PutSerial(dst, src, data, 0); } } + const auto edge_index = static_cast(row); + if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { + append_bulk_edge_group(scratch.out_groups, scratch.out_next, src, + edge_index); + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { + append_bulk_edge_group(scratch.in_groups, scratch.in_next, dst, + edge_index); + } } if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { for (const auto& [src, group] : scratch.out_groups) { @@ -641,70 +610,13 @@ void fill_bulk_edge_chunk_concurrent( scratch.filled_edge_count += static_cast(valid_edges); } -template -using BulkEdgeFillChunk = std::function&, - BulkEdgeWorkerScratch&, bool)>; - -template -struct BulkEdgeBuildOps { - bool needs_degree_count = false; - bool checks_out_single = false; - bool checks_in_single = false; - BulkEdgeCountChunk count_chunk; - std::function allocate_from_counts; - BulkEdgeFillChunk fill_chunk; -}; - template -BulkEdgeBuildOps make_bulk_edge_build_ops(OutWriter& out, - InWriter& in) { - BulkEdgeBuildOps ops; - constexpr bool kStoresAnyDirection = - OutWriter::kStrategy != EdgeStrategy::kNone || - InWriter::kStrategy != EdgeStrategy::kNone; - constexpr bool kNeedsAnyDegreeCount = - OutWriter::kStrategy == EdgeStrategy::kMultiple || - InWriter::kStrategy == EdgeStrategy::kMultiple; - ops.needs_degree_count = kNeedsAnyDegreeCount; - ops.checks_out_single = OutWriter::kStrategy == EdgeStrategy::kSingle; - ops.checks_in_single = InWriter::kStrategy == EdgeStrategy::kSingle; - ops.allocate_from_counts = [&]() { - if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { - out.AllocateFromCounts(); - } - if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { - in.AllocateFromCounts(); - } - }; - if constexpr (kNeedsAnyDegreeCount) { - ops.count_chunk = [&out, &in](BulkEdgeWorkerScratch& scratch, - bool concurrent) { - count_bulk_edge_chunk(scratch, out, in, concurrent); - }; - } - if constexpr (kStoresAnyDirection) { - ops.fill_chunk = [&out, &in](const BulkEdgeDataReader& data_reader, - BulkEdgeWorkerScratch& scratch, - bool concurrent) { - if (concurrent) { - fill_bulk_edge_chunk_concurrent(scratch, data_reader, out, in); - } else { - fill_bulk_edge_chunk_serial(scratch, data_reader, out, in); - } - }; - } - return ops; -} - -template uint64_t fill_bulk_edges(const std::shared_ptr& source, const IndexerType& src_indexer, const IndexerType& dst_indexer, const ChunkSourceOptions& options, - bool allow_concurrent_fill, std::vector& scratches, - const BulkEdgeBuildOps& ops) { - CHECK(!allow_concurrent_fill || !options.preserve_order); + OutWriter& out, InWriter& in) { auto supplier = source->Open(options); CHECK(supplier != nullptr); @@ -713,7 +625,7 @@ uint64_t fill_bulk_edges(const std::shared_ptr& source, for (auto& scratch : scratches) { scratch.filled_edge_count = 0; } - if (allow_concurrent_fill && worker_count > 1) { + if (!options.preserve_order && worker_count > 1) { auto fill = [&](int32_t worker, const std::shared_ptr& chunk) { CHECK_GE(worker, 0); CHECK_LT(worker, worker_count); @@ -721,7 +633,7 @@ uint64_t fill_bulk_edges(const std::shared_ptr& source, index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; BulkEdgeDataReader data_reader(data_column); - ops.fill_chunk(data_reader, scratch, true); + fill_bulk_edge_chunk_concurrent(scratch, data_reader, out, in); }; consume_supplier_indexed(*supplier, options, fill); @@ -731,7 +643,7 @@ uint64_t fill_bulk_edges(const std::shared_ptr& source, index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; BulkEdgeDataReader data_reader(data_column); - ops.fill_chunk(data_reader, scratch, false); + fill_bulk_edge_chunk_serial(scratch, data_reader, out, in); } } @@ -744,52 +656,56 @@ uint64_t fill_bulk_edges(const std::shared_ptr& source, return filled_edge_count; } -template -uint64_t build_bundled_edges_with_ops( - const BulkEdgeBuildOps& ops, const IndexerType& src_indexer, +template +uint64_t build_bundled_edges_with_writers( + OutWriter& out, InWriter& in, const IndexerType& src_indexer, const IndexerType& dst_indexer, const std::shared_ptr& source, int64_t source_bytes) { - // VertexTable owns its reserve policy. Edge storage consumes the resulting - // indexer capacity instead of duplicating PropertyGraph::Dump's policy. - if (!ops.fill_chunk) { - ops.allocate_from_counts(); + constexpr bool kStoresAnyDirection = + OutWriter::kStrategy != EdgeStrategy::kNone || + InWriter::kStrategy != EdgeStrategy::kNone; + if constexpr (!kStoresAnyDirection) { return 0; + } else { + // VertexTable owns its reserve policy. Edge storage consumes the resulting + // indexer capacity instead of duplicating PropertyGraph::Dump's policy. + const auto count_options = ResolveBulkBuildSourceOptions( + source_bytes, source->ParallelEnabled(), + BulkBuildWorkerStrategy::kBalancedProducerConsumer); + constexpr bool kNeedsDegreeCount = + OutWriter::kStrategy == EdgeStrategy::kMultiple || + InWriter::kStrategy == EdgeStrategy::kMultiple; + std::vector scratches; + if constexpr (kNeedsDegreeCount) { + scratches.resize(static_cast(count_options.consumer_count)); + count_bulk_edges(source, src_indexer, dst_indexer, count_options, + scratches, out, in); + } + constexpr bool kChecksOutSingle = + OutWriter::kStrategy == EdgeStrategy::kSingle; + constexpr bool kChecksInSingle = + InWriter::kStrategy == EdgeStrategy::kSingle; + const bool single_duplicate = std::any_of( + scratches.begin(), scratches.end(), [](const auto& scratch) { + return (kChecksOutSingle && scratch.out_single_duplicate) || + (kChecksInSingle && scratch.in_single_duplicate); + }); + const bool preserve_fill_order = (kChecksOutSingle || kChecksInSingle) && + (!kNeedsDegreeCount || single_duplicate); + const auto fill_options = + preserve_fill_order ? ChunkSourceOptions{} : count_options; + if (scratches.size() < static_cast(fill_options.consumer_count)) { + scratches.resize(static_cast(fill_options.consumer_count)); + } + if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { + out.AllocateFromCounts(); + } + if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { + in.AllocateFromCounts(); + } + return fill_bulk_edges(source, src_indexer, dst_indexer, + fill_options, scratches, out, in); } - const auto count_options = ResolveBulkBuildSourceOptions( - source_bytes, source->ParallelEnabled(), false, - BulkBuildWorkerStrategy::kBalancedProducerConsumer); - const bool needs_degree_count = ops.needs_degree_count; - std::vector scratches; - if (needs_degree_count) { - scratches.resize(static_cast(count_options.consumer_count)); - } - const bool has_single_direction = - ops.checks_out_single || ops.checks_in_single; - BulkEdgeCountSummary count_summary; - if (needs_degree_count) { - CHECK(static_cast(ops.count_chunk)); - count_summary = count_bulk_edges(source, src_indexer, dst_indexer, - count_options, scratches, ops.count_chunk); - } - const bool single_duplicate = - (ops.checks_out_single && count_summary.out_single_duplicate) || - (ops.checks_in_single && count_summary.in_single_duplicate); - const bool preserve_fill_order = - has_single_direction && (!needs_degree_count || single_duplicate); - const bool allow_concurrent_fill = !preserve_fill_order; - const auto fill_options = - preserve_fill_order - ? ResolveBulkBuildSourceOptions( - source_bytes, source->ParallelEnabled(), true, - BulkBuildWorkerStrategy::kBalancedProducerConsumer) - : count_options; - if (scratches.size() < static_cast(fill_options.consumer_count)) { - scratches.resize(static_cast(fill_options.consumer_count)); - } - ops.allocate_from_counts(); - return fill_bulk_edges(source, src_indexer, dst_indexer, - fill_options, allow_concurrent_fill, - scratches, ops); } template @@ -806,9 +722,9 @@ bool build_bundled_edges_typed(CsrBase* out_csr, CsrBase* in_csr, with_csr_bulk_writer(in_csr, [&](auto& in) { out.PrepareBuild(src_vertex_capacity); in.PrepareBuild(dst_vertex_capacity); - auto ops = make_bulk_edge_build_ops(out, in); - const auto filled_edge_count = build_bundled_edges_with_ops( - ops, src_indexer, dst_indexer, source, source_bytes); + const auto filled_edge_count = + build_bundled_edges_with_writers( + out, in, src_indexer, dst_indexer, source, source_bytes); out.Finish(filled_edge_count); in.Finish(filled_edge_count); }); @@ -851,9 +767,6 @@ bool internal::BundledEdgeCsrLoader::TryBuild( const IndexerType& src_indexer, const IndexerType& dst_indexer, const std::shared_ptr& source, int64_t source_bytes, vid_t src_vertex_capacity, vid_t dst_vertex_capacity) { - if (!source) { - return false; - } return build_bundled_edges(&out_csr, &in_csr, schema, src_indexer, dst_indexer, source, source_bytes, src_vertex_capacity, dst_vertex_capacity); diff --git a/src/storages/loader/loader_utils.cc b/src/storages/loader/loader_utils.cc index da6dcb87b..095f0c0ae 100644 --- a/src/storages/loader/loader_utils.cc +++ b/src/storages/loader/loader_utils.cc @@ -712,11 +712,16 @@ class CsvFileScanner { int64_t total = 0; bool in_quotes = false; // chunk 0 starts outside quotes const RowCounterState* last_selected = nullptr; - std::vector starts_inside(actual_threads, false); - std::vector rows_before(actual_threads, 0); + std::vector> safe_bounds; + if (build_ranges) { + safe_bounds.reserve(actual_threads + 1); + safe_bounds.emplace_back(0, 0); + } for (unsigned i = 0; i < actual_threads; ++i) { - starts_inside[i] = in_quotes; - rows_before[i] = total; + if (build_ranges && i > 0 && !in_quotes && + bounds[i] > safe_bounds.back().first && bounds[i] < file_size) { + safe_bounds.emplace_back(bounds[i], total); + } const RowCounterState& s = in_quotes ? results[i].inside : results[i].outside; total += s.count; @@ -736,18 +741,8 @@ class CsvFileScanner { return output; } - std::vector> safe_bounds; - safe_bounds.emplace_back(0, 0); - for (unsigned i = 1; i < actual_threads; ++i) { - // A nominal split is already positioned after a physical newline. It is - // a valid record boundary exactly when that newline was outside quotes. - // Unsafe splits are merged into the preceding range instead of rescanning - // the file tail to manufacture another boundary. - if (!starts_inside[i] && bounds[i] > safe_bounds.back().first && - bounds[i] < file_size) { - safe_bounds.emplace_back(bounds[i], rows_before[i]); - } - } + // A nominal split is a valid record boundary exactly when the preceding + // newline was outside quotes; unsafe splits remain in the previous range. safe_bounds.emplace_back(file_size, total); output.ranges.reserve(safe_bounds.size() - 1); for (size_t i = 0; i + 1 < safe_bounds.size(); ++i) { @@ -823,22 +818,20 @@ class CsvRangeStreamBuf final : public std::streambuf { return 0; } - std::streamsize total = 0; + const auto requested_count = static_cast(count); const auto buffered = static_cast(egptr() - gptr()); - const auto from_buffer = - std::min(buffered, static_cast(count - total)); + const auto from_buffer = std::min(buffered, requested_count); if (from_buffer > 0) { std::memcpy(destination, gptr(), from_buffer); gbump(static_cast(from_buffer)); - total += static_cast(from_buffer); } - if (total == count || remaining_ == 0) { - return total; + if (from_buffer == requested_count || remaining_ == 0) { + return static_cast(from_buffer); } - const auto requested = - std::min(remaining_, static_cast(count - total)); - file_.read(destination + total, static_cast(requested)); + const auto requested = std::min(remaining_, requested_count - from_buffer); + file_.read(destination + from_buffer, + static_cast(requested)); const auto read = static_cast(file_.gcount()); if (file_.bad()) { THROW_IO_EXCEPTION("Failed to read CSV range: " + file_path_); @@ -847,7 +840,7 @@ class CsvRangeStreamBuf final : public std::streambuf { if (read < requested) { remaining_ = 0; } - return total + static_cast(read); + return static_cast(from_buffer + read); } private: @@ -868,10 +861,6 @@ class CsvRangeStream final : public std::istream { CsvRangeStreamBuf buffer_; }; -} // namespace - -namespace { - std::shared_ptr build_csv_partition_plan( const std::vector& file_paths, const CsvReadConfig& config, int32_t producer_count) { @@ -927,16 +916,9 @@ std::shared_ptr build_csv_partition_plan( // intra-file speculative scans. scan_parallel() uses its caller for one // byte range, therefore the number of active scanners never exceeds H. const auto scan_budget = chunk_pipeline_detail::hardware_worker_count(); - std::vector scan_threads(file_paths.size(), 0); + auto scan_threads = range_counts; if (non_empty_files < static_cast(scan_budget)) { - size_t assigned_scanners = 0; - for (size_t i = 0; i < file_paths.size(); ++i) { - if (file_sizes[i] == 0) { - continue; - } - scan_threads[i] = std::max(1, range_counts[i]); - assigned_scanners += static_cast(scan_threads[i]); - } + size_t assigned_scanners = assigned_ranges; CHECK_LE(assigned_scanners, static_cast(scan_budget)); while (assigned_scanners < static_cast(scan_budget)) { size_t best = file_paths.size(); @@ -959,22 +941,14 @@ std::shared_ptr build_csv_partition_plan( ++scan_threads[best]; ++assigned_scanners; } - } else { - for (size_t i = 0; i < file_paths.size(); ++i) { - scan_threads[i] = file_sizes[i] == 0 ? 0 : 1; - } } std::vector scans(file_paths.size()); std::atomic next_file{0}; std::atomic scan_cancelled{false}; - std::mutex scan_error_mutex; std::exception_ptr scan_error; auto capture_scan_error = [&](std::exception_ptr error) { - bool expected = false; - if (scan_cancelled.compare_exchange_strong(expected, true, - std::memory_order_acq_rel)) { - std::lock_guard lock(scan_error_mutex); + if (!scan_cancelled.exchange(true, std::memory_order_acq_rel)) { scan_error = std::move(error); } }; @@ -990,8 +964,8 @@ std::shared_ptr build_csv_partition_plan( CsvFileScanner(file_paths[file_index], config.quoting, config.quote_char, config.delimiter, config.use_threads) .scan_with_ranges(static_cast(file_sizes[file_index]), - std::max(1, range_counts[file_index]), - std::max(1, scan_threads[file_index])); + range_counts[file_index], + scan_threads[file_index]); }; std::vector planners; @@ -1036,11 +1010,13 @@ std::shared_ptr build_csv_partition_plan( int64_t total = 0; for (size_t i = 0; i < file_paths.size(); ++i) { const auto& scan = scans[i]; - if (scan.row_count < 0 || - scan.row_count > std::numeric_limits::max() - total) { - total = kUnknownRowNum; - } else if (total != kUnknownRowNum) { - total += scan.row_count; + if (total != kUnknownRowNum) { + if (scan.row_count < 0 || + scan.row_count > std::numeric_limits::max() - total) { + total = kUnknownRowNum; + } else { + total += scan.row_count; + } } // SKIP is scoped per input file. Distribute it over record-aligned ranges @@ -1616,7 +1592,7 @@ class SourceBackedChunkSupplier final : public IDataChunkSupplier { } private: - std::shared_ptr OpenSupplier() const { + IDataChunkSupplier* OpenSupplier() const { std::lock_guard lock(mutex_); if (!supplier_) { supplier_ = source_->Open(); @@ -1627,7 +1603,7 @@ class SourceBackedChunkSupplier final : public IDataChunkSupplier { supplier_->Cancel(); } } - return supplier_; + return supplier_.get(); } std::shared_ptr source_; @@ -1831,7 +1807,7 @@ class ChainedCsvChunkSupplier final : public IDataChunkSupplier { std::shared_ptr GetNextChunk() override { while (next_file_ < file_paths_.size()) { if (!current_) { - current_ = std::make_shared( + current_ = std::make_unique( file_paths_[next_file_], config_, CsvRowCountMode::kUnknown); } auto chunk = current_->GetNextChunk(); @@ -1850,9 +1826,8 @@ class ChainedCsvChunkSupplier final : public IDataChunkSupplier { } int64_t total = 0; for (const auto& file_path : file_paths_) { - auto supplier = std::make_shared( - file_path, config_, CsvRowCountMode::kUnknown); - auto rows = supplier->RowNum(); + CSVChunkSupplier supplier(file_path, config_, CsvRowCountMode::kUnknown); + auto rows = supplier.RowNum(); if (rows < 0 || rows > std::numeric_limits::max() - total) { return kUnknownRowNum; } @@ -1865,7 +1840,7 @@ class ChainedCsvChunkSupplier final : public IDataChunkSupplier { private: std::vector file_paths_; CsvReadConfig config_; - std::shared_ptr current_; + std::unique_ptr current_; size_t next_file_ = 0; mutable int64_t row_num_ = kUnknownRowNum; }; @@ -1881,26 +1856,24 @@ std::shared_ptr make_data_chunk_supplier( } ChunkSourceOptions ResolveBulkBuildSourceOptions( - int64_t source_bytes, bool parallel_enabled, bool preserve_order, + int64_t source_bytes, bool parallel_enabled, BulkBuildWorkerStrategy worker_strategy) { constexpr int64_t kMinPartitionBytes = 64LL * 1024 * 1024; constexpr size_t kMaxQueuedChunks = 64; ChunkSourceOptions options; - options.preserve_order = preserve_order; + options.preserve_order = false; const auto workers = chunk_pipeline_detail::hardware_worker_count(); - if (!parallel_enabled || preserve_order || workers <= 1 || + if (!parallel_enabled || workers <= 1 || source_bytes < kDefaultBulkBuildMinBytes) { return options; } - const auto useful_partitions = std::max( - 1, source_bytes / kMinPartitionBytes + - (source_bytes % kMinPartitionBytes == 0 ? 0 : 1)); + const auto useful_partitions = (source_bytes - 1) / kMinPartitionBytes + 1; switch (worker_strategy) { case BulkBuildWorkerStrategy::kMaxProducers: { - options.producer_count = static_cast(std::min( - workers - 1, std::min(useful_partitions, workers))); + options.producer_count = + static_cast(std::min(useful_partitions, workers - 1)); break; } case BulkBuildWorkerStrategy::kBalancedProducerConsumer: { @@ -1912,7 +1885,6 @@ ChunkSourceOptions ResolveBulkBuildSourceOptions( break; } } - options.producer_count = std::max(1, options.producer_count); options.queue_capacity = std::clamp( static_cast(options.producer_count) * 2, 2, kMaxQueuedChunks); return options; From 3a74fce42c527c008b5386a04fe2de6e435996b9 Mon Sep 17 00:00:00 2001 From: "xiaolei.zl" Date: Wed, 22 Jul 2026 11:49:37 +0800 Subject: [PATCH 6/8] avoid using vector_t in BulkEdgeDataReader --- src/storages/loader/bundled_edge_csr_loader.cc | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/storages/loader/bundled_edge_csr_loader.cc b/src/storages/loader/bundled_edge_csr_loader.cc index b9d1793d2..4e86f3066 100644 --- a/src/storages/loader/bundled_edge_csr_loader.cc +++ b/src/storages/loader/bundled_edge_csr_loader.cc @@ -288,23 +288,21 @@ template class BulkEdgeDataReader { public: explicit BulkEdgeDataReader(const std::shared_ptr& column) - : column_(column.get()) { + : column_(column.get()), + value_column_(dynamic_cast*>(column.get())) { CHECK(column_ != nullptr); - if (auto* values = dynamic_cast*>(column_)) { - values_ = &values->data(); - } } EDATA_T Get(size_t row) const { - if (values_ != nullptr) { - return (*values_)[row]; + if (value_column_ != nullptr) { + return value_column_->get_value(row); } return column_->get_elem(row).template GetValue(); } private: const IContextColumn* column_; - const vector_t* values_ = nullptr; + const ValueColumn* value_column_; }; template <> From 04ab362910d7d5311fe7c154e5a56403fd372385 Mon Sep 17 00:00:00 2001 From: "xiaolei.zl" Date: Thu, 23 Jul 2026 18:12:40 +0800 Subject: [PATCH 7/8] use spill for csr loading to avoid read twice. --- .../function/import/csv_read_function.h | 2 +- .../neug/compiler/function/read_function.h | 6 +- .../execute/ops/batch/batch_insert_edge.h | 2 +- .../execute/ops/batch/batch_insert_vertex.h | 2 +- .../execute/ops/batch/batch_update_utils.h | 4 +- include/neug/main/query_processor.h | 7 +- include/neug/storages/graph/edge_table.h | 11 +- include/neug/storages/graph/graph_interface.h | 33 +- include/neug/storages/graph/property_graph.h | 6 +- include/neug/storages/graph/vertex_table.h | 10 +- .../storages/loader/chunk_pipeline_utils.h | 52 +- include/neug/storages/loader/loader_utils.h | 57 +- include/neug/transaction/insert_transaction.h | 7 +- include/neug/transaction/update_transaction.h | 7 +- include/neug/utils/io/read/csv/csv_reader.h | 4 +- .../execute/ops/batch/batch_insert_edge.cc | 5 +- .../execute/ops/batch/batch_insert_vertex.cc | 9 +- .../execute/ops/batch/batch_update_utils.cc | 14 +- src/main/query_processor.cc | 60 ++- src/storages/graph/edge_table.cc | 39 +- src/storages/graph/graph_interface.cc | 9 +- src/storages/graph/property_graph.cc | 16 +- src/storages/graph/vertex_table.cc | 27 +- .../loader/abstract_property_graph_loader.cc | 6 +- .../loader/bundled_edge_csr_loader.cc | 485 ++++++++++++++---- src/storages/loader/bundled_edge_csr_loader.h | 14 +- src/storages/loader/loader_utils.cc | 209 +++++--- src/transaction/insert_transaction.cc | 4 +- src/transaction/update_transaction.cc | 4 +- src/utils/io/read/csv/csv_reader.cc | 4 +- tests/storage/alter_property_test.cc | 6 +- tests/storage/test_edge_table.cc | 358 ++++++------- tests/storage/test_vertex_table.cc | 6 +- tests/unittest/utils.h | 22 +- tests/utils/test_reader.cc | 117 +++++ 35 files changed, 1083 insertions(+), 541 deletions(-) diff --git a/include/neug/compiler/function/import/csv_read_function.h b/include/neug/compiler/function/import/csv_read_function.h index f1a936443..f278e1ef3 100644 --- a/include/neug/compiler/function/import/csv_read_function.h +++ b/include/neug/compiler/function/import/csv_read_function.h @@ -143,7 +143,7 @@ struct CSVReadFunction { return ctx; } - static std::shared_ptr sourceFunc( + static std::unique_ptr sourceFunc( std::shared_ptr state, std::vector projected_columns) { if (!state) { diff --git a/include/neug/compiler/function/read_function.h b/include/neug/compiler/function/read_function.h index 298bdc4c6..42bc6fec6 100644 --- a/include/neug/compiler/function/read_function.h +++ b/include/neug/compiler/function/read_function.h @@ -35,9 +35,9 @@ namespace function { using read_exec_func_t = std::function state)>; -/// Creates a repeatable source for terminal ingestion. Storage may consume it -/// once through normal BatchAdd or reopen it for a staged bulk build. -using read_source_func_t = std::function( +/// Creates a configurable source for terminal ingestion. Storage selects the +/// projection and concurrency plan before opening it once. +using read_source_func_t = std::function( std::shared_ptr state, std::vector projected_columns)>; diff --git a/include/neug/execution/execute/ops/batch/batch_insert_edge.h b/include/neug/execution/execute/ops/batch/batch_insert_edge.h index d0630f3d2..0877a2101 100644 --- a/include/neug/execution/execute/ops/batch/batch_insert_edge.h +++ b/include/neug/execution/execute/ops/batch/batch_insert_edge.h @@ -41,7 +41,7 @@ class BatchInsertEdgeOprBuilder : public IOperatorBuilder { }; /// Fuses only a terminal, empty-sink COPY FROM plan. Storage chooses staged -/// build or normal BatchAdd from the supplied repeatable source. +/// build or normal BatchAdd before opening the supplied source once. class BatchInsertEdgeFromSourceOprBuilder : public IOperatorBuilder { public: neug::result Build(const Schema& schema, diff --git a/include/neug/execution/execute/ops/batch/batch_insert_vertex.h b/include/neug/execution/execute/ops/batch/batch_insert_vertex.h index 379cef23f..82fbee605 100644 --- a/include/neug/execution/execute/ops/batch/batch_insert_vertex.h +++ b/include/neug/execution/execute/ops/batch/batch_insert_vertex.h @@ -40,7 +40,7 @@ class BatchInsertVertexOprBuilder : public IOperatorBuilder { }; /// Fuses only a terminal, empty-sink COPY FROM plan. Storage chooses staged -/// build or normal BatchAdd from the supplied repeatable source. +/// build or normal BatchAdd before opening the supplied source once. class BatchInsertVertexFromSourceOprBuilder : public IOperatorBuilder { public: neug::result Build(const Schema& schema, diff --git a/include/neug/execution/execute/ops/batch/batch_update_utils.h b/include/neug/execution/execute/ops/batch/batch_update_utils.h index 2cf244c25..f8a2f50ab 100644 --- a/include/neug/execution/execute/ops/batch/batch_update_utils.h +++ b/include/neug/execution/execute/ops/batch/batch_update_utils.h @@ -18,6 +18,7 @@ #include "neug/common/types/graph_types.h" #include "neug/execution/common/context.h" +#include "neug/storages/loader/loader_utils.h" #include "neug/utils/property/types.h" namespace physical { @@ -35,7 +36,6 @@ class RepeatedPtrField; } // namespace google namespace neug { -class IDataChunkSupplier; class Schema; class StorageReadInterface; namespace function { @@ -79,7 +79,7 @@ bool resolve_vertex_label_id(const Schema& schema, const ::common::NameOrId& type, label_t& label_id); struct BatchInsertInput { - std::shared_ptr supplier; + std::unique_ptr data; Context output; }; diff --git a/include/neug/main/query_processor.h b/include/neug/main/query_processor.h index c41aefa75..a5982f6a2 100644 --- a/include/neug/main/query_processor.h +++ b/include/neug/main/query_processor.h @@ -72,14 +72,15 @@ class QueryProcessor { result>> check_and_retrieve_pipeline(const PropertyGraph& pg, const std::string& query_string, - const std::string& access_mode, - int32_t num_threads); + const std::string& access_mode); + + result resolve_thread_budget(int32_t requested_threads) const; result execute_internal( SnapshotGuard& guard, const std::string& query_string, std::shared_ptr cache_value, AccessMode access_mode, const execution::ParamsMap& parameters = {}, - int32_t num_threads = 0); + int32_t thread_budget = 0); result execute_explain_mode( const std::string& query_string, diff --git a/include/neug/storages/graph/edge_table.h b/include/neug/storages/graph/edge_table.h index 52d78f8e9..a9712695f 100644 --- a/include/neug/storages/graph/edge_table.h +++ b/include/neug/storages/graph/edge_table.h @@ -29,6 +29,7 @@ #include "neug/storages/csr/csr_base.h" #include "neug/storages/csr/csr_view.h" #include "neug/storages/graph/schema.h" +#include "neug/storages/loader/loader_utils.h" #include "neug/storages/module/module.h" #include "neug/utils/indexers.h" #include "neug/utils/property/table.h" @@ -40,9 +41,6 @@ class ModuleBroker; class CheckpointManifest; class PropertyGraph; -class IDataChunkSupplier; -class IDataChunkSource; - class EdgeTable { public: EdgeTable(std::shared_ptr meta) : meta_(meta) {} @@ -134,7 +132,8 @@ class EdgeTable { void BatchAddEdges(const IndexerType& src_indexer, const IndexerType& dst_indexer, - std::shared_ptr supplier); + std::unique_ptr source, + BulkLoadOptions options = {}); // Add edges in batch to the edge table. void BatchAddEdges(const std::vector& src_lid_list, @@ -193,8 +192,8 @@ class EdgeTable { private: bool TryBatchBuildEdges(const IndexerType& src_indexer, const IndexerType& dst_indexer, - const std::shared_ptr& source, - vid_t src_vertex_capacity, vid_t dst_vertex_capacity); + IDataChunkSource& source, vid_t src_vertex_capacity, + vid_t dst_vertex_capacity, BulkLoadOptions options); void dropAndCreateNewBundledCSR(Checkpoint& ckp, ColumnBase* prev_data_col); void dropAndCreateNewUnbundledCSR(Checkpoint& ckp, bool delete_property); diff --git a/include/neug/storages/graph/graph_interface.h b/include/neug/storages/graph/graph_interface.h index d604239a1..6f6601c05 100644 --- a/include/neug/storages/graph/graph_interface.h +++ b/include/neug/storages/graph/graph_interface.h @@ -22,6 +22,7 @@ #include "neug/storages/graph/graph_view.h" #include "neug/storages/graph/property_graph.h" #include "neug/storages/graph/schema.h" +#include "neug/storages/loader/loader_utils.h" #include "neug/utils/property/types.h" namespace neug { @@ -423,27 +424,27 @@ class StorageInsertInterface : virtual public IStorageInterface { const void*& prop) = 0; /** - * @brief Batch insert vertices from a record supplier. + * @brief Batch insert vertices from an operation-owned record source. * * @param v_label_id Vertex label for all records - * @param supplier Record batch data source + * @param source Record batch data source * @return Status indicating success or failure */ - virtual Status BatchAddVertices( - label_t v_label_id, std::shared_ptr supplier) = 0; + virtual Status BatchAddVertices(label_t v_label_id, + std::unique_ptr source) = 0; /** - * @brief Batch insert edges from a record supplier. + * @brief Batch insert edges from an operation-owned record source. * * @param src_label Source vertex label * @param dst_label Destination vertex label * @param edge_label Edge label - * @param supplier Record batch data source + * @param source Record batch data source * @return Status indicating success or failure */ - virtual Status BatchAddEdges( - label_t src_label, label_t dst_label, label_t edge_label, - std::shared_ptr supplier) = 0; + virtual Status BatchAddEdges(label_t src_label, label_t dst_label, + label_t edge_label, + std::unique_ptr source) = 0; }; /** @@ -620,12 +621,14 @@ class StorageAPUpdateInterface : public StorageUpdateInterface { public: explicit StorageAPUpdateInterface(PropertyGraph& graph, GraphView& view, timestamp_t timestamp, - neug::Allocator& alloc) + neug::Allocator& alloc, + BulkLoadOptions bulk_load_options = {}) : StorageUpdateInterface(view, timestamp), graph_(graph), mut_view_(view), alloc_(alloc), - timestamp_(timestamp) {} + timestamp_(timestamp), + bulk_load_options_(bulk_load_options) {} ~StorageAPUpdateInterface() {} Status UpdateVertexProperty(label_t label, vid_t lid, int col_id, @@ -646,11 +649,10 @@ class StorageAPUpdateInterface : public StorageUpdateInterface { Status DeleteEdges(label_t src_label, vid_t src, label_t dst_label, vid_t dst, label_t edge_label) override; void CreateCheckpoint() override; - Status BatchAddVertices( - label_t v_label_id, - std::shared_ptr supplier) override; + Status BatchAddVertices(label_t v_label_id, + std::unique_ptr source) override; Status BatchAddEdges(label_t src_label, label_t dst_label, label_t edge_label, - std::shared_ptr supplier) override; + std::unique_ptr source) override; Status BatchDeleteVertices(label_t v_label_id, const std::vector& vids) override; Status BatchDeleteEdges( @@ -681,6 +683,7 @@ class StorageAPUpdateInterface : public StorageUpdateInterface { GraphView& mut_view_; neug::Allocator& alloc_; timestamp_t timestamp_; + BulkLoadOptions bulk_load_options_; }; } // namespace neug diff --git a/include/neug/storages/graph/property_graph.h b/include/neug/storages/graph/property_graph.h index 4b12e6403..7bac6ebb3 100644 --- a/include/neug/storages/graph/property_graph.h +++ b/include/neug/storages/graph/property_graph.h @@ -291,10 +291,12 @@ class PropertyGraph { size_t capacity); Status BatchAddVertices(label_t v_label_id, - std::shared_ptr supplier); + std::unique_ptr source, + BulkLoadOptions options = {}); Status BatchAddEdges(label_t src_label, label_t dst_label, label_t edge_label, - std::shared_ptr supplier); + std::unique_ptr source, + BulkLoadOptions options = {}); Status BatchDeleteVertices(label_t v_label_id, const std::vector& vids); diff --git a/include/neug/storages/graph/vertex_table.h b/include/neug/storages/graph/vertex_table.h index d8f8d56d8..86cb1dfa6 100644 --- a/include/neug/storages/graph/vertex_table.h +++ b/include/neug/storages/graph/vertex_table.h @@ -18,6 +18,7 @@ #include "neug/common/types/value.h" #include "neug/storages/graph/schema.h" #include "neug/storages/graph/vertex_timestamp.h" +#include "neug/storages/loader/loader_utils.h" #include "neug/storages/module/module.h" #include "neug/utils/indexers.h" #include "neug/utils/property/table.h" @@ -27,8 +28,6 @@ namespace neug { class ModuleBroker; class CheckpointManifest; class Checkpoint; -class IDataChunkSupplier; -class IDataChunkSource; class VertexTableView; class VertexSet { @@ -258,7 +257,8 @@ class VertexTable { void Compact(timestamp_t ts = MAX_TIMESTAMP); - void BatchAddVertices(std::shared_ptr supplier); + void BatchAddVertices(std::unique_ptr source, + BulkLoadOptions options = {}); const VertexTimestamp& get_vertex_timestamp() const { return *v_ts_; } @@ -266,8 +266,8 @@ class VertexTable { Table& get_table() { return *table_; } private: - bool try_batch_build_vertices( - const std::shared_ptr& source); + bool try_batch_build_vertices(IDataChunkSource& source, + BulkLoadOptions options); void batch_add_vertices_impl(std::shared_ptr supplier); diff --git a/include/neug/storages/loader/chunk_pipeline_utils.h b/include/neug/storages/loader/chunk_pipeline_utils.h index 6bdb5fcbe..298da84ad 100644 --- a/include/neug/storages/loader/chunk_pipeline_utils.h +++ b/include/neug/storages/loader/chunk_pipeline_utils.h @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -96,14 +97,18 @@ inline void consume_chunk_pipeline_impl(IDataChunkSupplier& supplier, Consume&& consume) { consumer_count = std::max(1, consumer_count); if (consumer_count == 1) { + uint64_t next_row_ordinal = 0; while (auto chunk = supplier.GetNextChunk()) { - consume(0, chunk); + const auto row_count = chunk->row_num(); + consume(0, SequencedDataChunk{std::move(chunk), next_row_ordinal}); + CHECK_LE(row_count, + std::numeric_limits::max() - next_row_ordinal); + next_row_ordinal += static_cast(row_count); } return; } - chunk_pipeline_detail::BoundedQueue> queue( - queue_capacity); + chunk_pipeline_detail::BoundedQueue queue(queue_capacity); std::atomic cancelled{false}; std::mutex error_mutex; std::exception_ptr first_error; @@ -126,11 +131,19 @@ inline void consume_chunk_pipeline_impl(IDataChunkSupplier& supplier, std::thread producer([&] { try { + uint64_t next_row_ordinal = 0; while (!cancelled.load(std::memory_order_acquire)) { auto chunk = supplier.GetNextChunk(); - if (!chunk || !queue.Push(std::move(chunk))) { + if (!chunk) { + break; + } + const auto row_count = chunk->row_num(); + if (!queue.Push({std::move(chunk), next_row_ordinal})) { break; } + CHECK_LE(row_count, + std::numeric_limits::max() - next_row_ordinal); + next_row_ordinal += static_cast(row_count); } } catch (...) { capture_error(std::current_exception()); } queue.Close(); @@ -141,7 +154,7 @@ inline void consume_chunk_pipeline_impl(IDataChunkSupplier& supplier, for (int32_t i = 0; i < consumer_count; ++i) { consumers.emplace_back([&, i] { try { - std::shared_ptr chunk; + SequencedDataChunk chunk; while (!cancelled.load(std::memory_order_acquire) && queue.Pop(chunk)) { consume(i, chunk); } @@ -168,7 +181,11 @@ inline void consume_concurrent_supplier_impl(IDataChunkSupplier& supplier, CHECK(supplier.SupportsConcurrentGetNext()); consumer_count = std::max(1, consumer_count); if (consumer_count == 1) { - while (auto chunk = supplier.GetNextChunk()) { + while (true) { + auto chunk = supplier.GetNextChunkWithOrdinal(); + if (!chunk.chunk) { + break; + } consume(0, chunk); } return; @@ -195,8 +212,8 @@ inline void consume_concurrent_supplier_impl(IDataChunkSupplier& supplier, consumers.emplace_back([&, i] { try { while (!cancelled.load(std::memory_order_acquire)) { - auto chunk = supplier.GetNextChunk(); - if (!chunk) { + auto chunk = supplier.GetNextChunkWithOrdinal(); + if (!chunk.chunk) { break; } consume(i, chunk); @@ -221,13 +238,28 @@ template inline void consume_supplier_indexed(IDataChunkSupplier& supplier, const ChunkSourceOptions& options, Consume&& consume) { + const auto normalized = NormalizeChunkSourceOptions(options); + auto consumer_count = normalized.consumer_count; if (supplier.SupportsConcurrentGetNext()) { chunk_pipeline_detail::consume_concurrent_supplier_impl( - supplier, options.consumer_count, std::forward(consume)); + supplier, consumer_count, std::forward(consume)); return; } + + // A non-concurrent supplier needs an extra forwarding producer whenever + // several consumers are used. Account for that thread at the execution + // boundary as well; falling back to one inline consumer avoids the extra + // thread when the remaining budget is too small. + if (consumer_count > 1) { + const auto max_consumers_with_forwarder = + normalized.worker_budget - normalized.producer_count - 1; + consumer_count = std::min(consumer_count, max_consumers_with_forwarder); + if (consumer_count < 2) { + consumer_count = 1; + } + } chunk_pipeline_detail::consume_chunk_pipeline_impl( - supplier, options.consumer_count, options.queue_capacity, + supplier, consumer_count, normalized.queue_capacity, std::forward(consume)); } diff --git a/include/neug/storages/loader/loader_utils.h b/include/neug/storages/loader/loader_utils.h index 2425f3cd4..c499fc3a3 100644 --- a/include/neug/storages/loader/loader_utils.h +++ b/include/neug/storages/loader/loader_utils.h @@ -74,23 +74,29 @@ CsvReadConfig build_csv_read_config( class IDataChunkSource; +struct SequencedDataChunk { + std::shared_ptr chunk; + uint64_t first_row_ordinal = 0; +}; + class IDataChunkSupplier { public: virtual ~IDataChunkSupplier() = default; virtual std::shared_ptr GetNextChunk() = 0; virtual int64_t RowNum() const = 0; + /// Returns a chunk together with its stable source row ordinal. Concurrent + /// suppliers override this when their source advertises stable ordinals. + virtual SequencedDataChunk GetNextChunkWithOrdinal() { + return {GetNextChunk(), 0}; + } + virtual bool ProvidesStableRowOrdinals() const { return false; } + /// Whether GetNextChunk() may be called concurrently by several consumers. virtual bool SupportsConcurrentGetNext() const { return false; } /// Stops any background producers and wakes blocked GetNextChunk() calls. virtual void Cancel() {} - - /// Returns the repeatable source backing this supplier, when available. - /// Storage uses this hint to select staged bulk build internally. - virtual std::shared_ptr RepeatableSource() const { - return nullptr; - } }; struct ChunkSourceOptions { @@ -98,6 +104,8 @@ struct ChunkSourceOptions { /// many background parsing workers when input order need not be preserved. int32_t producer_count = 0; int32_t consumer_count = 1; + /// Maximum concurrently active workers for source planning and execution. + int32_t worker_budget = 1; size_t queue_capacity = 2; bool preserve_order = true; @@ -106,7 +114,12 @@ struct ChunkSourceOptions { std::vector projected_columns; }; -/// A repeatable source of data chunks. +/// Clamps source concurrency to a valid execution plan. The normalized result +/// always reserves at least one consumer and satisfies +/// producer_count + consumer_count <= worker_budget. +ChunkSourceOptions NormalizeChunkSourceOptions(ChunkSourceOptions options); + +/// A configurable source of data chunks. /// /// Keeping source setup separate from the supplier cursor lets storage consume /// the input directly without making execution::Context stateful or lazy. @@ -117,25 +130,39 @@ class IDataChunkSource { /// Opens a new supplier positioned at the beginning of the source using the /// requested projection and concurrency settings. virtual std::shared_ptr Open( - const ChunkSourceOptions& options = {}) const = 0; + const ChunkSourceOptions& options = {}) = 0; /// Returns the source size when cheaply known, otherwise -1. virtual int64_t EstimatedBytes() const { return -1; } /// Whether the source options permit parallel parsing and consumption. virtual bool ParallelEnabled() const { return true; } + + /// Whether parallel suppliers return globally ordered row ordinals. + virtual bool ProvidesStableRowOrdinals() const { return false; } }; -std::shared_ptr make_data_chunk_supplier( +/// Adapts an already-open supplier to an operation-owned source. +std::unique_ptr make_data_chunk_source( + std::shared_ptr supplier); +std::unique_ptr make_data_chunk_source( std::shared_ptr source); +std::shared_ptr open_data_chunk_source( + IDataChunkSource& source, const ChunkSourceOptions& options = {}); + +/// Operation-scoped resource limits for storage bulk loading. +struct BulkLoadOptions { + int32_t worker_budget = 1; +}; + enum class BulkBuildWorkerStrategy { kMaxProducers, kBalancedProducerConsumer, }; ChunkSourceOptions ResolveBulkBuildSourceOptions( - int64_t source_bytes, bool parallel_enabled, + int64_t source_bytes, bool parallel_enabled, int32_t worker_budget, BulkBuildWorkerStrategy worker_strategy); inline constexpr int64_t kUnknownRowNum = -1; @@ -184,25 +211,23 @@ class CSVChunkSupplier : public IDataChunkSupplier { struct CsvPartitionPlanCache; -/// Reopens the public CSV parser for each pass. Parallel suppliers share a -/// cached record-aligned partition plan so repeated passes do not repeat the -/// raw row-count scan. +/// Opens the public CSV parser with a storage-selected concurrency plan. class CSVChunkSource final : public IDataChunkSource { public: CSVChunkSource(std::vector file_paths, CsvReadConfig config, std::vector projected_columns = {}); std::shared_ptr Open( - const ChunkSourceOptions& options = {}) const override; + const ChunkSourceOptions& options = {}) override; int64_t EstimatedBytes() const override; bool ParallelEnabled() const override { return config_.use_threads; } + bool ProvidesStableRowOrdinals() const override { return true; } private: std::vector file_paths_; CsvReadConfig config_; std::vector projected_columns_; - // Source-local cache: file paths and partition-relevant config stay fixed - // across Open() calls; producer count selects the cached plan. + // Source-local cache keeps partition planning lazy until Open(). std::shared_ptr partition_plan_cache_; }; diff --git a/include/neug/transaction/insert_transaction.h b/include/neug/transaction/insert_transaction.h index 0e9efda1c..4aa1a1e25 100644 --- a/include/neug/transaction/insert_transaction.h +++ b/include/neug/transaction/insert_transaction.h @@ -252,12 +252,11 @@ class StorageTPInsertInterface : public StorageInsertInterface { return txn_.GetVertexIndex(label, id, index); } - Status BatchAddVertices( - label_t v_label_id, - std::shared_ptr supplier) override; + Status BatchAddVertices(label_t v_label_id, + std::unique_ptr source) override; Status BatchAddEdges(label_t src_label, label_t dst_label, label_t edge_label, - std::shared_ptr supplier) override; + std::unique_ptr source) override; private: InsertTransaction& txn_; diff --git a/include/neug/transaction/update_transaction.h b/include/neug/transaction/update_transaction.h index 2ad699e90..de96b72b4 100644 --- a/include/neug/transaction/update_transaction.h +++ b/include/neug/transaction/update_transaction.h @@ -201,11 +201,10 @@ class StorageTPUpdateInterface : public StorageUpdateInterface { // --- Batch methods --- void CreateCheckpoint() override; - Status BatchAddVertices( - label_t v_label_id, - std::shared_ptr supplier) override; + Status BatchAddVertices(label_t v_label_id, + std::unique_ptr source) override; Status BatchAddEdges(label_t src_label, label_t dst_label, label_t edge_label, - std::shared_ptr supplier) override; + std::unique_ptr source) override; Status BatchDeleteVertices(label_t v_label_id, const std::vector& vids) override; Status BatchDeleteEdges( diff --git a/include/neug/utils/io/read/csv/csv_reader.h b/include/neug/utils/io/read/csv/csv_reader.h index 71bede42f..b637aa83a 100644 --- a/include/neug/utils/io/read/csv/csv_reader.h +++ b/include/neug/utils/io/read/csv/csv_reader.h @@ -44,9 +44,9 @@ class CsvReader { void read(std::shared_ptr localState, execution::Context& ctx); - /// Creates a repeatable CSV source for direct COPY FROM bulk loading. + /// Creates a configurable CSV source for direct COPY FROM bulk loading. /// Returns nullptr when the read needs a row filter and must materialize. - std::shared_ptr createChunkSource( + std::unique_ptr createChunkSource( std::vector projected_columns = {}); result> inferSchema(); diff --git a/src/execution/execute/ops/batch/batch_insert_edge.cc b/src/execution/execute/ops/batch/batch_insert_edge.cc index ba0356586..d327f5d22 100644 --- a/src/execution/execute/ops/batch/batch_insert_edge.cc +++ b/src/execution/execute/ops/batch/batch_insert_edge.cc @@ -126,11 +126,12 @@ neug::result BatchInsertEdgeOpr::Eval( input = create_batch_insert_input(source_->state, *source_->read_function, mappings); } else { - input.supplier = create_data_chunk_supplier(ctx, mappings); + input.data = + make_data_chunk_source(create_data_chunk_supplier(ctx, mappings)); input.output = std::move(ctx); } RETURN_STATUS_ERROR_IF_NOT_OK(graph.BatchAddEdges( - src_label_id, dst_label_id, edge_label_id, std::move(input.supplier))); + src_label_id, dst_label_id, edge_label_id, std::move(input.data))); return neug::result(std::move(input.output)); } diff --git a/src/execution/execute/ops/batch/batch_insert_vertex.cc b/src/execution/execute/ops/batch/batch_insert_vertex.cc index 4f8b1143b..2bad3defb 100644 --- a/src/execution/execute/ops/batch/batch_insert_vertex.cc +++ b/src/execution/execute/ops/batch/batch_insert_vertex.cc @@ -64,18 +64,21 @@ neug::result BatchInsertVertexOpr::Eval( label_t vertex_label_id = 0; if (!resolve_vertex_label_id(graph.schema(), vertex_type_, vertex_label_id)) { RETURN_STATUS_ERROR(StatusCode::ERR_INVALID_ARGUMENT, - "Failed to resolve vertex type for BatchInsertVertex"); + "Failed to resolve vertex type " + + vertex_type_.ShortDebugString() + + " for BatchInsertVertex"); } BatchInsertInput input; if (source_) { input = create_batch_insert_input(source_->state, *source_->read_function, property_mappings_); } else { - input.supplier = create_data_chunk_supplier(ctx, property_mappings_); + input.data = make_data_chunk_source( + create_data_chunk_supplier(ctx, property_mappings_)); input.output = std::move(ctx); } RETURN_STATUS_ERROR_IF_NOT_OK( - graph.BatchAddVertices(vertex_label_id, std::move(input.supplier))); + graph.BatchAddVertices(vertex_label_id, std::move(input.data))); return neug::result(std::move(input.output)); } diff --git a/src/execution/execute/ops/batch/batch_update_utils.cc b/src/execution/execute/ops/batch/batch_update_utils.cc index 330fd5f6e..1051067aa 100644 --- a/src/execution/execute/ops/batch/batch_update_utils.cc +++ b/src/execution/execute/ops/batch/batch_update_utils.cc @@ -375,8 +375,14 @@ BatchInsertSource build_batch_insert_source(const physical::PhysicalPlan& plan, auto catalog = neug::main::MetadataRegistry::getCatalog(); auto registered_function = catalog->getFunctionWithSignature(source.extension_name()); - return {std::move(state), - registered_function->ptrCast()}; + auto* read_function = + dynamic_cast(registered_function); + if (read_function == nullptr) { + THROW_INVALID_ARGUMENT_EXCEPTION( + "Batch insert source is not a read function: " + + source.extension_name()); + } + return {std::move(state), read_function}; } BatchInsertInput create_batch_insert_input( @@ -392,7 +398,7 @@ BatchInsertInput create_batch_insert_input( auto source = read_function.sourceFunc(shared_state, std::move(projected_columns)); if (source) { - return {make_data_chunk_supplier(std::move(source)), Context{}}; + return {std::move(source), Context{}}; } } @@ -400,7 +406,7 @@ BatchInsertInput create_batch_insert_input( auto output = read_function.execFunc(shared_state); auto supplier = create_data_chunk_supplier(output, prop_mappings); output.tag_ids.clear(); - return {std::move(supplier), std::move(output)}; + return {make_data_chunk_source(std::move(supplier)), std::move(output)}; } std::vector match_files_with_pattern( diff --git a/src/main/query_processor.cc b/src/main/query_processor.cc index fb8d89dc7..8a503bd99 100644 --- a/src/main/query_processor.cc +++ b/src/main/query_processor.cc @@ -14,6 +14,9 @@ */ #include "neug/main/query_processor.h" + +#include + #include "neug/execution/common/context.h" #include "neug/execution/common/operators/retrieve/sink.h" #include "neug/execution/execute/plan_parser.h" @@ -25,21 +28,9 @@ namespace neug { result>> -QueryProcessor::check_and_retrieve_pipeline(const PropertyGraph& pg, - const std::string& query_string, - const std::string& user_access_mode, - int32_t num_threads) { - if (num_threads == 0) { - num_threads = max_thread_num_; - } - if (num_threads > max_thread_num_) { - num_threads = max_thread_num_; - } - if (num_threads < 1) { - RETURN_ERROR(neug::Status(neug::StatusCode::ERR_INVALID_ARGUMENT, - "Number of threads must be greater than 0")); - } - +QueryProcessor::check_and_retrieve_pipeline( + const PropertyGraph& pg, const std::string& query_string, + const std::string& user_access_mode) { auto access_mode = user_access_mode.empty() ? planner_->analyzeMode(query_string) : ParseAccessMode(user_access_mode); @@ -58,23 +49,36 @@ QueryProcessor::check_and_retrieve_pipeline(const PropertyGraph& pg, return std::make_pair(access_mode, cache_value); } +result QueryProcessor::resolve_thread_budget( + int32_t requested_threads) const { + if (requested_threads < 0 || max_thread_num_ < 1) { + RETURN_ERROR(neug::Status(neug::StatusCode::ERR_INVALID_ARGUMENT, + "Number of threads must be greater than 0")); + } + if (requested_threads == 0) { + return max_thread_num_; + } + return std::min(requested_threads, max_thread_num_); +} + result QueryProcessor::execute( const std::string& query_string, const std::string& user_access_mode, const execution::ParamsMap& parameters, int32_t num_threads) { SnapshotGuard guard(snapshot_store_); - GS_AUTO(access_mode_pipeline, check_and_retrieve_pipeline( - *guard.get().mutable_graph(), query_string, - user_access_mode, num_threads)); + GS_AUTO(thread_budget, resolve_thread_budget(num_threads)); + GS_AUTO(access_mode_pipeline, + check_and_retrieve_pipeline(*guard.get().mutable_graph(), + query_string, user_access_mode)); if (need_exclusive_lock(access_mode_pipeline.first)) { std::unique_lock lock(mutex_); return execute_internal(guard, query_string, access_mode_pipeline.second, access_mode_pipeline.first, parameters, - num_threads); + thread_budget); } else { std::shared_lock lock(mutex_); return execute_internal(guard, query_string, access_mode_pipeline.second, access_mode_pipeline.first, parameters, - num_threads); + thread_budget); } } @@ -83,9 +87,10 @@ result QueryProcessor::execute(const std::string& query_string, const rapidjson::Value& parameters, int32_t num_threads) { SnapshotGuard guard(snapshot_store_); - GS_AUTO(access_mode_pipeline, check_and_retrieve_pipeline( - *guard.get().mutable_graph(), query_string, - user_access_mode, num_threads)); + GS_AUTO(thread_budget, resolve_thread_budget(num_threads)); + GS_AUTO(access_mode_pipeline, + check_and_retrieve_pipeline(*guard.get().mutable_graph(), + query_string, user_access_mode)); const auto& param_types = access_mode_pipeline.second->params_type; execution::ParamsMap params_map; @@ -104,12 +109,12 @@ result QueryProcessor::execute(const std::string& query_string, std::unique_lock lock(mutex_); return execute_internal(guard, query_string, access_mode_pipeline.second, access_mode_pipeline.first, params_map, - num_threads); + thread_budget); } else { std::shared_lock lock(mutex_); return execute_internal(guard, query_string, access_mode_pipeline.second, access_mode_pipeline.first, params_map, - num_threads); + thread_budget); } } @@ -117,10 +122,11 @@ result QueryProcessor::execute(const std::string& query_string, result QueryProcessor::execute_internal( SnapshotGuard& guard, const std::string& query_string, std::shared_ptr cache_value, AccessMode access_mode, - const execution::ParamsMap& parameters, int32_t num_threads) { + const execution::ParamsMap& parameters, int32_t thread_budget) { auto& slot = guard.get(); auto& pg = *slot.mutable_graph(); - StorageAPUpdateInterface graph(pg, slot.mutable_view(), 0, allocator_); + StorageAPUpdateInterface graph(pg, slot.mutable_view(), 0, allocator_, + BulkLoadOptions{thread_budget}); google::protobuf::Arena arena; neug::QueryResponse* response = diff --git a/src/storages/graph/edge_table.cc b/src/storages/graph/edge_table.cc index 999cada08..b689b344a 100644 --- a/src/storages/graph/edge_table.cc +++ b/src/storages/graph/edge_table.cc @@ -772,15 +772,29 @@ std::pair EdgeTable::AddEdge( void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, const IndexerType& dst_indexer, - std::shared_ptr supplier) { - CHECK(supplier != nullptr); + std::unique_ptr source, + BulkLoadOptions options) { + CHECK(source != nullptr); + constexpr auto kMaxVertexCapacity = std::numeric_limits::max(); + if (src_indexer.capacity() > kMaxVertexCapacity) { + THROW_RUNTIME_ERROR( + "Source vertex indexer capacity exceeds addressable " + "range: " + + std::to_string(src_indexer.capacity())); + } + if (dst_indexer.capacity() > kMaxVertexCapacity) { + THROW_RUNTIME_ERROR( + "Destination vertex indexer capacity exceeds " + "addressable range: " + + std::to_string(dst_indexer.capacity())); + } const auto src_vertex_capacity = static_cast(src_indexer.capacity()); const auto dst_vertex_capacity = static_cast(dst_indexer.capacity()); - auto source = supplier->RepeatableSource(); - if (source && TryBatchBuildEdges(src_indexer, dst_indexer, source, - src_vertex_capacity, dst_vertex_capacity)) { + if (TryBatchBuildEdges(src_indexer, dst_indexer, *source, src_vertex_capacity, + dst_vertex_capacity, options)) { return; } + auto supplier = open_data_chunk_source(*source); // Keep fallback COPY paths aligned with the vertex table's actual capacity, // while leaving completely unloaded edge tables lazy until persistence. @@ -846,17 +860,19 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, } } -bool EdgeTable::TryBatchBuildEdges( - const IndexerType& src_indexer, const IndexerType& dst_indexer, - const std::shared_ptr& source, vid_t src_vertex_capacity, - vid_t dst_vertex_capacity) { +bool EdgeTable::TryBatchBuildEdges(const IndexerType& src_indexer, + const IndexerType& dst_indexer, + IDataChunkSource& source, + vid_t src_vertex_capacity, + vid_t dst_vertex_capacity, + BulkLoadOptions options) { if (!meta_ || !meta_->is_bundled() || !out_csr_ || !in_csr_ || out_csr_->edge_num() != 0 || in_csr_->edge_num() != 0 || (meta_->oe_strategy != EdgeStrategy::kNone && !meta_->oe_mutable) || (meta_->ie_strategy != EdgeStrategy::kNone && !meta_->ie_mutable)) { return false; } - const auto source_bytes = source->EstimatedBytes(); + const auto source_bytes = source.EstimatedBytes(); if (!ShouldUseBulkBuild(source_bytes)) { return false; } @@ -865,7 +881,8 @@ bool EdgeTable::TryBatchBuildEdges( staged.Init(ckp_, memory_level_); if (!internal::BundledEdgeCsrLoader::TryBuild( *staged.out_csr_, *staged.in_csr_, *meta_, src_indexer, dst_indexer, - source, source_bytes, src_vertex_capacity, dst_vertex_capacity)) { + source, source_bytes, src_vertex_capacity, dst_vertex_capacity, *ckp_, + options)) { return false; } Swap(staged); diff --git a/src/storages/graph/graph_interface.cc b/src/storages/graph/graph_interface.cc index 15b4ce0db..b274d046f 100644 --- a/src/storages/graph/graph_interface.cc +++ b/src/storages/graph/graph_interface.cc @@ -116,8 +116,9 @@ Status StorageAPUpdateInterface::DeleteEdges(label_t src_label, vid_t src, } Status StorageAPUpdateInterface::BatchAddVertices( - label_t v_label_id, std::shared_ptr supplier) { - auto status = graph_.BatchAddVertices(v_label_id, std::move(supplier)); + label_t v_label_id, std::unique_ptr source) { + auto status = graph_.BatchAddVertices(v_label_id, std::move(source), + bulk_load_options_); if (status.ok()) { mut_view_.Rebuild(graph_); } @@ -126,9 +127,9 @@ Status StorageAPUpdateInterface::BatchAddVertices( Status StorageAPUpdateInterface::BatchAddEdges( label_t src_label, label_t dst_label, label_t edge_label, - std::shared_ptr supplier) { + std::unique_ptr source) { auto status = graph_.BatchAddEdges(src_label, dst_label, edge_label, - std::move(supplier)); + std::move(source), bulk_load_options_); if (status.ok()) { mut_view_.Rebuild(graph_); } diff --git a/src/storages/graph/property_graph.cc b/src/storages/graph/property_graph.cc index 13a0220a5..3977ef989 100644 --- a/src/storages/graph/property_graph.cc +++ b/src/storages/graph/property_graph.cc @@ -134,22 +134,24 @@ Status PropertyGraph::EnsureCapacity(label_t src_label, label_t dst_label, return neug::Status::OK(); } -Status PropertyGraph::BatchAddVertices( - label_t v_label, std::shared_ptr supplier) { +Status PropertyGraph::BatchAddVertices(label_t v_label, + std::unique_ptr source, + BulkLoadOptions options) { RETURN_IF_NOT_OK(vertex_label_check(v_label)); - vertex_tables_[v_label].BatchAddVertices(std::move(supplier)); + vertex_tables_[v_label].BatchAddVertices(std::move(source), options); return neug::Status::OK(); } -Status PropertyGraph::BatchAddEdges( - label_t src_v_label, label_t dst_v_label, label_t e_label, - std::shared_ptr supplier) { +Status PropertyGraph::BatchAddEdges(label_t src_v_label, label_t dst_v_label, + label_t e_label, + std::unique_ptr source, + BulkLoadOptions options) { RETURN_IF_NOT_OK(edge_triplet_check(src_v_label, dst_v_label, e_label)); size_t index = schema_.generate_edge_label(src_v_label, dst_v_label, e_label); assert(edge_tables_.count(index) > 0); edge_tables_.at(index).BatchAddEdges( vertex_tables_.at(src_v_label).get_indexer(), - vertex_tables_.at(dst_v_label).get_indexer(), std::move(supplier)); + vertex_tables_.at(dst_v_label).get_indexer(), std::move(source), options); return neug::Status::OK(); } diff --git a/src/storages/graph/vertex_table.cc b/src/storages/graph/vertex_table.cc index c16334a39..496b00f93 100644 --- a/src/storages/graph/vertex_table.cc +++ b/src/storages/graph/vertex_table.cc @@ -48,30 +48,29 @@ void VertexTable::Init(std::shared_ptr ckp, MemoryLevel level) { v_ts_->Open(*ckp_, ModuleDescriptor{}, level); } -void VertexTable::BatchAddVertices( - std::shared_ptr supplier) { - CHECK(supplier != nullptr); - auto source = supplier->RepeatableSource(); - if (source && ShouldUseBulkBuild(source->EstimatedBytes()) && - try_batch_build_vertices(source)) { +void VertexTable::BatchAddVertices(std::unique_ptr source, + BulkLoadOptions options) { + CHECK(source != nullptr); + if (ShouldUseBulkBuild(source->EstimatedBytes()) && + try_batch_build_vertices(*source, options)) { return; } - batch_add_vertices_impl(std::move(supplier)); + batch_add_vertices_impl(open_data_chunk_source(*source)); } -bool VertexTable::try_batch_build_vertices( - const std::shared_ptr& source) { +bool VertexTable::try_batch_build_vertices(IDataChunkSource& source, + BulkLoadOptions options) { if (Size() != 0) { return false; } auto staged = VertexTable(vertex_schema_); staged.Init(ckp_, memory_level_); - const auto source_bytes = source->EstimatedBytes(); - auto options = - ResolveBulkBuildSourceOptions(source_bytes, source->ParallelEnabled(), - BulkBuildWorkerStrategy::kMaxProducers); - auto supplier = source->Open(options); + const auto source_bytes = source.EstimatedBytes(); + auto source_options = ResolveBulkBuildSourceOptions( + source_bytes, source.ParallelEnabled(), options.worker_budget, + BulkBuildWorkerStrategy::kMaxProducers); + auto supplier = open_data_chunk_source(source, source_options); if (!supplier) { return false; } diff --git a/src/storages/loader/abstract_property_graph_loader.cc b/src/storages/loader/abstract_property_graph_loader.cc index c462a66d1..179fbbf6b 100644 --- a/src/storages/loader/abstract_property_graph_loader.cc +++ b/src/storages/loader/abstract_property_graph_loader.cc @@ -29,7 +29,8 @@ void AbstractPropertyGraphLoader::addVerticesToVertexTable( auto supplier = createVertexChunkSupplier(v_label_id, label_name, v_file, pk_type, pk_name, pk_ind, loading_config_, 0); - graph_.BatchAddVertices(v_label_id, supplier); + graph_.BatchAddVertices(v_label_id, + make_data_chunk_source(std::move(supplier))); } } @@ -115,7 +116,8 @@ void AbstractPropertyGraphLoader::addEdgesToEdgeTable( << " from file: " << e_file; auto supplier = createEdgeChunkSupplier( src_label_id, dst_label_id, e_label_id, e_file, loading_config_, 0); - graph_.BatchAddEdges(src_label_id, dst_label_id, e_label_id, supplier); + graph_.BatchAddEdges(src_label_id, dst_label_id, e_label_id, + make_data_chunk_source(std::move(supplier))); } } diff --git a/src/storages/loader/bundled_edge_csr_loader.cc b/src/storages/loader/bundled_edge_csr_loader.cc index 4e86f3066..218ced9fa 100644 --- a/src/storages/loader/bundled_edge_csr_loader.cc +++ b/src/storages/loader/bundled_edge_csr_loader.cc @@ -20,17 +20,24 @@ #include #include #include +#include +#include #include #include +#include #include +#include +#include #include #include #include "neug/common/columns/value_columns.h" #include "neug/common/types/container_types.h" +#include "neug/storages/checkpoint.h" #include "neug/storages/csr/mutable_csr.h" #include "neug/storages/loader/chunk_pipeline_utils.h" #include "neug/storages/loader/loader_utils.h" +#include "neug/utils/exception/exception.h" #include "neug/utils/property/types.h" namespace neug { @@ -290,7 +297,9 @@ class BulkEdgeDataReader { explicit BulkEdgeDataReader(const std::shared_ptr& column) : column_(column.get()), value_column_(dynamic_cast*>(column.get())) { - CHECK(column_ != nullptr); + if (column_ == nullptr) { + THROW_SCHEMA_MISMATCH("Bundled edge property column is missing"); + } } EDATA_T Get(size_t row) const { @@ -314,6 +323,162 @@ class BulkEdgeDataReader { EmptyType Get(size_t /*row*/) const { return EmptyType(); } }; +template +struct BulkEdgeSpillRecord { + vid_t src; + vid_t dst; + EDATA_T data; +}; + +template +class BulkEdgeSpillSegment { + public: + using record_t = BulkEdgeSpillRecord; + + // Retained across ordered runs so each spill segment opens its file and + // allocates its replay buffer once. Adjacent ranges continue reading from + // the current position without another seek. + class Reader { + public: + Reader(std::string path, uint64_t record_count) + : path_(std::move(path)), + record_count_(record_count), + input_(path_, std::ios::binary), + records_(kRecordsPerBuffer) { + if (!input_) { + THROW_IO_EXCEPTION("Failed to open edge spill file: " + path_); + } + } + + Reader(const Reader&) = delete; + Reader& operator=(const Reader&) = delete; + + template + void ReplayRange(uint64_t offset, uint64_t count, Consume&& consume) { + CHECK_LE(offset, record_count_); + CHECK_LE(count, record_count_ - offset); + if (next_record_offset_ != offset) { + Seek(offset); + } + + auto remaining = count; + while (remaining > 0) { + const auto requested = + static_cast(std::min(remaining, records_.size())); + input_.read(reinterpret_cast(records_.data()), + static_cast(requested * sizeof(record_t))); + const auto bytes = input_.gcount(); + if (bytes != + static_cast(requested * sizeof(record_t))) { + THROW_IO_EXCEPTION("Truncated edge spill file: " + path_); + } + next_record_offset_ += requested; + consume(records_.data(), requested); + remaining -= requested; + } + } + + private: + void Seek(uint64_t offset) { + if (offset > + static_cast(std::numeric_limits::max()) / + sizeof(record_t)) { + THROW_IO_EXCEPTION("Edge spill offset exceeds stream capacity: " + + path_); + } + input_.clear(); + input_.seekg(static_cast(offset * sizeof(record_t))); + if (!input_) { + THROW_IO_EXCEPTION("Failed to seek edge spill file: " + path_); + } + next_record_offset_ = offset; + } + + std::string path_; + uint64_t record_count_; + std::ifstream input_; + std::vector records_; + uint64_t next_record_offset_ = std::numeric_limits::max(); + }; + + explicit BulkEdgeSpillSegment(Checkpoint& checkpoint) + : file_(checkpoint.CreateRuntimeFile()), + output_(file_.path(), std::ios::binary | std::ios::trunc) { + static_assert(std::is_trivially_copyable_v); + if (!output_) { + THROW_IO_EXCEPTION("Failed to create edge spill file: " + file_.path()); + } + buffer_.reserve(kRecordsPerBuffer); + } + + ~BulkEdgeSpillSegment() = default; + BulkEdgeSpillSegment(const BulkEdgeSpillSegment&) = delete; + BulkEdgeSpillSegment& operator=(const BulkEdgeSpillSegment&) = delete; + + void Append(vid_t src, vid_t dst, const EDATA_T& data) { + CHECK(!finalized_); + CHECK_LT(record_count_, std::numeric_limits::max()); + buffer_.push_back({src, dst, data}); + ++record_count_; + if (buffer_.size() == kRecordsPerBuffer) { + Flush(); + } + } + + uint64_t RecordCount() const { return record_count_; } + + void Finalize() { + if (finalized_) { + return; + } + Flush(); + output_.close(); + if (output_.fail()) { + THROW_IO_EXCEPTION("Failed to finalize edge spill file: " + file_.path()); + } + finalized_ = true; + } + + template + void Replay(Consume&& consume) const { + auto reader = OpenReader(); + reader->ReplayRange(0, record_count_, std::forward(consume)); + } + + std::unique_ptr OpenReader() const { + CHECK(finalized_); + return std::make_unique(file_.path(), record_count_); + } + + private: + void Flush() { + if (buffer_.empty()) { + return; + } + output_.write( + reinterpret_cast(buffer_.data()), + static_cast(buffer_.size() * sizeof(record_t))); + if (!output_) { + THROW_IO_EXCEPTION("Failed to write edge spill file: " + file_.path()); + } + buffer_.clear(); + } + + static constexpr size_t kRecordsPerBuffer = 4096; + CheckpointFileManager::RuntimeFileHandle file_; + std::ofstream output_; + std::vector buffer_; + uint64_t record_count_ = 0; + bool finalized_ = false; +}; + +struct BulkEdgeSpillRun { + size_t segment = 0; + uint64_t record_offset = 0; + uint64_t record_count = 0; + uint64_t first_row_ordinal = 0; +}; + constexpr uint32_t kInvalidBulkEdgeIndex = std::numeric_limits::max(); struct BulkEdgeGroup { @@ -455,33 +620,65 @@ void count_bulk_edge_chunk(BulkEdgeWorkerScratch& scratch, OutWriter& out, } } -template -void count_bulk_edges(const std::shared_ptr& source, - const IndexerType& src_indexer, - const IndexerType& dst_indexer, - ChunkSourceOptions options, - std::vector& scratches, - OutWriter& out, InWriter& in) { - // Degree accumulation and single-slot uniqueness checks are commutative, so - // this pass never needs input order even when fill may later fall back to the - // ordered last-write-wins path. - // The degree pass only needs endpoint OIDs, so edge properties are not - // parsed, typed, allocated, and discarded during the first pass. - options.projected_columns = {0, 1}; - auto supplier = source->Open(options); +template +struct BulkEdgeMaterialization { + std::vector>> segments; + std::vector> runs; +}; + +template +BulkEdgeMaterialization materialize_bulk_edges( + IDataChunkSource& source, const IndexerType& src_indexer, + const IndexerType& dst_indexer, const ChunkSourceOptions& options, + Checkpoint& checkpoint, std::vector& scratches, + OutWriter& out, InWriter& in) { + auto supplier = open_data_chunk_source(source, options); CHECK(supplier != nullptr); + if (options.consumer_count > 1 && source.ProvidesStableRowOrdinals()) { + CHECK(supplier->ProvidesStableRowOrdinals()); + } - const auto worker_count = options.consumer_count; + const auto worker_count = std::max(1, options.consumer_count); CHECK_EQ(scratches.size(), static_cast(worker_count)); + BulkEdgeMaterialization result; + result.segments.reserve(static_cast(worker_count)); + result.runs.resize(static_cast(worker_count)); + for (int32_t worker = 0; worker < worker_count; ++worker) { + result.segments.push_back( + std::make_unique>(checkpoint)); + } + const bool concurrent = worker_count > 1; - auto count = [&](int32_t worker, const std::shared_ptr& chunk) { + auto materialize = [&](int32_t worker, const SequencedDataChunk& sequenced) { CHECK_GE(worker, 0); CHECK_LT(worker, worker_count); + const auto& chunk = sequenced.chunk; auto& scratch = scratches[static_cast(worker)]; index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); count_bulk_edge_chunk(scratch, out, in, concurrent); + const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; + BulkEdgeDataReader data_reader(data_column); + auto& segment = *result.segments[static_cast(worker)]; + const auto record_offset = segment.RecordCount(); + for (size_t row = 0; row < scratch.src_lids.size(); ++row) { + const auto src = scratch.src_lids[row]; + const auto dst = scratch.dst_lids[row]; + if (is_valid_bulk_edge(src, dst)) { + segment.Append(src, dst, data_reader.Get(row)); + } + } + const auto record_count = segment.RecordCount() - record_offset; + if (record_count > 0) { + result.runs[static_cast(worker)].push_back( + {static_cast(worker), record_offset, record_count, + sequenced.first_row_ordinal}); + } }; - consume_supplier_indexed(*supplier, options, count); + consume_supplier_indexed(*supplier, options, materialize); + for (auto& segment : result.segments) { + segment->Finalize(); + } + return result; } void append_bulk_edge_group(flat_hash_map& groups, @@ -500,10 +697,10 @@ void append_bulk_edge_group(flat_hash_map& groups, ++group.count; } -template +template void fill_bulk_edge_chunk_serial(BulkEdgeWorkerScratch& scratch, - const BulkEdgeDataReader& data_reader, - OutWriter& out, InWriter& in) { + const DataReader& data_reader, OutWriter& out, + InWriter& in) { size_t valid_edges = 0; for (size_t row = 0; row < scratch.src_lids.size(); ++row) { const auto src = scratch.src_lids[row]; @@ -525,11 +722,10 @@ void fill_bulk_edge_chunk_serial(BulkEdgeWorkerScratch& scratch, scratch.filled_edge_count += static_cast(valid_edges); } -template -void fill_bulk_edge_chunk_concurrent( - BulkEdgeWorkerScratch& scratch, - const BulkEdgeDataReader& data_reader, OutWriter& out, - InWriter& in) { +template +void fill_bulk_edge_chunk_concurrent(BulkEdgeWorkerScratch& scratch, + const DataReader& data_reader, + OutWriter& out, InWriter& in) { CHECK_LE(scratch.src_lids.size(), static_cast(std::numeric_limits::max())); constexpr bool kDirectOut = OutWriter::kStrategy == EdgeStrategy::kSingle; @@ -608,40 +804,110 @@ void fill_bulk_edge_chunk_concurrent( scratch.filled_edge_count += static_cast(valid_edges); } -template -uint64_t fill_bulk_edges(const std::shared_ptr& source, - const IndexerType& src_indexer, - const IndexerType& dst_indexer, - const ChunkSourceOptions& options, - std::vector& scratches, - OutWriter& out, InWriter& in) { - auto supplier = source->Open(options); - CHECK(supplier != nullptr); +template +class BulkEdgeSpillDataReader { + public: + explicit BulkEdgeSpillDataReader(const BulkEdgeSpillRecord* records) + : records_(records) {} + + EDATA_T Get(size_t row) const { return records_[row].data; } - const auto worker_count = options.consumer_count; - CHECK_GE(scratches.size(), static_cast(worker_count)); + private: + const BulkEdgeSpillRecord* records_; +}; + +template +uint64_t replay_bulk_edges( + const BulkEdgeMaterialization& materialization, bool concurrent, + bool preserve_input_order, std::vector& scratches, + OutWriter& out, InWriter& in) { + const auto& segments = materialization.segments; + CHECK_EQ(scratches.size(), segments.size()); for (auto& scratch : scratches) { scratch.filled_edge_count = 0; } - if (!options.preserve_order && worker_count > 1) { - auto fill = [&](int32_t worker, const std::shared_ptr& chunk) { - CHECK_GE(worker, 0); - CHECK_LT(worker, worker_count); - auto& scratch = scratches[static_cast(worker)]; - index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); - const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; - BulkEdgeDataReader data_reader(data_column); - fill_bulk_edge_chunk_concurrent(scratch, data_reader, out, in); - }; - - consume_supplier_indexed(*supplier, options, fill); + + auto replay_segment = [&](size_t worker) { + auto& scratch = scratches[worker]; + segments[worker]->Replay( + [&](const BulkEdgeSpillRecord* records, size_t count) { + scratch.src_lids.resize(count); + scratch.dst_lids.resize(count); + for (size_t row = 0; row < count; ++row) { + scratch.src_lids[row] = records[row].src; + scratch.dst_lids[row] = records[row].dst; + } + BulkEdgeSpillDataReader data_reader(records); + if (concurrent) { + fill_bulk_edge_chunk_concurrent(scratch, data_reader, out, in); + } else { + fill_bulk_edge_chunk_serial(scratch, data_reader, out, in); + } + }); + }; + + if (preserve_input_order) { + std::vector ordered_runs; + for (const auto& runs : materialization.runs) { + ordered_runs.insert(ordered_runs.end(), runs.begin(), runs.end()); + } + std::stable_sort(ordered_runs.begin(), ordered_runs.end(), + [](const auto& lhs, const auto& rhs) { + return lhs.first_row_ordinal < rhs.first_row_ordinal; + }); + using spill_reader_t = typename BulkEdgeSpillSegment::Reader; + std::vector> readers(segments.size()); + for (const auto& run : ordered_runs) { + CHECK_LT(run.segment, segments.size()); + auto& scratch = scratches[run.segment]; + auto& reader = readers[run.segment]; + if (!reader) { + reader = segments[run.segment]->OpenReader(); + } + reader->ReplayRange( + run.record_offset, run.record_count, + [&](const BulkEdgeSpillRecord* records, size_t count) { + scratch.src_lids.resize(count); + scratch.dst_lids.resize(count); + for (size_t row = 0; row < count; ++row) { + scratch.src_lids[row] = records[row].src; + scratch.dst_lids[row] = records[row].dst; + } + BulkEdgeSpillDataReader data_reader(records); + fill_bulk_edge_chunk_serial(scratch, data_reader, out, in); + }); + } + } else if (concurrent && segments.size() > 1) { + std::atomic cancelled{false}; + std::mutex error_mutex; + std::exception_ptr first_error; + std::vector workers; + workers.reserve(segments.size()); + for (size_t worker = 0; worker < segments.size(); ++worker) { + workers.emplace_back([&, worker] { + try { + if (!cancelled.load(std::memory_order_acquire)) { + replay_segment(worker); + } + } catch (...) { + bool expected = false; + if (cancelled.compare_exchange_strong(expected, true, + std::memory_order_acq_rel)) { + std::lock_guard lock(error_mutex); + first_error = std::current_exception(); + } + } + }); + } + for (auto& worker : workers) { + worker.join(); + } + if (first_error) { + std::rethrow_exception(first_error); + } } else { - auto& scratch = scratches.front(); - while (auto chunk = supplier->GetNextChunk()) { - index_bulk_edge_endpoints(chunk, src_indexer, dst_indexer, scratch); - const auto data_column = chunk->col_num() > 2 ? chunk->get(2) : nullptr; - BulkEdgeDataReader data_reader(data_column); - fill_bulk_edge_chunk_serial(scratch, data_reader, out, in); + for (size_t worker = 0; worker < segments.size(); ++worker) { + replay_segment(worker); } } @@ -655,64 +921,66 @@ uint64_t fill_bulk_edges(const std::shared_ptr& source, } template -uint64_t build_bundled_edges_with_writers( - OutWriter& out, InWriter& in, const IndexerType& src_indexer, - const IndexerType& dst_indexer, - const std::shared_ptr& source, int64_t source_bytes) { +uint64_t build_bundled_edges_with_writers(OutWriter& out, InWriter& in, + const IndexerType& src_indexer, + const IndexerType& dst_indexer, + IDataChunkSource& source, + int64_t source_bytes, + Checkpoint& checkpoint, + BulkLoadOptions bulk_load_options) { constexpr bool kStoresAnyDirection = OutWriter::kStrategy != EdgeStrategy::kNone || InWriter::kStrategy != EdgeStrategy::kNone; if constexpr (!kStoresAnyDirection) { return 0; } else { - // VertexTable owns its reserve policy. Edge storage consumes the resulting - // indexer capacity instead of duplicating PropertyGraph::Dump's policy. - const auto count_options = ResolveBulkBuildSourceOptions( - source_bytes, source->ParallelEnabled(), - BulkBuildWorkerStrategy::kBalancedProducerConsumer); - constexpr bool kNeedsDegreeCount = - OutWriter::kStrategy == EdgeStrategy::kMultiple || - InWriter::kStrategy == EdgeStrategy::kMultiple; - std::vector scratches; - if constexpr (kNeedsDegreeCount) { - scratches.resize(static_cast(count_options.consumer_count)); - count_bulk_edges(source, src_indexer, dst_indexer, count_options, - scratches, out, in); - } - constexpr bool kChecksOutSingle = - OutWriter::kStrategy == EdgeStrategy::kSingle; - constexpr bool kChecksInSingle = + constexpr bool kHasSingleDirection = + OutWriter::kStrategy == EdgeStrategy::kSingle || InWriter::kStrategy == EdgeStrategy::kSingle; - const bool single_duplicate = std::any_of( - scratches.begin(), scratches.end(), [](const auto& scratch) { - return (kChecksOutSingle && scratch.out_single_duplicate) || - (kChecksInSingle && scratch.in_single_duplicate); - }); - const bool preserve_fill_order = (kChecksOutSingle || kChecksInSingle) && - (!kNeedsDegreeCount || single_duplicate); - const auto fill_options = - preserve_fill_order ? ChunkSourceOptions{} : count_options; - if (scratches.size() < static_cast(fill_options.consumer_count)) { - scratches.resize(static_cast(fill_options.consumer_count)); + auto source_options = ResolveBulkBuildSourceOptions( + source_bytes, source.ParallelEnabled(), bulk_load_options.worker_budget, + BulkBuildWorkerStrategy::kBalancedProducerConsumer); + if constexpr (kHasSingleDirection) { + if (!source.ProvidesStableRowOrdinals()) { + // Without stable ordinals a one-shot parallel source cannot recover + // last-write-wins order if a single-edge duplicate is discovered. + source_options = ChunkSourceOptions{}; + source_options.worker_budget = + std::max(1, bulk_load_options.worker_budget); + } } + std::vector scratches; + scratches.resize(static_cast( + std::max(1, source_options.consumer_count))); + auto materialization = materialize_bulk_edges( + source, src_indexer, dst_indexer, source_options, checkpoint, scratches, + out, in); if constexpr (OutWriter::kStrategy == EdgeStrategy::kMultiple) { out.AllocateFromCounts(); } if constexpr (InWriter::kStrategy == EdgeStrategy::kMultiple) { in.AllocateFromCounts(); } - return fill_bulk_edges(source, src_indexer, dst_indexer, - fill_options, scratches, out, in); + const bool single_duplicate = + kHasSingleDirection && + std::any_of(scratches.begin(), scratches.end(), + [](const auto& scratch) { + return scratch.out_single_duplicate || + scratch.in_single_duplicate; + }); + const bool concurrent_replay = + source_options.consumer_count > 1 && !single_duplicate; + return replay_bulk_edges(materialization, concurrent_replay, + single_duplicate, scratches, out, in); } } template -bool build_bundled_edges_typed(CsrBase* out_csr, CsrBase* in_csr, - const IndexerType& src_indexer, - const IndexerType& dst_indexer, - const std::shared_ptr& source, - int64_t source_bytes, vid_t src_vertex_capacity, - vid_t dst_vertex_capacity) { +bool build_bundled_edges_typed( + CsrBase* out_csr, CsrBase* in_csr, const IndexerType& src_indexer, + const IndexerType& dst_indexer, IDataChunkSource& source, + int64_t source_bytes, vid_t src_vertex_capacity, vid_t dst_vertex_capacity, + Checkpoint& checkpoint, BulkLoadOptions bulk_load_options) { bool built = false; const bool out_supported = with_csr_bulk_writer(out_csr, [&](auto& out) { @@ -722,7 +990,8 @@ bool build_bundled_edges_typed(CsrBase* out_csr, CsrBase* in_csr, in.PrepareBuild(dst_vertex_capacity); const auto filled_edge_count = build_bundled_edges_with_writers( - out, in, src_indexer, dst_indexer, source, source_bytes); + out, in, src_indexer, dst_indexer, source, source_bytes, + checkpoint, bulk_load_options); out.Finish(filled_edge_count); in.Finish(filled_edge_count); }); @@ -735,9 +1004,10 @@ bool build_bundled_edges(CsrBase* out_csr, CsrBase* in_csr, const EdgeSchema& schema, const IndexerType& src_indexer, const IndexerType& dst_indexer, - const std::shared_ptr& source, - int64_t source_bytes, vid_t src_vertex_capacity, - vid_t dst_vertex_capacity) { + IDataChunkSource& source, int64_t source_bytes, + vid_t src_vertex_capacity, vid_t dst_vertex_capacity, + Checkpoint& checkpoint, + BulkLoadOptions bulk_load_options) { const auto property_type = schema.properties.empty() ? DataTypeId::kEmpty : schema.properties[0].id(); @@ -746,13 +1016,15 @@ bool build_bundled_edges(CsrBase* out_csr, CsrBase* in_csr, case DataTypeId::enum_val: \ return build_bundled_edges_typed( \ out_csr, in_csr, src_indexer, dst_indexer, source, source_bytes, \ - src_vertex_capacity, dst_vertex_capacity); + src_vertex_capacity, dst_vertex_capacity, checkpoint, \ + bulk_load_options); FOR_EACH_DATA_TYPE_NO_STRING(TYPE_DISPATCHER) #undef TYPE_DISPATCHER case DataTypeId::kEmpty: return build_bundled_edges_typed( out_csr, in_csr, src_indexer, dst_indexer, source, source_bytes, - src_vertex_capacity, dst_vertex_capacity); + src_vertex_capacity, dst_vertex_capacity, checkpoint, + bulk_load_options); default: return false; } @@ -763,11 +1035,12 @@ bool build_bundled_edges(CsrBase* out_csr, CsrBase* in_csr, bool internal::BundledEdgeCsrLoader::TryBuild( CsrBase& out_csr, CsrBase& in_csr, const EdgeSchema& schema, const IndexerType& src_indexer, const IndexerType& dst_indexer, - const std::shared_ptr& source, int64_t source_bytes, - vid_t src_vertex_capacity, vid_t dst_vertex_capacity) { - return build_bundled_edges(&out_csr, &in_csr, schema, src_indexer, - dst_indexer, source, source_bytes, - src_vertex_capacity, dst_vertex_capacity); + IDataChunkSource& source, int64_t source_bytes, vid_t src_vertex_capacity, + vid_t dst_vertex_capacity, Checkpoint& checkpoint, + BulkLoadOptions options) { + return build_bundled_edges( + &out_csr, &in_csr, schema, src_indexer, dst_indexer, source, source_bytes, + src_vertex_capacity, dst_vertex_capacity, checkpoint, options); } } // namespace neug diff --git a/src/storages/loader/bundled_edge_csr_loader.h b/src/storages/loader/bundled_edge_csr_loader.h index 495590a29..7971262fb 100644 --- a/src/storages/loader/bundled_edge_csr_loader.h +++ b/src/storages/loader/bundled_edge_csr_loader.h @@ -24,13 +24,15 @@ namespace neug { class CsrBase; +class Checkpoint; class IDataChunkSource; +struct BulkLoadOptions; namespace internal { -/// Builds the outgoing and incoming CSR pair for a bundled edge table directly -/// from a repeatable chunk source. EdgeTable owns staging and publication; this -/// class owns source planning and the CSR bulk-build protocol. +/// Builds the outgoing and incoming CSR pair for a bundled edge table from a +/// one-shot chunk input. EdgeTable owns staging and publication; this class +/// owns source planning, the spill intermediate, and the CSR build protocol. class BundledEdgeCsrLoader { public: template @@ -43,10 +45,10 @@ class BundledEdgeCsrLoader { /// Callers must pass fresh, unpublished CSR instances. static bool TryBuild(CsrBase& out_csr, CsrBase& in_csr, const EdgeSchema& schema, const IndexerType& src_indexer, - const IndexerType& dst_indexer, - const std::shared_ptr& source, + const IndexerType& dst_indexer, IDataChunkSource& source, int64_t source_bytes, vid_t src_vertex_capacity, - vid_t dst_vertex_capacity); + vid_t dst_vertex_capacity, Checkpoint& checkpoint, + BulkLoadOptions options); }; } // namespace internal diff --git a/src/storages/loader/loader_utils.cc b/src/storages/loader/loader_utils.cc index 095f0c0ae..61917eb88 100644 --- a/src/storages/loader/loader_utils.cc +++ b/src/storages/loader/loader_utils.cc @@ -31,6 +31,7 @@ #include #include #include +#include #include #include #include @@ -68,6 +69,7 @@ struct CsvRangeTask { std::string file_path; CsvPartitionRange range; int64_t skip_rows = 0; + uint64_t first_row_ordinal = 0; }; struct CsvPartitionPlan { @@ -78,7 +80,7 @@ struct CsvPartitionPlan { struct CsvPartitionPlanCache { std::shared_ptr GetOrCreate( const std::vector& file_paths, const CsvReadConfig& config, - int32_t producer_count); + int32_t producer_count, int32_t worker_budget); private: struct Entry { @@ -87,7 +89,7 @@ struct CsvPartitionPlanCache { }; std::mutex mutex_; - std::unordered_map> entries_; + std::map, std::unique_ptr> entries_; }; namespace { @@ -460,6 +462,8 @@ struct CsvScanResult { std::vector ranges; }; +constexpr size_t kMinCsvScanBytesPerWorker = 4 << 20; // 4 MB + /// Scans CSV bytes to count rows and optionally produce record-aligned ranges. /// csv-parser always recognizes RFC-style doubled quotes, so the scanner uses /// the same effective dialect instead of the legacy double_quote option. @@ -482,14 +486,13 @@ class CsvFileScanner { if (file_size == 0) return 0; - constexpr size_t kMinChunkSize = 4 << 20; // 4 MB unsigned num_threads = use_threads_ ? std::thread::hardware_concurrency() : 1; if (num_threads == 0) num_threads = 1; - if (file_size < kMinChunkSize * num_threads) { - num_threads = - std::max(1u, static_cast(file_size / kMinChunkSize)); + if (file_size < kMinCsvScanBytesPerWorker * num_threads) { + num_threads = std::max( + 1u, static_cast(file_size / kMinCsvScanBytesPerWorker)); } if (num_threads <= 1) return count_single(file_size); @@ -505,8 +508,11 @@ class CsvFileScanner { const auto target_partitions = std::min( file_size, static_cast(std::max(1, requested_partitions))); - const auto workers = static_cast(std::min( - file_size, static_cast(std::max(1, scan_threads)))); + const auto useful_workers = + std::max(1, file_size / kMinCsvScanBytesPerWorker); + const auto workers = static_cast( + std::min({file_size, useful_workers, + static_cast(std::max(1, scan_threads))})); if (workers == 1) { result.row_count = count_single(file_size); result.ranges.push_back({0, file_size, 0}); @@ -863,7 +869,7 @@ class CsvRangeStream final : public std::istream { std::shared_ptr build_csv_partition_plan( const std::vector& file_paths, const CsvReadConfig& config, - int32_t producer_count) { + int32_t producer_count, int32_t worker_budget) { CHECK_GE(producer_count, 1); auto plan = std::make_shared(); @@ -915,7 +921,7 @@ std::shared_ptr build_csv_partition_plan( // concurrently. With fewer files, the same H budget is divided among // intra-file speculative scans. scan_parallel() uses its caller for one // byte range, therefore the number of active scanners never exceeds H. - const auto scan_budget = chunk_pipeline_detail::hardware_worker_count(); + const auto scan_budget = std::max(1, worker_budget); auto scan_threads = range_counts; if (non_empty_files < static_cast(scan_budget)) { size_t assigned_scanners = assigned_ranges; @@ -1008,6 +1014,7 @@ std::shared_ptr build_csv_partition_plan( } int64_t total = 0; + uint64_t next_row_ordinal = 0; for (size_t i = 0; i < file_paths.size(); ++i) { const auto& scan = scans[i]; if (total != kUnknownRowNum) { @@ -1033,7 +1040,13 @@ std::shared_ptr build_csv_partition_plan( const auto range_rows = range_end_row - range.start_row; const auto range_skip = std::min(remaining_skip, range_rows); remaining_skip -= range_skip; - plan->tasks.push_back({file_paths[i], range, range_skip}); + plan->tasks.push_back( + {file_paths[i], range, range_skip, next_row_ordinal}); + const auto output_rows = range_rows - range_skip; + CHECK_GE(output_rows, 0); + CHECK_LE(static_cast(output_rows), + std::numeric_limits::max() - next_row_ordinal); + next_row_ordinal += static_cast(output_rows); } } plan->row_count = total; @@ -1048,20 +1061,22 @@ std::shared_ptr build_csv_partition_plan( std::shared_ptr CsvPartitionPlanCache::GetOrCreate( const std::vector& file_paths, const CsvReadConfig& config, - int32_t producer_count) { - producer_count = std::clamp( - producer_count, 1, chunk_pipeline_detail::hardware_worker_count()); + int32_t producer_count, int32_t worker_budget) { + worker_budget = std::max(1, worker_budget); + producer_count = std::clamp(producer_count, 1, + std::max(1, worker_budget - 1)); Entry* entry; { std::lock_guard lock(mutex_); - auto& cached = entries_[producer_count]; + auto& cached = entries_[{producer_count, worker_budget}]; if (!cached) { cached = std::make_unique(); } entry = cached.get(); } std::call_once(entry->once, [&] { - entry->plan = build_csv_partition_plan(file_paths, config, producer_count); + entry->plan = build_csv_partition_plan(file_paths, config, producer_count, + worker_budget); }); return entry->plan; } @@ -1558,58 +1573,53 @@ int64_t CSVChunkSupplier::RowNum() const { namespace { -class SourceBackedChunkSupplier final : public IDataChunkSupplier { +class SupplierChunkSource final : public IDataChunkSource { public: - explicit SourceBackedChunkSupplier(std::shared_ptr source) - : source_(std::move(source)) { - CHECK(source_ != nullptr); - } + explicit SupplierChunkSource(std::shared_ptr supplier) + : supplier_(std::move(supplier)) {} - std::shared_ptr GetNextChunk() override { - return OpenSupplier()->GetNextChunk(); + std::shared_ptr Open( + const ChunkSourceOptions& /*options*/) override { + if (!supplier_) { + THROW_INVALID_ARGUMENT_EXCEPTION("Data chunk source is empty"); + } + return std::move(supplier_); } - int64_t RowNum() const override { return OpenSupplier()->RowNum(); } + bool ParallelEnabled() const override { return false; } - bool SupportsConcurrentGetNext() const override { - return OpenSupplier()->SupportsConcurrentGetNext(); - } + private: + std::shared_ptr supplier_; +}; - void Cancel() override { - std::shared_ptr supplier; - { - std::lock_guard lock(mutex_); - cancelled_ = true; - supplier = supplier_; - } - if (supplier) { - supplier->Cancel(); +class SharedChunkSource final : public IDataChunkSource { + public: + explicit SharedChunkSource(std::shared_ptr source) + : source_(std::move(source)) {} + + std::shared_ptr Open( + const ChunkSourceOptions& options) override { + if (!source_) { + THROW_INVALID_ARGUMENT_EXCEPTION("Data chunk source is empty"); } + auto source = std::move(source_); + return source->Open(options); } - std::shared_ptr RepeatableSource() const override { - return source_; + int64_t EstimatedBytes() const override { + return source_ ? source_->EstimatedBytes() : -1; } - private: - IDataChunkSupplier* OpenSupplier() const { - std::lock_guard lock(mutex_); - if (!supplier_) { - supplier_ = source_->Open(); - if (!supplier_) { - THROW_INTERNAL_EXCEPTION("Data source returned a null supplier"); - } - if (cancelled_) { - supplier_->Cancel(); - } - } - return supplier_.get(); + bool ParallelEnabled() const override { + return source_ && source_->ParallelEnabled(); } + bool ProvidesStableRowOrdinals() const override { + return source_ && source_->ProvidesStableRowOrdinals(); + } + + private: std::shared_ptr source_; - mutable std::mutex mutex_; - mutable std::shared_ptr supplier_; - mutable bool cancelled_ = false; }; CsvReadConfig project_csv_config(const CsvReadConfig& config, @@ -1663,7 +1673,8 @@ class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { config_(std::move(config)), producer_count_(std::clamp( options.producer_count, 1, - chunk_pipeline_detail::hardware_worker_count())), + std::max(1, options.worker_budget - 1))), + worker_budget_(std::max(1, options.worker_budget)), plan_cache_(std::move(plan_cache)), queue_(options.queue_capacity) {} @@ -1673,14 +1684,18 @@ class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { } std::shared_ptr GetNextChunk() override { + return GetNextChunkWithOrdinal().chunk; + } + + SequencedDataChunk GetNextChunkWithOrdinal() override { EnsureScanned(); StartWorkers(); - std::shared_ptr chunk; + SequencedDataChunk chunk; if (queue_.Pop(chunk)) { return chunk; } RethrowError(); - return nullptr; + return {}; } int64_t RowNum() const override { @@ -1689,6 +1704,7 @@ class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { } bool SupportsConcurrentGetNext() const override { return true; } + bool ProvidesStableRowOrdinals() const override { return true; } void Cancel() override { stop_.store(true, std::memory_order_release); @@ -1699,7 +1715,8 @@ class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { void EnsureScanned() { std::call_once(scan_once_, [&] { CHECK(plan_cache_ != nullptr); - plan_ = plan_cache_->GetOrCreate(file_paths_, config_, producer_count_); + plan_ = plan_cache_->GetOrCreate(file_paths_, config_, producer_count_, + worker_budget_); CHECK(plan_ != nullptr); }); } @@ -1734,11 +1751,19 @@ class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { range_config.use_threads = false; CsvSupplierRuntime runtime(task.file_path, range_config, CsvRowCountMode::kUnknown, task.range); + auto next_row_ordinal = task.first_row_ordinal; while (!stop_.load(std::memory_order_acquire)) { auto chunk = runtime.get_next_chunk(); - if (!chunk || !queue_.Push(std::move(chunk))) { + if (!chunk) { break; } + const auto row_count = chunk->row_num(); + if (!queue_.Push({std::move(chunk), next_row_ordinal})) { + break; + } + CHECK_LE(row_count, + std::numeric_limits::max() - next_row_ordinal); + next_row_ordinal += static_cast(row_count); } } } catch (...) { SetError(std::current_exception()); } @@ -1784,9 +1809,10 @@ class PartitionedCsvChunkSupplier final : public IDataChunkSupplier { std::vector file_paths_; CsvReadConfig config_; int32_t producer_count_; + int32_t worker_budget_; std::shared_ptr plan_cache_; std::shared_ptr plan_; - chunk_pipeline_detail::BoundedQueue> queue_; + chunk_pipeline_detail::BoundedQueue queue_; std::once_flag scan_once_; std::once_flag workers_once_; std::vector workers_; @@ -1847,23 +1873,65 @@ class ChainedCsvChunkSupplier final : public IDataChunkSupplier { } // namespace -std::shared_ptr make_data_chunk_supplier( +std::unique_ptr make_data_chunk_source( + std::shared_ptr supplier) { + if (!supplier) { + return nullptr; + } + return std::make_unique(std::move(supplier)); +} + +std::unique_ptr make_data_chunk_source( std::shared_ptr source) { if (!source) { return nullptr; } - return std::make_shared(std::move(source)); + return std::make_unique(std::move(source)); +} + +std::shared_ptr open_data_chunk_source( + IDataChunkSource& source, const ChunkSourceOptions& options) { + auto supplier = source.Open(NormalizeChunkSourceOptions(options)); + if (!supplier) { + THROW_INTERNAL_EXCEPTION("Data source returned a null supplier"); + } + return supplier; +} + +ChunkSourceOptions NormalizeChunkSourceOptions(ChunkSourceOptions options) { + options.worker_budget = std::max(1, options.worker_budget); + options.producer_count = std::max(0, options.producer_count); + options.consumer_count = std::max(1, options.consumer_count); + options.queue_capacity = std::max(1, options.queue_capacity); + + if (options.preserve_order || options.worker_budget == 1) { + options.producer_count = 0; + options.consumer_count = 1; + } else { + options.producer_count = + std::min(options.producer_count, options.worker_budget - 1); + options.consumer_count = std::min( + options.consumer_count, options.worker_budget - options.producer_count); + } + + CHECK_GE(options.worker_budget, 1); + CHECK_GE(options.producer_count, 0); + CHECK_GE(options.consumer_count, 1); + CHECK_LE(options.producer_count + options.consumer_count, + options.worker_budget); + return options; } ChunkSourceOptions ResolveBulkBuildSourceOptions( - int64_t source_bytes, bool parallel_enabled, + int64_t source_bytes, bool parallel_enabled, int32_t worker_budget, BulkBuildWorkerStrategy worker_strategy) { constexpr int64_t kMinPartitionBytes = 64LL * 1024 * 1024; constexpr size_t kMaxQueuedChunks = 64; ChunkSourceOptions options; + const auto workers = std::max(1, worker_budget); + options.worker_budget = workers; options.preserve_order = false; - const auto workers = chunk_pipeline_detail::hardware_worker_count(); if (!parallel_enabled || workers <= 1 || source_bytes < kDefaultBulkBuildMinBytes) { return options; @@ -1887,7 +1955,7 @@ ChunkSourceOptions ResolveBulkBuildSourceOptions( } options.queue_capacity = std::clamp( static_cast(options.producer_count) * 2, 2, kMaxQueuedChunks); - return options; + return NormalizeChunkSourceOptions(std::move(options)); } CSVChunkSource::CSVChunkSource(std::vector file_paths, @@ -1899,15 +1967,16 @@ CSVChunkSource::CSVChunkSource(std::vector file_paths, partition_plan_cache_(std::make_shared()) {} std::shared_ptr CSVChunkSource::Open( - const ChunkSourceOptions& options) const { + const ChunkSourceOptions& requested_options) { if (file_paths_.empty()) { THROW_INVALID_ARGUMENT_EXCEPTION("CSV chunk source has no input paths"); } + const auto options = NormalizeChunkSourceOptions(requested_options); auto open_config = project_csv_config( config_, compose_projection(projected_columns_, options.projected_columns)); - if (open_config.use_threads && options.producer_count > 0 && - !options.preserve_order) { + if (open_config.use_threads && options.worker_budget > 1 && + options.producer_count > 0 && !options.preserve_order) { return std::make_shared( file_paths_, std::move(open_config), options, partition_plan_cache_); } @@ -2198,7 +2267,7 @@ void set_column_from_value_column( if (vids[k] >= std::numeric_limits::max()) continue; if (value_col) { - write(vids[k], value_col->data()[k]); + write(vids[k], value_col->get_value(k)); } else { auto val = ctx_col->get_elem(k); if (!val.IsNull()) @@ -2210,7 +2279,7 @@ void set_column_from_value_column( if (vids[k] >= std::numeric_limits::max()) continue; if (value_col) { - typed->set_value(vids[k], value_col->data()[k]); + typed->set_value(vids[k], value_col->get_value(k)); } else { auto val = ctx_col->get_elem(k); if (!val.IsNull()) diff --git a/src/transaction/insert_transaction.cc b/src/transaction/insert_transaction.cc index 2f55e2578..1c3b976de 100644 --- a/src/transaction/insert_transaction.cc +++ b/src/transaction/insert_transaction.cc @@ -275,7 +275,7 @@ void InsertTransaction::create_id_indexer_if_not_exists(label_t label) { } Status StorageTPInsertInterface::BatchAddVertices( - label_t v_label_id, std::shared_ptr supplier) { + label_t v_label_id, std::unique_ptr source) { LOG(ERROR) << "BatchAddVertices is not supported in TP mode currently."; return Status(StatusCode::ERR_NOT_SUPPORTED, "BatchAddVertices is not supported in TP mode currently."); @@ -283,7 +283,7 @@ Status StorageTPInsertInterface::BatchAddVertices( Status StorageTPInsertInterface::BatchAddEdges( label_t src_label, label_t dst_label, label_t edge_label, - std::shared_ptr supplier) { + std::unique_ptr source) { LOG(ERROR) << "BatchAddEdges is not supported in TP mode currently."; return Status(StatusCode::ERR_NOT_SUPPORTED, "BatchAddEdges is not supported in TP mode currently."); diff --git a/src/transaction/update_transaction.cc b/src/transaction/update_transaction.cc index d1aa1063c..997901c2c 100644 --- a/src/transaction/update_transaction.cc +++ b/src/transaction/update_transaction.cc @@ -1192,7 +1192,7 @@ void StorageTPUpdateInterface::CreateCheckpoint() { } Status StorageTPUpdateInterface::BatchAddVertices( - label_t v_label_id, std::shared_ptr supplier) { + label_t v_label_id, std::unique_ptr source) { LOG(ERROR) << "BatchAddVertices is not supported in TP mode currently."; return Status(StatusCode::ERR_NOT_SUPPORTED, "BatchAddVertices is not supported in TP mode currently."); @@ -1200,7 +1200,7 @@ Status StorageTPUpdateInterface::BatchAddVertices( Status StorageTPUpdateInterface::BatchAddEdges( label_t src_label, label_t dst_label, label_t edge_label, - std::shared_ptr supplier) { + std::unique_ptr source) { LOG(ERROR) << "BatchAddEdges is not supported in TP mode currently."; return Status(StatusCode::ERR_NOT_SUPPORTED, "BatchAddEdges is not supported in TP mode currently."); diff --git a/src/utils/io/read/csv/csv_reader.cc b/src/utils/io/read/csv/csv_reader.cc index 1c9bccb94..a27d6fb1e 100644 --- a/src/utils/io/read/csv/csv_reader.cc +++ b/src/utils/io/read/csv/csv_reader.cc @@ -485,7 +485,7 @@ CsvReader::CsvReader(std::shared_ptr sharedState, CsvReader::~CsvReader() = default; -std::shared_ptr CsvReader::createChunkSource( +std::unique_ptr CsvReader::createChunkSource( std::vector projected_columns) { if (!sharedState_) { THROW_INVALID_ARGUMENT_EXCEPTION("SharedState is null"); @@ -512,7 +512,7 @@ std::shared_ptr CsvReader::createChunkSource( if (paths.empty()) { THROW_INVALID_ARGUMENT_EXCEPTION("No file paths provided"); } - return std::make_shared(paths, std::move(read_config), + return std::make_unique(paths, std::move(read_config), std::move(projected_columns)); } diff --git a/tests/storage/alter_property_test.cc b/tests/storage/alter_property_test.cc index 0bfce9a36..db78a4586 100644 --- a/tests/storage/alter_property_test.cc +++ b/tests/storage/alter_property_test.cc @@ -116,7 +116,7 @@ void testLoadVertexBatch(PropertyGraph& graph, std::string vertex_type_name, } } auto supplier = std::make_shared(v_file, std::move(config)); - CHECK(graph.BatchAddVertices(v_label, supplier).ok()); + CHECK(graph.BatchAddVertices(v_label, make_data_chunk_source(supplier)).ok()); } void testLoadEdgeBatch(PropertyGraph& graph, std::string src_vertex_type, @@ -194,7 +194,9 @@ void testLoadEdgeBatch(PropertyGraph& graph, std::string src_vertex_type, } } auto supplier = std::make_shared(e_file, std::move(config)); - CHECK(graph.BatchAddEdges(src_label_id, dst_label_id, e_label_id, supplier) + CHECK(graph + .BatchAddEdges(src_label_id, dst_label_id, e_label_id, + make_data_chunk_source(supplier)) .ok()); } diff --git a/tests/storage/test_edge_table.cc b/tests/storage/test_edge_table.cc index 80627167d..86daac734 100644 --- a/tests/storage/test_edge_table.cc +++ b/tests/storage/test_edge_table.cc @@ -32,6 +32,7 @@ #include "neug/storages/graph/edge_table.h" #include "neug/storages/loader/loader_utils.h" #include "neug/storages/module_descriptor.h" +#include "neug/utils/exception/exception.h" #include "unittest/utils.h" namespace neug { @@ -193,12 +194,14 @@ class EdgeTableTest : public ::testing::Test { void BatchInsert(std::vector>&& chunks) { auto supplier = std::make_shared(std::move(chunks)); - edge_table->BatchAddEdges(src_indexer, dst_indexer, supplier); + edge_table->BatchAddEdges(src_indexer, dst_indexer, + make_data_chunk_source(std::move(supplier))); } void BatchBuild(std::shared_ptr source) { edge_table->BatchAddEdges(src_indexer, dst_indexer, - make_data_chunk_supplier(std::move(source))); + make_data_chunk_source(std::move(source)), + BulkLoadOptions{4}); } void BatchBuild(std::vector> chunks, @@ -832,10 +835,9 @@ TEST_F(EdgeTableTest, BatchBuildEdgesHandlesParallelVariants) { BatchBuild(source); EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); - EXPECT_EQ(source->OpenCount(), 2); - ASSERT_EQ(source->OpenedProjections().size(), 2); - EXPECT_EQ(source->OpenedProjections()[0], (std::vector{0, 1})); - EXPECT_TRUE(source->OpenedProjections()[1].empty()); + EXPECT_EQ(source->OpenCount(), 1); + ASSERT_EQ(source->OpenedProjections().size(), 1); + EXPECT_TRUE(source->OpenedProjections()[0].empty()); std::vector> expected; expected.reserve(kEdgeNum); for (size_t row = 0; row < kEdgeNum; ++row) { @@ -927,7 +929,7 @@ TEST_F(EdgeTableTest, BatchBuildEdgesHandlesInputsWithNoStoredEdges) { EXPECT_EQ(edge_table->EdgeNum(), 0); }; - run(edge_label_int_, {}, 2); + run(edge_label_int_, {}, 1); run(edge_label_none_, MakeIntEdgeChunks(std::vector{0, 1, 2}, std::vector{1, 2, 3}, @@ -935,6 +937,18 @@ TEST_F(EdgeTableTest, BatchBuildEdgesHandlesInputsWithNoStoredEdges) { 0); } +TEST_F(EdgeTableTest, BatchBuildMissingBundledPropertyReportsSchemaMismatch) { + auto ckp = make_checkpoint(workspace()); + InitEdgeTable(ckp, 2, 2, edge_label_int_); + auto chunks = convert_to_data_chunks( + {split_column_to_chunks(std::vector{0}, 1), + split_column_to_chunks(std::vector{1}, 1)}); + + EXPECT_THROW(BatchBuild(std::move(chunks)), + exception::SchemaMismatchException); + EXPECT_EQ(edge_table->EdgeNum(), 0); +} + TEST_F(EdgeTableTest, SecondBatchFallsBackWhenIncomingCsrIsNotEmpty) { auto ckp = make_checkpoint(workspace()); constexpr int64_t kVertexNum = 3; @@ -946,7 +960,7 @@ TEST_F(EdgeTableTest, SecondBatchFallsBackWhenIncomingCsrIsNotEmpty) { MakeIntEdgeSource(std::vector{0}, std::vector{0}, std::vector{10}, 1); BatchBuild(first_source); - EXPECT_EQ(first_source->OpenCount(), 2); + EXPECT_EQ(first_source->OpenCount(), 1); auto second_source = MakeIntEdgeSource(std::vector{1}, std::vector{1}, @@ -957,11 +971,7 @@ TEST_F(EdgeTableTest, SecondBatchFallsBackWhenIncomingCsrIsNotEmpty) { ExpectIncomingEdges({{0, 0, 10}, {1, 1, 20}}); } -TEST_F(EdgeTableTest, SingleEdgeBulkBuildFillsUniqueSlotsAcrossChunks) { - if (std::thread::hardware_concurrency() < 4) { - GTEST_SKIP() << "Concurrent bulk fill needs at least two consumer workers"; - } - +TEST_F(EdgeTableTest, SingleEdgeBulkBuildPreservesOrderAcrossChunks) { struct SupplierActivity { std::atomic active{0}; std::atomic max_active{0}; @@ -1011,15 +1021,11 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildFillsUniqueSlotsAcrossChunks) { kForceBulkBuildBytes); BatchBuild(source); - ASSERT_EQ(source->OpenedOptions().size(), 2); - EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); - EXPECT_FALSE(source->OpenedOptions()[1].preserve_order); - EXPECT_EQ(source->OpenedOptions()[0].projected_columns, - (std::vector{0, 1})); - EXPECT_TRUE(source->OpenedOptions()[1].projected_columns.empty()); - ASSERT_EQ(activities.size(), 2); - EXPECT_GT(activities[0]->max_active.load(), 1); - EXPECT_GT(activities[1]->max_active.load(), 1); + ASSERT_EQ(source->OpenedOptions().size(), 1); + EXPECT_TRUE(source->OpenedOptions()[0].preserve_order); + EXPECT_TRUE(source->OpenedOptions()[0].projected_columns.empty()); + ASSERT_EQ(activities.size(), 1); + EXPECT_EQ(activities[0]->max_active.load(), 1); EXPECT_EQ(edge_table->EdgeNum(), kSrcNum); std::vector> expected; @@ -1041,7 +1047,7 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildHandlesEndpointAndDirectionVariants) { std::vector{10, 20, 30, 40}, 4); BatchBuild(source); - EXPECT_EQ(source->OpenCount(), 2); + EXPECT_EQ(source->OpenCount(), 1); EXPECT_EQ(edge_table->EdgeNum(), 2); const std::vector> expected = { {0, 0, 10}, {2, 2, 40}}; @@ -1071,20 +1077,7 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildHandlesEndpointAndDirectionVariants) { preserves_both_single_directions(); } -TEST_F(EdgeTableTest, SingleEdgeBulkBuildDetectsCrossWorkerDuplicate) { - if (std::thread::hardware_concurrency() < 4) { - GTEST_SKIP() - << "Cross-worker duplicate detection needs two consumer workers"; - } - - struct BarrierState { - std::atomic next{0}; - std::mutex mutex_; - std::condition_variable cv_; - size_t arrived_ = 0; - bool cancelled_ = false; - }; - +TEST_F(EdgeTableTest, SingleEdgeBulkBuildDuplicateUsesLastValue) { auto ckp = make_checkpoint(workspace()); InitEdgeTable(ckp, 1, 2, edge_label_single_); auto chunks = @@ -1095,53 +1088,20 @@ TEST_F(EdgeTableTest, SingleEdgeBulkBuildDetectsCrossWorkerDuplicate) { std::make_shared>>( std::move(chunks)); auto source = std::make_shared( - [shared_chunks](const ChunkSourceOptions& options, size_t) { - if (options.preserve_order) { - return std::shared_ptr( - std::make_shared(*shared_chunks, true)); - } - auto state = std::make_shared(); + [shared_chunks](const ChunkSourceOptions&, size_t) { return std::shared_ptr( - std::make_shared( - [shared_chunks, state] { - const auto index = - state->next.fetch_add(1, std::memory_order_relaxed); - if (index >= shared_chunks->size()) { - return std::shared_ptr{}; - } - if (index < 2) { - std::unique_lock lock(state->mutex_); - ++state->arrived_; - state->cv_.notify_all(); - state->cv_.wait(lock, [&] { - return state->arrived_ >= 2 || state->cancelled_; - }); - if (state->cancelled_) { - return std::shared_ptr{}; - } - } - return (*shared_chunks)[index]; - }, - 2, true, - [state] { - { - std::lock_guard lock(state->mutex_); - state->cancelled_ = true; - } - state->cv_.notify_all(); - })); + std::make_shared(*shared_chunks, true)); }, kForceBulkBuildBytes); BatchBuild(source); - ASSERT_EQ(source->OpenedOptions().size(), 2); - EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); - EXPECT_TRUE(source->OpenedOptions()[1].preserve_order); + ASSERT_EQ(source->OpenedOptions().size(), 1); + EXPECT_TRUE(source->OpenedOptions()[0].preserve_order); EXPECT_EQ(edge_table->EdgeNum(), 2); ExpectOutgoingEdges({{0, 1, 20}}); } -TEST_F(EdgeTableTest, BatchBuildEdgesFromPartitionedCsvInTwoPasses) { +TEST_F(EdgeTableTest, BatchBuildEdgesFromPartitionedCsvWithSpill) { auto ckp = make_checkpoint(workspace()); constexpr int64_t kSrcNum = 100; constexpr int64_t kDstNum = 80; @@ -1181,9 +1141,7 @@ TEST_F(EdgeTableTest, BatchBuildEdgesFromPartitionedCsvInTwoPasses) { OpenEdgeTableInMemory(ckp, CheckpointManifest(), kSrcNum, kDstNum); auto csv_source = std::make_shared( std::vector{csv_path.string()}, std::move(config)); - // Force the large-file planner while keeping the fixture small. The wrapper - // forwards both Open() calls to the same CSV source, so its cached partition - // plan is reused by the count and fill passes. + // Force the large-file planner while keeping the fixture small. auto source = std::make_shared( [csv_source](const ChunkSourceOptions& options, size_t) { return csv_source->Open(options); @@ -1192,43 +1150,104 @@ TEST_F(EdgeTableTest, BatchBuildEdgesFromPartitionedCsvInTwoPasses) { BatchBuild(source); EXPECT_EQ(edge_table->EdgeNum(), kEdgeNum); - EXPECT_EQ(source->OpenCount(), 2); + EXPECT_EQ(source->OpenCount(), 1); + ASSERT_EQ(source->OpenedOptions().size(), 1); + const auto& opened = source->OpenedOptions().front(); + EXPECT_EQ(opened.worker_budget, 4); + EXPECT_LE(opened.producer_count + opened.consumer_count, + opened.worker_budget); ExpectOutgoingEdges(expected, true); ExpectIncomingEdges(std::move(expected), true); } -TEST_F(EdgeTableTest, BatchBuildFailuresDoNotPublishPartialCsr) { - for (size_t failure_pass : {size_t{0}, size_t{1}}) { - SCOPED_TRACE(::testing::Message() << "failure pass " << failure_pass); - auto ckp = make_checkpoint(workspace()); - InitEdgeTable(ckp, 1, 1, edge_label_int_); - auto chunk = MakeIntEdgeChunk(); +TEST_F(EdgeTableTest, + SingleAndMultipleBulkBuildUsesParallelSpillAndOrderedReplay) { + auto ckp = make_checkpoint(workspace()); + constexpr int64_t kSrcNum = 64; + constexpr int64_t kDstNum = 97; + constexpr size_t kEdgeNum = 4096; + std::vector> incoming; + incoming.reserve(kEdgeNum); + std::vector> outgoing(kSrcNum); + + const auto csv_path = temp_dir() / "parallel-single-edges.csv"; + { + std::ofstream output(csv_path, std::ios::binary); + ASSERT_TRUE(output.is_open()); + output << "src|dst|data\n"; + for (size_t row = 0; row < kEdgeNum; ++row) { + const auto src = static_cast(row % kSrcNum); + const auto dst = static_cast((row * 31) % kDstNum); + const auto data = static_cast(row); + output << src << '|' << dst << '|' << data << '\n'; + incoming.emplace_back(src, dst, data); + outgoing[static_cast(src)] = {src, dst, data}; + } + } - auto source = std::make_shared( - [chunk, failure_pass](const ChunkSourceOptions&, size_t open_index) { - if (open_index != failure_pass) { - return std::shared_ptr( - std::make_shared( - std::vector>{chunk})); - } - auto remaining = std::make_shared>(chunk); - return std::shared_ptr( - std::make_shared( - [remaining] { - if (*remaining) { - return std::exchange(*remaining, - std::shared_ptr{}); - } - throw std::runtime_error("injected bulk build failure"); - }, - 1)); - }, - kForceBulkBuildBytes); + CsvReadConfig config; + config.delimiter = '|'; + config.skip_rows = 1; + config.chunk_size = 5; + config.column_names = {"src", "dst", "data"}; + config.include_columns = config.column_names; + config.column_types.emplace("src", DataType(DataTypeId::kInt64)); + config.column_types.emplace("dst", DataType(DataTypeId::kInt64)); + config.column_types.emplace("data", DataType(DataTypeId::kInt32)); - EXPECT_THROW(BatchBuild(source), std::runtime_error); - EXPECT_EQ(source->OpenCount(), failure_pass + 1); - EXPECT_EQ(edge_table->EdgeNum(), 0); - } + InitIndexers(*ckp, kSrcNum, kDstNum); + ConstructEdgeTable(src_label_, dst_label_, edge_label_single_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), kSrcNum, kDstNum); + auto csv_source = std::make_shared( + std::vector{csv_path.string()}, std::move(config)); + auto source = std::make_shared( + [csv_source](const ChunkSourceOptions& options, size_t) { + return csv_source->Open(options); + }, + kForceBulkBuildBytes, true, true); + BatchBuild(source); + + ASSERT_EQ(source->OpenedOptions().size(), 1); + const auto& opened = source->OpenedOptions().front(); + EXPECT_FALSE(opened.preserve_order); + EXPECT_GT(opened.producer_count, 0); + EXPECT_GT(opened.consumer_count, 1); + EXPECT_LE(opened.producer_count + opened.consumer_count, + opened.worker_budget); + ExpectOutgoingEdges(outgoing, true); + ExpectIncomingEdges(std::move(incoming), true); +} + +TEST_F(EdgeTableTest, BatchBuildFailuresDoNotPublishPartialCsr) { + auto ckp = make_checkpoint(workspace()); + InitEdgeTable(ckp, 1, 1, edge_label_int_); + auto chunk = MakeIntEdgeChunk(); + const auto runtime_file_count = [&] { + return static_cast( + std::distance(std::filesystem::directory_iterator(ckp->runtime_dir()), + std::filesystem::directory_iterator())); + }; + const auto files_before = runtime_file_count(); + auto source = std::make_shared( + [chunk](const ChunkSourceOptions&, size_t) { + auto remaining = std::make_shared>(chunk); + return std::shared_ptr( + std::make_shared( + [remaining] { + if (*remaining) { + return std::exchange(*remaining, + std::shared_ptr{}); + } + throw std::runtime_error("injected bulk build failure"); + }, + 1)); + }, + kForceBulkBuildBytes); + + EXPECT_THROW(BatchBuild(source), std::runtime_error); + EXPECT_EQ(source->OpenCount(), 1); + EXPECT_EQ(edge_table->EdgeNum(), 0); + EXPECT_EQ(runtime_file_count(), files_before); } TEST_F(EdgeTableTest, @@ -1237,20 +1256,13 @@ TEST_F(EdgeTableTest, GTEST_SKIP() << "Concurrent supplier cancellation requires two workers"; } - auto run = [&](neug::label_t edge_label, size_t failure_pass) { - SCOPED_TRACE(::testing::Message() << "failure pass " << failure_pass); + auto run = [&](neug::label_t edge_label) { auto ckp = make_checkpoint(workspace()); InitEdgeTable(ckp, 1, 1, edge_label); auto chunk = MakeIntEdgeChunk(); std::atomic cancel_count{0}; auto source = std::make_shared( - [chunk, failure_pass, &cancel_count](const ChunkSourceOptions&, - size_t open_index) { - if (open_index != failure_pass) { - return std::shared_ptr( - std::make_shared( - std::vector>{chunk})); - } + [chunk, &cancel_count](const ChunkSourceOptions&, size_t) { auto next_call = std::make_shared>(0); return std::shared_ptr( std::make_shared( @@ -1274,87 +1286,51 @@ TEST_F(EdgeTableTest, kForceBulkBuildBytes); EXPECT_THROW(BatchBuild(source), std::runtime_error); - EXPECT_EQ(source->OpenCount(), failure_pass + 1); + EXPECT_EQ(source->OpenCount(), 1); EXPECT_EQ(cancel_count.load(std::memory_order_relaxed), 1); EXPECT_EQ(edge_table->EdgeNum(), 0); - if (failure_pass == 0) { - ASSERT_EQ(source->OpenedOptions().size(), 1); - EXPECT_FALSE(source->OpenedOptions()[0].preserve_order); - EXPECT_EQ(source->OpenedOptions()[0].projected_columns, - (std::vector{0, 1})); - } + ASSERT_EQ(source->OpenedOptions().size(), 1); + EXPECT_TRUE(source->OpenedOptions()[0].projected_columns.empty()); }; - run(edge_label_int_, 1); - if (std::thread::hardware_concurrency() >= 4) { - run(edge_label_single_, 0); - } + run(edge_label_int_); } -TEST_F(EdgeTableTest, - NonConcurrentSecondPassFailureCancelsBlockedSupplierAndDoesNotPublish) { - if (std::thread::hardware_concurrency() < 4) { - GTEST_SKIP() << "Bounded pipeline cancellation requires two consumers"; - } - - struct BlockingSupplierState { - std::shared_ptr chunk; - std::mutex mutex; - std::condition_variable cancelled_cv; - bool cancelled = false; - std::atomic cancel_count{0}; - std::atomic timed_out{false}; - }; +TEST_F(EdgeTableTest, OperationOwnedSourceRejectsSecondOpen) { + auto source = + MakeIntEdgeSource(std::vector{0}, std::vector{0}, + std::vector{42}, 1); + auto input = make_data_chunk_source(source); + auto supplier = open_data_chunk_source(*input); + ASSERT_NE(supplier, nullptr); + EXPECT_THROW(open_data_chunk_source(*input), + exception::InvalidArgumentException); + EXPECT_EQ(source->OpenCount(), 1); +} - auto ckp = make_checkpoint(workspace()); - InitEdgeTable(ckp, 1, 1, edge_label_int_); - auto valid_chunk = MakeIntEdgeChunk(); - // The source indexer expects int64 endpoints. The uint64 source column makes - // a fill consumer throw after the producer has requested its next chunk. - auto invalid_chunk = MakeIntEdgeChunk(uint64_t{0}); - std::shared_ptr second_pass_state; - auto source = std::make_shared( - [valid_chunk, invalid_chunk, &second_pass_state]( - const ChunkSourceOptions&, size_t open_index) { - if (open_index == 0) { - return std::shared_ptr( - std::make_shared( - std::vector>{valid_chunk})); - } - second_pass_state = std::make_shared(); - second_pass_state->chunk = invalid_chunk; +TEST_F(EdgeTableTest, SourceOpenClampsWorkerBudget) { + auto source = + std::make_shared([](const ChunkSourceOptions&, size_t) { return std::shared_ptr( std::make_shared( - [state = second_pass_state] { - if (state->chunk) { - return std::exchange(state->chunk, nullptr); - } - std::unique_lock lock(state->mutex); - if (!state->cancelled_cv.wait_for( - lock, std::chrono::seconds(2), - [&] { return state->cancelled; })) { - state->timed_out.store(true, std::memory_order_relaxed); - } - return std::shared_ptr{}; - }, - 1, false, - [state = second_pass_state] { - { - std::lock_guard lock(state->mutex); - state->cancelled = true; - } - state->cancel_count.fetch_add(1, std::memory_order_relaxed); - state->cancelled_cv.notify_all(); - })); - }, - kForceBulkBuildBytes); + [] { return std::shared_ptr{}; }, 0)); + }); + ChunkSourceOptions options{ + .producer_count = 8, + .consumer_count = 7, + .worker_budget = 4, + .preserve_order = false, + }; - EXPECT_THROW(BatchBuild(source), std::exception); - EXPECT_EQ(source->OpenCount(), 2); - ASSERT_NE(second_pass_state, nullptr); - EXPECT_EQ(second_pass_state->cancel_count.load(std::memory_order_relaxed), 1); - EXPECT_FALSE(second_pass_state->timed_out.load(std::memory_order_relaxed)); - EXPECT_EQ(edge_table->EdgeNum(), 0); + auto supplier = open_data_chunk_source(*source, options); + ASSERT_NE(supplier, nullptr); + ASSERT_EQ(source->OpenedOptions().size(), 1); + const auto& opened = source->OpenedOptions().front(); + EXPECT_EQ(opened.worker_budget, 4); + EXPECT_EQ(opened.producer_count, 3); + EXPECT_EQ(opened.consumer_count, 1); + EXPECT_LE(opened.producer_count + opened.consumer_count, + opened.worker_budget); } TEST_F(EdgeTableTest, BatchInsertPathsUseIndexerVertexCapacity) { @@ -2214,7 +2190,7 @@ TYPED_TEST(EdgeTableToolsTest, TestBatchAddEdges) { EdgeTable e_table = EdgeTable(edge_schema); OpenEdgeTableLegacy(e_table, temp_ckp, neug::CheckpointManifest(), MemoryLevel::kInMemory); - e_table.BatchAddEdges(indexer, indexer, suppliers[0]); + e_table.BatchAddEdges(indexer, indexer, make_data_chunk_source(suppliers[0])); EXPECT_EQ(e_table.EdgeNum(), 10); EXPECT_EQ(e_table.PropTableSize(), 0); EXPECT_EQ(e_table.Capacity(), neug::CsrBase::INFINITE_CAPACITY); @@ -2267,7 +2243,7 @@ TYPED_TEST(EdgeTableToolsTest, TestAddProperties) { EdgeTable e_table = EdgeTable(edge_schema); OpenEdgeTableLegacy(e_table, temp_ckp, neug::CheckpointManifest(), MemoryLevel::kInMemory); - e_table.BatchAddEdges(indexer, indexer, suppliers[0]); + e_table.BatchAddEdges(indexer, indexer, make_data_chunk_source(suppliers[0])); EXPECT_EQ(e_table.EdgeNum(), 10); EXPECT_EQ(e_table.PropTableSize(), 0); EXPECT_EQ(e_table.Capacity(), neug::CsrBase::INFINITE_CAPACITY); diff --git a/tests/storage/test_vertex_table.cc b/tests/storage/test_vertex_table.cc index 3d60870f7..758d2990b 100644 --- a/tests/storage/test_vertex_table.cc +++ b/tests/storage/test_vertex_table.cc @@ -687,7 +687,7 @@ TEST_F(VertexTableTest, VertexTableResizeTest) { auto data_chunks = generate_data_chunks(10000); std::shared_ptr batch_supplier = std::make_shared(std::move(data_chunks)); - table.BatchAddVertices(batch_supplier); + table.BatchAddVertices(neug::make_data_chunk_source(batch_supplier)); EXPECT_EQ(table.VertexNum(), 10000); EXPECT_EQ(table.LidNum(), 10000); @@ -707,7 +707,7 @@ TEST_F(VertexTableTest, VertexTableResizeTest) { } } -TEST_F(VertexTableTest, InsertVerticesFromRepeatableSource) { +TEST_F(VertexTableTest, InsertVerticesFromOneShotSource) { neug::VertexTable table(schema_.get_vertex_schema(v_label_id_)); auto ckp = make_checkpoint(Workspace()); OpenVertexTableLegacy(table, ckp, neug::CheckpointManifest(), memory_level_); @@ -715,7 +715,7 @@ TEST_F(VertexTableTest, InsertVerticesFromRepeatableSource) { constexpr size_t kVertexNum = 4097; auto source = std::make_shared( generate_data_chunks(kVertexNum), kForceBulkBuildBytes); - table.BatchAddVertices(make_data_chunk_supplier(source)); + table.BatchAddVertices(neug::make_data_chunk_source(source)); EXPECT_EQ(table.VertexNum(), kVertexNum); EXPECT_EQ(table.LidNum(), kVertexNum); diff --git a/tests/unittest/utils.h b/tests/unittest/utils.h index 2ec50a8fc..3d3780bba 100644 --- a/tests/unittest/utils.h +++ b/tests/unittest/utils.h @@ -120,7 +120,7 @@ class GeneratedChunkSource final : public neug::IDataChunkSource { : chunks_(std::move(chunks)), estimated_bytes_(estimated_bytes) {} std::shared_ptr Open( - const neug::ChunkSourceOptions& options) const override { + const neug::ChunkSourceOptions& options) override { ++open_count_; opened_projections_.push_back(options.projected_columns); auto chunks = chunks_; @@ -153,8 +153,8 @@ class GeneratedChunkSource final : public neug::IDataChunkSource { private: std::vector> chunks_; int64_t estimated_bytes_; - mutable size_t open_count_ = 0; - mutable std::vector> opened_projections_; + size_t open_count_ = 0; + std::vector> opened_projections_; }; class TestChunkSource final : public neug::IDataChunkSource { @@ -163,19 +163,24 @@ class TestChunkSource final : public neug::IDataChunkSource { const neug::ChunkSourceOptions&, size_t)>; explicit TestChunkSource(Factory factory, int64_t estimated_bytes = -1, - bool parallel_enabled = true) + bool parallel_enabled = true, + bool stable_row_ordinals = false) : factory_(std::move(factory)), estimated_bytes_(estimated_bytes), - parallel_enabled_(parallel_enabled) {} + parallel_enabled_(parallel_enabled), + stable_row_ordinals_(stable_row_ordinals) {} std::shared_ptr Open( - const neug::ChunkSourceOptions& options) const override { + const neug::ChunkSourceOptions& options) override { opened_options_.push_back(options); return factory_(options, open_count_++); } int64_t EstimatedBytes() const override { return estimated_bytes_; } bool ParallelEnabled() const override { return parallel_enabled_; } + bool ProvidesStableRowOrdinals() const override { + return stable_row_ordinals_; + } size_t OpenCount() const { return open_count_; } const std::vector& OpenedOptions() const { @@ -186,8 +191,9 @@ class TestChunkSource final : public neug::IDataChunkSource { Factory factory_; int64_t estimated_bytes_; bool parallel_enabled_; - mutable size_t open_count_ = 0; - mutable std::vector opened_options_; + bool stable_row_ordinals_; + size_t open_count_ = 0; + std::vector opened_options_; }; template diff --git a/tests/utils/test_reader.cc b/tests/utils/test_reader.cc index 73d142815..53e351e22 100644 --- a/tests/utils/test_reader.cc +++ b/tests/utils/test_reader.cc @@ -17,6 +17,7 @@ #include #include +#include #include #include "neug/storages/loader/chunk_pipeline_utils.h" @@ -32,6 +33,7 @@ ChunkSourceOptions parallel_source_options( std::vector projected_columns = {}) { return { .producer_count = producer_count, + .worker_budget = std::max(1, producer_count + 1), .queue_capacity = queue_capacity, .preserve_order = false, .projected_columns = std::move(projected_columns), @@ -67,6 +69,58 @@ std::vector read_sorted_column( } // namespace +TEST(ChunkSourceOptionsTest, BulkBuildPlannerHonorsWorkerBudget) { + constexpr int64_t kLargeInput = 1024LL * 1024 * 1024; + + const auto serial = ResolveBulkBuildSourceOptions( + kLargeInput, true, 1, BulkBuildWorkerStrategy::kBalancedProducerConsumer); + EXPECT_EQ(serial.worker_budget, 1); + EXPECT_EQ(serial.producer_count, 0); + EXPECT_EQ(serial.consumer_count, 1); + + const auto balanced = ResolveBulkBuildSourceOptions( + kLargeInput, true, 4, BulkBuildWorkerStrategy::kBalancedProducerConsumer); + EXPECT_EQ(balanced.worker_budget, 4); + EXPECT_GT(balanced.producer_count, 0); + EXPECT_GE(balanced.consumer_count, 1); + EXPECT_LE(balanced.producer_count + balanced.consumer_count, + balanced.worker_budget); + + const auto clamped = ResolveBulkBuildSourceOptions( + kLargeInput, true, 0, BulkBuildWorkerStrategy::kMaxProducers); + EXPECT_EQ(clamped.worker_budget, 1); + EXPECT_EQ(clamped.producer_count, 0); +} + +TEST(ChunkSourceOptionsTest, ExecutionBoundaryClampsMalformedWorkerCounts) { + ChunkSourceOptions malformed{ + .producer_count = 8, + .consumer_count = 7, + .worker_budget = 4, + .queue_capacity = 0, + .preserve_order = false, + }; + const auto normalized = NormalizeChunkSourceOptions(malformed); + EXPECT_EQ(normalized.worker_budget, 4); + EXPECT_EQ(normalized.producer_count, 3); + EXPECT_EQ(normalized.consumer_count, 1); + EXPECT_EQ(normalized.queue_capacity, 1); + EXPECT_LE(normalized.producer_count + normalized.consumer_count, + normalized.worker_budget); + + malformed.worker_budget = 0; + const auto serial = NormalizeChunkSourceOptions(malformed); + EXPECT_EQ(serial.worker_budget, 1); + EXPECT_EQ(serial.producer_count, 0); + EXPECT_EQ(serial.consumer_count, 1); + + malformed.worker_budget = 8; + malformed.preserve_order = true; + const auto ordered = NormalizeChunkSourceOptions(malformed); + EXPECT_EQ(ordered.producer_count, 0); + EXPECT_EQ(ordered.consumer_count, 1); +} + // Test 1: Basic CSV reading with default options TEST_F(ReaderTest, TestBasicCsvRead) { // Create test CSV file @@ -116,6 +170,69 @@ TEST_F(ReaderTest, CsvChunkSourceReusesPartitionPlanAcrossParallelOpens) { std::filesystem::rename(moved_path, csv_path); } +TEST_F(ReaderTest, CsvChunkSourceSeparatesPartitionPlansByWorkerBudget) { + createCsvFile("budgeted-plan.csv", "id|name\n1|Alice\n2|Bob\n3|Carol\n"); + auto config = csv_config( + {{"id", DataTypeId::kInt32}, {"name", DataTypeId::kVarchar}}, 1); + const auto csv_path = + std::filesystem::path(ARROW_READER_TEST_DIR) / "budgeted-plan.csv"; + const auto moved_path = + std::filesystem::path(ARROW_READER_TEST_DIR) / "budgeted-plan.moved"; + CSVChunkSource source({csv_path.string()}, std::move(config)); + const auto first_options = parallel_source_options(4, 4); + + auto first = source.Open(first_options); + ASSERT_NE(first, nullptr); + EXPECT_EQ(first->RowNum(), 4); + + std::filesystem::rename(csv_path, moved_path); + auto second_options = first_options; + second_options.worker_budget = 8; + auto second = source.Open(second_options); + ASSERT_NE(second, nullptr); + EXPECT_ANY_THROW(second->RowNum()); + std::filesystem::rename(moved_path, csv_path); +} + +TEST_F(ReaderTest, PartitionedCsvProvidesStableRowOrdinals) { + constexpr size_t kRowCount = 257; + std::ostringstream csv; + csv << "id|value\n"; + for (size_t row = 0; row < kRowCount; ++row) { + csv << row << '|' << row * 7 << '\n'; + } + createCsvFile("stable-ordinals.csv", csv.str()); + + auto config = csv_config( + {{"id", DataTypeId::kInt64}, {"value", DataTypeId::kInt64}}, 1); + config.chunk_size = 3; + const auto csv_path = + std::filesystem::path(ARROW_READER_TEST_DIR) / "stable-ordinals.csv"; + CSVChunkSource source({csv_path.string()}, std::move(config)); + auto supplier = source.Open(parallel_source_options(4, 8)); + ASSERT_NE(supplier, nullptr); + ASSERT_TRUE(supplier->SupportsConcurrentGetNext()); + + std::vector chunks; + while (true) { + auto chunk = supplier->GetNextChunkWithOrdinal(); + if (!chunk.chunk) { + break; + } + chunks.push_back(std::move(chunk)); + } + std::sort(chunks.begin(), chunks.end(), [](const auto& lhs, const auto& rhs) { + return lhs.first_row_ordinal < rhs.first_row_ordinal; + }); + + uint64_t expected_ordinal = 0; + for (const auto& sequenced : chunks) { + EXPECT_EQ(sequenced.first_row_ordinal, expected_ordinal); + expected_ordinal += sequenced.chunk->row_num(); + } + EXPECT_EQ(expected_ordinal, kRowCount); +} + TEST_F(ReaderTest, PartitionedCsvHandlesQuotedRecordBoundaries) { { SCOPED_TRACE("quoted records"); From 4cebeea9b7ca9ea798005c382207fd86f68aae52 Mon Sep 17 00:00:00 2001 From: "xiaolei.zl" Date: Mon, 27 Jul 2026 10:10:12 +0800 Subject: [PATCH 8/8] resolve review comments --- .../storages/loader/bundled_edge_csr_loader.h | 0 .../neug/storages/loader/chunk_pipeline_utils.h | 5 ----- src/main/query_processor.cc | 14 +++++++++++--- src/storages/graph/edge_table.cc | 3 +-- src/storages/loader/bundled_edge_csr_loader.cc | 2 +- 5 files changed, 13 insertions(+), 11 deletions(-) rename {src => include/neug}/storages/loader/bundled_edge_csr_loader.h (100%) diff --git a/src/storages/loader/bundled_edge_csr_loader.h b/include/neug/storages/loader/bundled_edge_csr_loader.h similarity index 100% rename from src/storages/loader/bundled_edge_csr_loader.h rename to include/neug/storages/loader/bundled_edge_csr_loader.h diff --git a/include/neug/storages/loader/chunk_pipeline_utils.h b/include/neug/storages/loader/chunk_pipeline_utils.h index 298da84ad..1b24fa3be 100644 --- a/include/neug/storages/loader/chunk_pipeline_utils.h +++ b/include/neug/storages/loader/chunk_pipeline_utils.h @@ -35,11 +35,6 @@ namespace neug { namespace chunk_pipeline_detail { -inline int32_t hardware_worker_count() { - auto workers = static_cast(std::thread::hardware_concurrency()); - return workers <= 0 ? 1 : workers; -} - template class BoundedQueue { public: diff --git a/src/main/query_processor.cc b/src/main/query_processor.cc index 8a503bd99..5a2b4b5bd 100644 --- a/src/main/query_processor.cc +++ b/src/main/query_processor.cc @@ -51,9 +51,17 @@ QueryProcessor::check_and_retrieve_pipeline( result QueryProcessor::resolve_thread_budget( int32_t requested_threads) const { - if (requested_threads < 0 || max_thread_num_ < 1) { - RETURN_ERROR(neug::Status(neug::StatusCode::ERR_INVALID_ARGUMENT, - "Number of threads must be greater than 0")); + // Zero requests the configured default, so only negative counts are + // rejected here. + if (requested_threads < 0) { + RETURN_ERROR( + neug::Status(neug::StatusCode::ERR_INVALID_ARGUMENT, + "Number of threads must be non-negative (0 means the " + "configured default)")); + } + if (max_thread_num_ < 1) { + RETURN_ERROR(neug::Status(neug::StatusCode::ERR_INTERNAL_ERROR, + "Max thread number must be greater than 0")); } if (requested_threads == 0) { return max_thread_num_; diff --git a/src/storages/graph/edge_table.cc b/src/storages/graph/edge_table.cc index b689b344a..f26418611 100644 --- a/src/storages/graph/edge_table.cc +++ b/src/storages/graph/edge_table.cc @@ -37,14 +37,13 @@ #include "neug/storages/csr/csr_view_utils.h" #include "neug/storages/csr/immutable_csr.h" #include "neug/storages/csr/mutable_csr.h" +#include "neug/storages/loader/bundled_edge_csr_loader.h" #include "neug/storages/loader/loader_utils.h" #include "neug/storages/module/type_name.h" #include "neug/storages/module_descriptor.h" #include "neug/utils/io/file/file_utils.h" #include "neug/utils/property/types.h" -#include "../loader/bundled_edge_csr_loader.h" - namespace neug { void filterInvalidEdges(std::vector& src_lid, diff --git a/src/storages/loader/bundled_edge_csr_loader.cc b/src/storages/loader/bundled_edge_csr_loader.cc index 218ced9fa..5da43e944 100644 --- a/src/storages/loader/bundled_edge_csr_loader.cc +++ b/src/storages/loader/bundled_edge_csr_loader.cc @@ -13,7 +13,7 @@ * limitations under the License. */ -#include "bundled_edge_csr_loader.h" +#include "neug/storages/loader/bundled_edge_csr_loader.h" #include