diff --git a/include/neug/compiler/function/import/csv_read_function.h b/include/neug/compiler/function/import/csv_read_function.h index 9a55756f2..f278e1ef3 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)); @@ -111,27 +112,53 @@ 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); return ctx; } + static std::unique_ptr sourceFunc( + std::shared_ptr state, + std::vector projected_columns) { + 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); + resolvePaths(source_state); + return createReader(source_state) + ->createChunkSource(std::move(projected_columns)); + } + static std::shared_ptr sniffFunc( const reader::FileSchema& schema) { auto state = std::make_shared(); @@ -143,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) { @@ -170,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) { @@ -184,4 +199,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..42bc6fec6 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 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)>; + // 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..0877a2101 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,23 @@ class BatchInsertEdgeOprBuilder : public IOperatorBuilder { } }; +/// Fuses only a terminal, empty-sink COPY FROM plan. Storage chooses staged +/// build or normal BatchAdd before opening the supplied source once. +class BatchInsertEdgeFromSourceOprBuilder : public IOperatorBuilder { + public: + 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..82fbee605 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,23 @@ class BatchInsertVertexOprBuilder : public IOperatorBuilder { } }; +/// Fuses only a terminal, empty-sink COPY FROM plan. Storage chooses staged +/// build or normal BatchAdd before opening the supplied source once. +class BatchInsertVertexFromSourceOprBuilder : public IOperatorBuilder { + public: + 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..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,11 +18,16 @@ #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 { +class PhysicalPlan; class PropertyMapping; -} +} // namespace physical +namespace common { +class NameOrId; +} // namespace common namespace google { namespace protobuf { template @@ -31,9 +36,14 @@ class RepeatedPtrField; } // namespace google namespace neug { -class IDataChunkSupplier; class Schema; class StorageReadInterface; +namespace function { +struct ReadFunction; +} +namespace reader { +struct ReadSharedState; +} namespace execution { namespace ops { @@ -65,6 +75,29 @@ std::shared_ptr create_data_chunk_supplier( const Context& ctx, const std::vector>& prop_mappings); +bool resolve_vertex_label_id(const Schema& schema, + const ::common::NameOrId& type, label_t& label_id); + +struct BatchInsertInput { + std::unique_ptr data; + Context output; +}; + +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); std::vector> create_csv_chunk_suppliers( 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/csr/mutable_csr.h b/include/neug/storages/csr/mutable_csr.h index 61a685502..881213686 100644 --- a/include/neug/storages/csr/mutable_csr.h +++ b/include/neug/storages/csr/mutable_csr.h @@ -43,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. @@ -50,6 +54,12 @@ static_assert( sizeof(std::atomic) == sizeof(int), "atomic must have the same size as int on supported platforms"); +namespace mutable_csr_detail { + +int capacity_with_reserve(int degree); + +} // namespace mutable_csr_detail + template class MutableCsr : public TypedCsrBase { public: @@ -235,6 +245,8 @@ class MutableCsr : public TypedCsrBase { } private: + friend class internal::BundledEdgeCsrLoader; + std::unique_ptr locks_; std::shared_ptr adj_list_buffer_; std::shared_ptr degree_list_; @@ -382,6 +394,8 @@ class SingleMutableCsr : public TypedCsrBase { } private: + friend class internal::BundledEdgeCsrLoader; + std::shared_ptr nbr_list_; std::atomic edge_num_{0}; CsrPrefetchPolicy prefetch_policy_; diff --git a/include/neug/storages/graph/edge_table.h b/include/neug/storages/graph/edge_table.h index 56ac97a6a..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,8 +41,6 @@ class ModuleBroker; class CheckpointManifest; class PropertyGraph; -class IDataChunkSupplier; - class EdgeTable { public: EdgeTable(std::shared_ptr meta) : meta_(meta) {} @@ -133,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, @@ -190,6 +190,11 @@ class EdgeTable { void DetachInAdjlist(vid_t vid, Allocator& alloc); private: + bool TryBatchBuildEdges(const IndexerType& src_indexer, + const IndexerType& dst_indexer, + 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 8ca36284a..86cb1dfa6 100644 --- a/include/neug/storages/graph/vertex_table.h +++ b/include/neug/storages/graph/vertex_table.h @@ -257,7 +257,8 @@ class VertexTable { void Compact(timestamp_t ts = MAX_TIMESTAMP); - void insert_vertices(std::shared_ptr suppliers); + void BatchAddVertices(std::unique_ptr source, + BulkLoadOptions options = {}); 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(IDataChunkSource& source, + BulkLoadOptions options); + + 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/bundled_edge_csr_loader.h b/include/neug/storages/loader/bundled_edge_csr_loader.h new file mode 100644 index 000000000..7971262fb --- /dev/null +++ b/include/neug/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 Checkpoint; +class IDataChunkSource; +struct BulkLoadOptions; + +namespace internal { + +/// 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 + class MutableWriter; + + template + class SingleMutableWriter; + + /// 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, IDataChunkSource& source, + int64_t source_bytes, vid_t src_vertex_capacity, + vid_t dst_vertex_capacity, Checkpoint& checkpoint, + BulkLoadOptions options); +}; + +} // namespace internal +} // namespace neug 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..1b24fa3be --- /dev/null +++ b/include/neug/storages/loader/chunk_pipeline_utils.h @@ -0,0 +1,261 @@ +/** 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 { + +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; +}; + +/// 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, + int32_t consumer_count, + size_t queue_capacity, + Consume&& consume) { + consumer_count = std::max(1, consumer_count); + if (consumer_count == 1) { + uint64_t next_row_ordinal = 0; + while (auto chunk = supplier.GetNextChunk()) { + 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); + 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 { + uint64_t next_row_ordinal = 0; + while (!cancelled.load(std::memory_order_acquire)) { + auto chunk = supplier.GetNextChunk(); + 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(); + }); + + std::vector consumers; + consumers.reserve(static_cast(consumer_count)); + for (int32_t i = 0; i < consumer_count; ++i) { + consumers.emplace_back([&, i] { + try { + SequencedDataChunk 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); + } +} + +/// 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_impl(IDataChunkSupplier& supplier, + int32_t consumer_count, + Consume&& consume) { + CHECK(supplier.SupportsConcurrentGetNext()); + consumer_count = std::max(1, consumer_count); + if (consumer_count == 1) { + while (true) { + auto chunk = supplier.GetNextChunkWithOrdinal(); + if (!chunk.chunk) { + break; + } + 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.GetNextChunkWithOrdinal(); + if (!chunk.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 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) { + const auto normalized = NormalizeChunkSourceOptions(options); + auto consumer_count = normalized.consumer_count; + if (supplier.SupportsConcurrentGetNext()) { + chunk_pipeline_detail::consume_concurrent_supplier_impl( + 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, consumer_count, normalized.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 32b33aab9..c499fc3a3 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 @@ -68,32 +72,164 @@ CsvReadConfig build_csv_read_config( const std::unordered_map& csv_options, const std::vector& column_types); +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() {} +}; + +struct ChunkSourceOptions { + /// 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; + /// Maximum concurrently active workers for source planning and execution. + int32_t worker_budget = 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; +}; + +/// 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. +class IDataChunkSource { + public: + virtual ~IDataChunkSource() = default; + + /// 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 = {}) = 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; } +}; + +/// 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, int32_t worker_budget, + 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; +} + +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_; }; -using CSVStreamChunkSupplier = CSVChunkSupplier; -using CSVTableChunkSupplier = CSVChunkSupplier; +struct CsvPartitionPlanCache; + +/// 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 = {}) 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 keeps partition planning lazy until Open(). + std::shared_ptr partition_plan_cache_; +}; void fillVertexReaderMeta(label_t v_label, const std::string& v_label_name, const std::string& v_file, 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_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..b637aa83a 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,11 @@ class CsvReader { void read(std::shared_ptr localState, execution::Context& ctx); + /// Creates a configurable CSV source for direct COPY FROM bulk loading. + /// Returns nullptr when the read needs a row filter and must materialize. + std::unique_ptr createChunkSource( + std::vector projected_columns = {}); + 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..d327f5d22 100644 --- a/src/execution/execute/ops/batch/batch_insert_edge.cc +++ b/src/execution/execute/ops/batch/batch_insert_edge.cc @@ -14,6 +14,7 @@ */ #include "neug/execution/execute/ops/batch/batch_insert_edge.h" +#include "neug/compiler/function/read_function.h" #include "neug/execution/common/context.h" #include "neug/execution/execute/ops/batch/batch_update_utils.h" #include "neug/storages/graph/graph_interface.h" @@ -21,6 +22,7 @@ #include "neug/utils/result.h" #include +#include #include #include @@ -35,37 +37,16 @@ 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, 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: " @@ -82,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 @@ -94,13 +72,15 @@ class BatchInsertEdgeOpr : public IOperator { public: BatchInsertEdgeOpr( physical::EdgeType edge_type, - std::vector> prop_mappings, - std::vector> src_vertex_bindings, - std::vector> dst_vertex_bindings) + std::vector> property_mappings, + std::vector> source_mappings, + std::vector> destination_mappings, + std::optional source = std::nullopt) : 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)) {} + property_mappings_(std::move(property_mappings)), + source_mappings_(std::move(source_mappings)), + destination_mappings_(std::move(destination_mappings)), + source_(std::move(source)) {} std::string get_operator_name() const override { return "BatchInsertEdgeOpr"; @@ -111,8 +91,9 @@ class BatchInsertEdgeOpr : public IOperator { private: physical::EdgeType edge_type_; - std::vector> prop_mappings_, - src_vertex_bindings_, dst_vertex_bindings_; + std::vector> property_mappings_, + source_mappings_, destination_mappings_; + std::optional source_; }; neug::result BatchInsertEdgeOpr::Eval( @@ -131,30 +112,33 @@ 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); + 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, + mappings); + } else { + input.data = + make_data_chunk_source(create_data_chunk_supplier(ctx, mappings)); + input.output = std::move(ctx); } - 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)); + RETURN_STATUS_ERROR_IF_NOT_OK(graph.BatchAddEdges( + src_label_id, dst_label_id, edge_label_id, std::move(input.data))); + return neug::result(std::move(input.output)); } 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()) { @@ -174,7 +158,37 @@ 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; + if (!is_terminal_batch_insert(plan, op_idx)) { + 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()) { + THROW_INTERNAL_EXCEPTION( + "BatchInsertEdgeFromSourceOprBuilder: edge type is not set"); + } + + auto source = build_batch_insert_source(plan, op_idx); + + 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(), + destination_mappings); + physical::EdgeType edge_type; + edge_type.CopyFrom(edge_pb.edge_type()); + 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)), + 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 b4454f3a8..2bad3defb 100644 --- a/src/execution/execute/ops/batch/batch_insert_vertex.cc +++ b/src/execution/execute/ops/batch/batch_insert_vertex.cc @@ -14,13 +14,14 @@ */ #include "neug/execution/execute/ops/batch/batch_insert_vertex.h" +#include "neug/compiler/function/read_function.h" #include "neug/execution/common/context.h" #include "neug/execution/execute/ops/batch/batch_update_utils.h" #include "neug/storages/graph/graph_interface.h" #include "neug/utils/exception/exception.h" #include "neug/utils/result.h" -#include +#include #include #include @@ -34,10 +35,12 @@ namespace ops { class BatchInsertVertexOpr : public IOperator { public: BatchInsertVertexOpr( - common::NameOrId vertex_type, - std::vector> prop_mappings) + ::common::NameOrId vertex_type, + 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"; @@ -47,8 +50,9 @@ class BatchInsertVertexOpr : public IOperator { Context&& ctx, OprTimer* timer) override; private: - common::NameOrId vertex_type_; - std::vector> prop_mappings_; + ::common::NameOrId vertex_type_; + std::vector> property_mappings_; + std::optional source_; }; neug::result BatchInsertVertexOpr::Eval( @@ -58,36 +62,30 @@ 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; + if (!resolve_vertex_label_id(graph.schema(), vertex_type_, vertex_label_id)) { + RETURN_STATUS_ERROR(StatusCode::ERR_INVALID_ARGUMENT, + "Failed to resolve vertex type " + + vertex_type_.ShortDebugString() + + " for BatchInsertVertex"); } - default: - THROW_INVALID_ARGUMENT_EXCEPTION( - "BatchInsertVertexOpr: invalid vertex_type: " + - vertex_type_.DebugString()); + BatchInsertInput input; + if (source_) { + input = create_batch_insert_input(source_->state, *source_->read_function, + property_mappings_); + } else { + input.data = make_data_chunk_source( + create_data_chunk_supplier(ctx, property_mappings_)); + input.output = std::move(ctx); } - 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)); + graph.BatchAddVertices(vertex_label_id, std::move(input.data))); + return neug::result(std::move(input.output)); } 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,11 +94,36 @@ 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); + ctx_meta); +} + +neug::result BatchInsertVertexFromSourceOprBuilder::Build( + const Schema& schema, const ContextMeta& ctx_meta, + const physical::PhysicalPlan& plan, int op_idx) { + (void) schema; + if (!is_terminal_batch_insert(plan, op_idx)) { + 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()) { + THROW_INTERNAL_EXCEPTION( + "BatchInsertVertexFromSourceOprBuilder: vertex type is not set"); + } + + 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(vertex_type), + std::move(property_mappings), std::move(source)), + ctx_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 4339f2947..1051067aa 100644 --- a/src/execution/execute/ops/batch/batch_update_utils.cc +++ b/src/execution/execute/ops/batch/batch_update_utils.cc @@ -22,6 +22,7 @@ #include #include +#include #include #include #include @@ -30,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" @@ -338,6 +342,73 @@ std::shared_ptr create_data_chunk_supplier( return std::make_shared(std::move(projected_chunks)); } +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; + } +} + +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()); + 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( + 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); + } + auto source = + read_function.sourceFunc(shared_state, std::move(projected_columns)); + if (source) { + return {std::move(source), Context{}}; + } + } + + 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 {make_data_chunk_source(std::move(supplier)), std::move(output)}; +} + 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/main/query_processor.cc b/src/main/query_processor.cc index fb8d89dc7..5a2b4b5bd 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,44 @@ 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 { + // 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_; + } + 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 +95,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 +117,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 +130,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/csr/mutable_csr.cc b/src/storages/csr/mutable_csr.cc index 5152fcf95..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 = std::ceil(new_degree * NeugDBConfig::DEFAULT_RESERVE_RATIO); + 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 98664b616..f26418611 100644 --- a/src/storages/graph/edge_table.cc +++ b/src/storages/graph/edge_table.cc @@ -21,17 +21,23 @@ #include #include +#include #include #include +#include +#include #include +#include #include #include +#include #include "neug/common/columns/value_columns.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/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" @@ -765,9 +771,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()); + 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()); + 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. + 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(); @@ -779,6 +810,7 @@ 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) { auto chunk = supplier->GetNextChunk(); if (chunk == nullptr) { @@ -805,7 +837,6 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, } } } - std::vector valid_flags; filterInvalidEdges(src_lid, dst_lid, valid_flags); size_t new_size = table_idx_.load() + src_lid.size(); if (new_size >= Capacity()) { @@ -828,6 +859,35 @@ void EdgeTable::BatchAddEdges(const IndexerType& src_indexer, } } +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(); + if (!ShouldUseBulkBuild(source_bytes)) { + return false; + } + + EdgeTable staged(meta_); + 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, *ckp_, + options)) { + return false; + } + Swap(staged); + return true; +} + void EdgeTable::BatchAddEdges( const std::vector& src_lid_list, const std::vector& dst_lid_list, diff --git a/src/storages/graph/graph_interface.cc b/src/storages/graph/graph_interface.cc index eba93c230..b274d046f 100644 --- a/src/storages/graph/graph_interface.cc +++ b/src/storages/graph/graph_interface.cc @@ -116,15 +116,24 @@ 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)); + 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_); + } + return status; } 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)); + std::unique_ptr source) { + auto status = graph_.BatchAddEdges(src_label, dst_label, edge_label, + std::move(source), bulk_load_options_); + if (status.ok()) { + mut_view_.Rebuild(graph_); + } + return status; } Status StorageAPUpdateInterface::BatchDeleteVertices( diff --git a/src/storages/graph/property_graph.cc b/src/storages/graph/property_graph.cc index 54d9b1631..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].insert_vertices(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(), 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 e1150fc22..496b00f93 100644 --- a/src/storages/graph/vertex_table.cc +++ b/src/storages/graph/vertex_table.cc @@ -15,7 +15,10 @@ #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" @@ -45,21 +48,78 @@ void VertexTable::Init(std::shared_ptr ckp, MemoryLevel level) { v_ts_->Open(*ckp_, ModuleDescriptor{}, level); } -void VertexTable::insert_vertices( +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(open_data_chunk_source(*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 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; + } + + 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) { + 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); + } + }; + auto row_nums = supplier->RowNum(); 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"); } + reserve_checkpoint_headroom(indexer_->size() + row_count); while (true) { auto chunk = supplier->GetNextChunk(); if (chunk == nullptr) { @@ -85,17 +145,12 @@ void VertexTable::insert_vertices( // Capacity check for actual batch size. size_t chunk_rows = chunk->row_num(); - 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); + if (chunk_rows > std::numeric_limits::max() - indexer_->size()) { + THROW_RUNTIME_ERROR("Vertex row count overflow"); } - + size_t new_size = indexer_->size() + chunk_rows; + reserve_checkpoint_headroom(new_size); auto vids = insert_primary_keys(pk_col); - 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); 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 new file mode 100644 index 000000000..5da43e944 --- /dev/null +++ b/src/storages/loader/bundled_edge_csr_loader.cc @@ -0,0 +1,1046 @@ +/** 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 "neug/storages/loader/bundled_edge_csr_loader.h" + +#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.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 { + +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 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 { + +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()), + value_column_(dynamic_cast*>(column.get())) { + if (column_ == nullptr) { + THROW_SCHEMA_MISMATCH("Bundled edge property column is missing"); + } + } + + EDATA_T Get(size_t row) const { + if (value_column_ != nullptr) { + return value_column_->get_value(row); + } + return column_->get_elem(row).template GetValue(); + } + + private: + const IContextColumn* column_; + const ValueColumn* value_column_; +}; + +template <> +class BulkEdgeDataReader { + public: + explicit BulkEdgeDataReader( + const std::shared_ptr& /*column*/) {} + + 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 { + 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(); +} + +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)); + } + } + } +} + +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 = 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 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, materialize); + for (auto& segment : result.segments) { + segment->Finalize(); + } + return result; +} + +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 fill_bulk_edge_chunk_serial(BulkEdgeWorkerScratch& scratch, + 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]; + 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 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; + constexpr bool kDirectIn = InWriter::kStrategy == EdgeStrategy::kSingle; + 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.PutSerial(src, dst, data, 0); + } + if constexpr (kDirectIn) { + 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) { + 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 +class BulkEdgeSpillDataReader { + public: + explicit BulkEdgeSpillDataReader(const BulkEdgeSpillRecord* records) + : records_(records) {} + + EDATA_T Get(size_t row) const { return records_[row].data; } + + 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; + } + + 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 { + for (size_t worker = 0; worker < segments.size(); ++worker) { + replay_segment(worker); + } + } + + 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_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 { + constexpr bool kHasSingleDirection = + OutWriter::kStrategy == EdgeStrategy::kSingle || + InWriter::kStrategy == EdgeStrategy::kSingle; + 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(); + } + 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, 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) { + const bool in_supported = + with_csr_bulk_writer(in_csr, [&](auto& in) { + out.PrepareBuild(src_vertex_capacity); + in.PrepareBuild(dst_vertex_capacity); + const auto filled_edge_count = + build_bundled_edges_with_writers( + out, in, src_indexer, dst_indexer, source, source_bytes, + checkpoint, bulk_load_options); + 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, + 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(); + 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, 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, checkpoint, + bulk_load_options); + default: + return false; + } +} + +} // namespace + +bool internal::BundledEdgeCsrLoader::TryBuild( + CsrBase& out_csr, CsrBase& in_csr, const EdgeSchema& schema, + 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 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/loader_utils.cc b/src/storages/loader/loader_utils.cc index 4c2d0e8ba..61917eb88 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,39 @@ 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; + uint64_t first_row_ordinal = 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, int32_t worker_budget); + + private: + struct Entry { + std::once_flag once; + std::shared_ptr plan; + }; + + std::mutex mutex_; + std::map, std::unique_ptr> entries_; +}; + namespace { constexpr size_t kDefaultCsvChunkRows = 4096; @@ -80,6 +122,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; } @@ -413,43 +456,88 @@ 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. -class CsvRowCountCounter { +/// Result of scanning a CSV file without materializing its fields. +struct CsvScanResult { + int64_t row_count = 0; + 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. +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) + 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), - double_quote_(double_quote), - delimiter_(delimiter) {} + 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; - 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) { - 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); - return count_parallel(file_size, num_threads); + return scan_parallel(file_size, num_threads, false).row_count; + } + + 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 = std::min( + file_size, + static_cast(std::max(1, requested_partitions))); + 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}); + 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: @@ -458,6 +546,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 @@ -474,6 +574,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); @@ -501,11 +604,28 @@ 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) { - 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; } @@ -517,8 +637,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) { @@ -543,7 +663,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]); @@ -556,7 +676,8 @@ class CsvRowCountCounter { } /// Parallel scan with speculative dual state machines. - int64_t count_parallel(size_t file_size, unsigned num_threads) const { + 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 +693,41 @@ 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> 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) { + 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; @@ -599,21 +741,351 @@ class CsvRowCountCounter { if (last_selected && last_selected->has_content) ++total; - return total; + CsvScanResult output; + output.row_count = total; + if (!build_ranges) { + return output; + } + + // 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) { + output.ranges.push_back({safe_bounds[i].first, safe_bounds[i + 1].first, + safe_bounds[i].second}); + } + return output; } std::string file_path_; bool quoting_; char quote_char_; - bool double_quote_; char delimiter_; + bool use_threads_; }; +class CsvRangeStreamBuf final : public std::streambuf { + public: + 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); + } + 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 (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; + } + + const auto requested_count = static_cast(count); + const auto buffered = static_cast(egptr() - gptr()); + const auto from_buffer = std::min(buffered, requested_count); + if (from_buffer > 0) { + std::memcpy(destination, gptr(), from_buffer); + gbump(static_cast(from_buffer)); + } + + if (from_buffer == requested_count || remaining_ == 0) { + return static_cast(from_buffer); + } + 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_); + } + remaining_ -= read; + if (read < requested) { + remaining_ = 0; + } + return static_cast(from_buffer + read); + } + + private: + std::ifstream file_; + std::string file_path_; + size_t remaining_ = 0; + std::array buffer_{}; +}; + +class CsvRangeStream final : public std::istream { + public: + CsvRangeStream(const std::string& file_path, size_t start, size_t end) + : std::istream(nullptr), buffer_(file_path, start, end) { + rdbuf(&buffer_); + } + + private: + CsvRangeStreamBuf buffer_; +}; + +std::shared_ptr build_csv_partition_plan( + const std::vector& file_paths, const CsvReadConfig& config, + int32_t producer_count, int32_t worker_budget) { + 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; + } + } + + 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 = std::max(1, worker_budget); + auto scan_threads = range_counts; + if (non_empty_files < static_cast(scan_budget)) { + 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(); + 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; + } + } + + std::vector scans(file_paths.size()); + std::atomic next_file{0}; + std::atomic scan_cancelled{false}; + std::exception_ptr scan_error; + auto capture_scan_error = [&](std::exception_ptr error) { + if (!scan_cancelled.exchange(true, std::memory_order_acq_rel)) { + scan_error = std::move(error); + } + }; + auto scan_file = [&](size_t file_index) { + if (file_sizes[file_index] == 0) { + return; + } + 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]), + range_counts[file_index], + 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; + uint64_t next_row_ordinal = 0; + for (size_t i = 0; i < file_paths.size(); ++i) { + const auto& scan = scans[i]; + 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 + // 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, 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; + 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, 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, 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, + worker_budget); + }); + 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 +1099,23 @@ 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), + 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_ = CsvFileScanner(file_path, quoting_, quote_char_, delimiter_, + use_threads_) + .count_rows(); + } + if (range_) { + csv_format_.threading(false); + } reset_reader(); } @@ -696,14 +1178,30 @@ struct CsvSupplierRuntime { return chunk; } - int64_t row_num() const { return row_num_; } + int64_t row_num() const { + if (row_num_ == kUnknownRowNum) { + if (range_) { + return kUnknownRowNum; + } + row_num_ = CsvFileScanner(file_path_, quoting_, quote_char_, delimiter_, + use_threads_) + .count_rows(); + } + 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 +1230,14 @@ struct CsvSupplierRuntime { size_t chunk_size_ = kDefaultCsvChunkRows; bool escaping_ = false; char escape_char_ = '\\'; - int64_t row_num_ = 0; + bool quoting_ = false; + char quote_char_ = '"'; + 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 +1507,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 +1547,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 +1564,448 @@ 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 SupplierChunkSource final : public IDataChunkSource { + public: + explicit SupplierChunkSource(std::shared_ptr supplier) + : supplier_(std::move(supplier)) {} + + std::shared_ptr Open( + const ChunkSourceOptions& /*options*/) override { + if (!supplier_) { + THROW_INVALID_ARGUMENT_EXCEPTION("Data chunk source is empty"); + } + return std::move(supplier_); + } + + bool ParallelEnabled() const override { return false; } + + private: + std::shared_ptr supplier_; +}; + +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); + } + + int64_t EstimatedBytes() const override { + return source_ ? source_->EstimatedBytes() : -1; + } + + bool ParallelEnabled() const override { + return source_ && source_->ParallelEnabled(); + } + + bool ProvidesStableRowOrdinals() const override { + return source_ && source_->ProvidesStableRowOrdinals(); + } + + private: + std::shared_ptr source_; +}; + +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; +} + +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, + 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, + 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) {} + + ~PartitionedCsvChunkSupplier() override { + Cancel(); + JoinWorkers(); + } + + std::shared_ptr GetNextChunk() override { + return GetNextChunkWithOrdinal().chunk; + } + + SequencedDataChunk GetNextChunkWithOrdinal() override { + EnsureScanned(); + StartWorkers(); + SequencedDataChunk chunk; + if (queue_.Pop(chunk)) { + return chunk; + } + RethrowError(); + return {}; + } + + int64_t RowNum() const override { + const_cast(this)->EnsureScanned(); + return plan_->row_count; + } + + bool SupportsConcurrentGetNext() const override { return true; } + bool ProvidesStableRowOrdinals() 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_, + worker_budget_); + 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() { + 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; + 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) { + 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()); } + if (active_workers_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + 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_; + int32_t worker_budget_; + 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}; + 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_unique( + 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_) { + CSVChunkSupplier supplier(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::unique_ptr current_; + size_t next_file_ = 0; + mutable int64_t row_num_ = kUnknownRowNum; +}; + +} // namespace + +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_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, 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; + if (!parallel_enabled || workers <= 1 || + source_bytes < kDefaultBulkBuildMinBytes) { + return options; + } + + const auto useful_partitions = (source_bytes - 1) / kMinPartitionBytes + 1; + switch (worker_strategy) { + case BulkBuildWorkerStrategy::kMaxProducers: { + options.producer_count = + static_cast(std::min(useful_partitions, workers - 1)); + 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.queue_capacity = std::clamp( + static_cast(options.producer_count) * 2, 2, kMaxQueuedChunks); + return NormalizeChunkSourceOptions(std::move(options)); +} + +CSVChunkSource::CSVChunkSource(std::vector file_paths, + 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 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.worker_budget > 1 && + options.producer_count > 0 && !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, @@ -1314,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()) @@ -1326,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/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..a27d6fb1e 100644 --- a/src/utils/io/read/csv/csv_reader.cc +++ b/src/utils/io/read/csv/csv_reader.cc @@ -485,6 +485,37 @@ CsvReader::CsvReader(std::shared_ptr sharedState, CsvReader::~CsvReader() = default; +std::unique_ptr CsvReader::createChunkSource( + std::vector projected_columns) { + 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_unique(paths, std::move(read_config), + std::move(projected_columns)); +} + void CsvReader::read(std::shared_ptr /*localState*/, execution::Context& ctx) { if (!sharedState_) { 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_copy_temp.cc b/tests/storage/test_copy_temp.cc index 4b0ffae81..8eeac087a 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,62 @@ TEST_F(CopyTempTest, NodeDefaultPrimaryKey) { conn->Close(); } +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"; + 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); + + // 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 \"" + edges + "\" (header = true);"); + ASSERT_TRUE(fallback) << fallback.error().ToString(); + } + 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(), 3); + 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..86daac734 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" @@ -24,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 { @@ -68,12 +77,34 @@ 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"); + 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"); 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"); + 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()); @@ -125,9 +156,58 @@ 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); + 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_source(std::move(source)), + BulkLoadOptions{4}); + } + + void BatchBuild(std::vector> chunks, + int64_t estimated_bytes = kForceBulkBuildBytes) { + BatchBuild(std::make_shared(std::move(chunks), + estimated_bytes)); } size_t ExpectedBatchInsertCapacity(size_t inserted_edge_num) const { @@ -216,6 +296,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)) { @@ -237,9 +359,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_, edge_label_out_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 +820,554 @@ TEST_F(EdgeTableTest, TestBatchAddEdgesBundled) { ASSERT_EQ(dsts.size(), edge_num + more_edge_num); } +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(), 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) { + expected.emplace_back(srcs[row], dsts[row], data[row]); + } + ExpectOutgoingEdges(std::move(expected), true); + }; + + 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); + }; + + 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); + }; + + with_property(); + without_properties(); + supernode(); +} + +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); + }; + + run(edge_label_int_, {}, 1); + 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, 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; + InitIndexers(*ckp, kVertexNum, kVertexNum); + ConstructEdgeTable(src_label_, dst_label_, edge_label_out_none_); + OpenEdgeTableInMemory(ckp, CheckpointManifest(), kVertexNum, kVertexNum); + + auto first_source = + MakeIntEdgeSource(std::vector{0}, std::vector{0}, + std::vector{10}, 1); + BatchBuild(first_source); + EXPECT_EQ(first_source->OpenCount(), 1); + + auto second_source = + MakeIntEdgeSource(std::vector{1}, std::vector{1}, + std::vector{20}, 1); + BatchBuild(second_source); + EXPECT_EQ(second_source->OpenCount(), 1); + + ExpectIncomingEdges({{0, 0, 10}, {1, 1, 20}}); +} + +TEST_F(EdgeTableTest, SingleEdgeBulkBuildPreservesOrderAcrossChunks) { + struct SupplierActivity { + std::atomic active{0}; + std::atomic max_active{0}; + }; + + 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 = MakeIntEdgeChunks(srcs, dsts, data, 128); + ASSERT_EQ(chunks.size(), 128); + + InitEdgeTable(ckp, kSrcNum, kDstNum, edge_label_single_); + 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(), 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; + expected.reserve(srcs.size()); + for (size_t row = 0; row < srcs.size(); ++row) { + expected.emplace_back(srcs[row], dsts[row], data[row]); + } + ExpectOutgoingEdges(expected); + ExpectIncomingEdges(std::move(expected), true); +} + +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(), 1); + EXPECT_EQ(edge_table->EdgeNum(), 2); + const std::vector> expected = { + {0, 0, 10}, {2, 2, 40}}; + ExpectOutgoingEdges(expected); + ExpectIncomingEdges(expected); + }; + + 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, SingleEdgeBulkBuildDuplicateUsesLastValue) { + auto ckp = make_checkpoint(workspace()); + 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>>( + std::move(chunks)); + auto source = std::make_shared( + [shared_chunks](const ChunkSourceOptions&, size_t) { + return std::shared_ptr( + std::make_shared(*shared_chunks, true)); + }, + kForceBulkBuildBytes); + BatchBuild(source); + + 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, BatchBuildEdgesFromPartitionedCsvWithSpill) { + 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. + 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(), 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, + 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}; + } + } + + 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)); + + 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, + ConcurrentBatchBuildFailuresCancelSupplierAndDoNotPublish) { + if (std::thread::hardware_concurrency() < 2) { + GTEST_SKIP() << "Concurrent supplier cancellation requires two workers"; + } + + 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, &cancel_count](const ChunkSourceOptions&, size_t) { + 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 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(), 1); + EXPECT_EQ(cancel_count.load(std::memory_order_relaxed), 1); + EXPECT_EQ(edge_table->EdgeNum(), 0); + ASSERT_EQ(source->OpenedOptions().size(), 1); + EXPECT_TRUE(source->OpenedOptions()[0].projected_columns.empty()); + }; + + run(edge_label_int_); +} + +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); +} + +TEST_F(EdgeTableTest, SourceOpenClampsWorkerBudget) { + auto source = + std::make_shared([](const ChunkSourceOptions&, size_t) { + return std::shared_ptr( + std::make_shared( + [] { return std::shared_ptr{}; }, 0)); + }); + ChunkSourceOptions options{ + .producer_count = 8, + .consumer_count = 7, + .worker_budget = 4, + .preserve_order = false, + }; + + 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) { + 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)); + } + + 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); + }; + + constexpr neug::vid_t kBulkVertexNum = 4097; + run(true, kBulkVertexNum, kBulkVertexNum + kBulkVertexNum / 4); + run(false, 16, 128); +} + TEST_F(EdgeTableTest, TestBatchAddEdgesUnbundled) { auto ckp = make_checkpoint(workspace()); int64_t src_num = 100; @@ -716,7 +1389,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)); @@ -1514,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); @@ -1567,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_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 a6b4adeb7..758d2990b 100644 --- a/tests/storage/test_vertex_table.cc +++ b/tests/storage/test_vertex_table.cc @@ -687,10 +687,11 @@ 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(neug::make_data_chunk_source(batch_supplier)); 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,27 @@ TEST_F(VertexTableTest, VertexTableResizeTest) { } } +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_); + + constexpr size_t kVertexNum = 4097; + auto source = std::make_shared( + generate_data_chunks(kVertexNum), kForceBulkBuildBytes); + table.BatchAddVertices(neug::make_data_chunk_source(source)); + + EXPECT_EQ(table.VertexNum(), kVertexNum); + EXPECT_EQ(table.LidNum(), kVertexNum); + EXPECT_EQ(table.Capacity(), kVertexNum + kVertexNum / 4); + EXPECT_EQ(source->OpenCount(), 1); + 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..3d3780bba 100644 --- a/tests/unittest/utils.h +++ b/tests/unittest/utils.h @@ -18,8 +18,11 @@ #include #include +#include #include #include +#include +#include #include #include #include @@ -41,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; } @@ -69,6 +78,122 @@ 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 { + public: + explicit GeneratedChunkSource( + std::vector> chunks, + int64_t estimated_bytes = -1) + : chunks_(std::move(chunks)), estimated_bytes_(estimated_bytes) {} + + std::shared_ptr Open( + const neug::ChunkSourceOptions& options) override { + ++open_count_; + opened_projections_.push_back(options.projected_columns); + 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); + } + + 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_; + size_t open_count_ = 0; + 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, + bool stable_row_ordinals = false) + : factory_(std::move(factory)), + estimated_bytes_(estimated_bytes), + parallel_enabled_(parallel_enabled), + stable_row_ordinals_(stable_row_ordinals) {} + + std::shared_ptr Open( + 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 { + return opened_options_; + } + + private: + Factory factory_; + int64_t estimated_bytes_; + bool parallel_enabled_; + 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 90a4db953..53e351e22 100644 --- a/tests/utils/test_reader.cc +++ b/tests/utils/test_reader.cc @@ -15,9 +15,112 @@ #include "test_reader.h" +#include +#include +#include +#include + +#include "neug/storages/loader/chunk_pipeline_utils.h" +#include "neug/storages/loader/loader_utils.h" + 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, + .worker_budget = std::max(1, producer_count + 1), + .queue_capacity = queue_capacity, + .preserve_order = false, + .projected_columns = std::move(projected_columns), + }; +} + +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(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 @@ -45,6 +148,287 @@ TEST_F(ReaderTest, TestBasicCsvRead) { EXPECT_EQ(ctx.row_num(), 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, 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"); + 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"))); + } + + { + 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); + + EXPECT_EQ(read_sorted_column(supplier), + (std::vector{2, 3})); + } + + { + 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})); + } +} + +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); + + 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}})); + } + + { + 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()); + } + + { + 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"); + auto config = csv_config( + {{"id", DataTypeId::kInt32}, {"name", DataTypeId::kVarchar}}, 1); + config.quoting = true; + config.escaping = false; + config.chunk_size = 1; + auto path = [](const char* file) { + return std::string(ARROW_READER_TEST_DIR) + "/" + file; + }; + CSVChunkSource source({path("part-a.csv"), path("part-b.csv")}, config); + + auto supplier = source.Open(parallel_source_options()); + ASSERT_NE(supplier, nullptr); + EXPECT_EQ(supplier->RowNum(), + 6); // Two headers are counted as reserve hints. + + 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"); + auto config = csv_config( + {{"id", DataTypeId::kInt32}, {"name", DataTypeId::kVarchar}}, 1); + config.quoting = true; + config.escaping = false; + config.chunk_size = 1; + CSVChunkSource source( + {std::string(ARROW_READER_TEST_DIR) + "/partition-error.csv"}, config); + + auto supplier = source.Open(parallel_source_options(4, 2)); + ASSERT_NE(supplier, nullptr); + EXPECT_ANY_THROW({ + while (supplier->GetNextChunk()) {} + }); +} + // 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"