diff --git a/doc/source/development/operator_streams.md b/doc/source/development/operator_streams.md new file mode 100644 index 000000000..613827b10 --- /dev/null +++ b/doc/source/development/operator_streams.md @@ -0,0 +1,69 @@ +# Operator batch streams + +Every concrete `IOperator::Eval` accepts and returns `Stream` +directly. Initialization errors are deferred to `Next()`. A stream is a move-only synchronous pull +interface. `Next()` produces one batch, EOF, or a terminal error. An empty +batch retains its schema and is not EOF. Destroying the stream releases its +cursor without reading the rest of its input. + +`Next()` returns `result>` directly. ContextChunk +owns its DataChunk and anonymous execution head; Stream has no additional Batch +wrapper. Readers and storage suppliers keep their DataChunk interface and adapt +once at the Source/BatchInsert boundary. Between operators, the same ContextChunk +is moved directly through the kernels. Stream tag_ids retain output aliases. + +## Execution and state + +- Source copies its read configuration per execution and opens its reader on + first demand. CSV, JSON and JSONL provide the same supplier interface. +- Project, Select, Unfold, vertex/edge expansion, path operations, intersection + and the triangle kernel use `map_chunks`. One pull invokes the kernel on one + batch and returns that batch directly; no Context is constructed. +- Limit owns its cross-batch skip and remaining-row counters in the returned + stream. Once satisfied, it releases upstream without pulling another batch. +- Union buffers its shared input for branch replay, then concatenates branch + streams on demand. It does not collect or flatten branch outputs. +- The existing sorting, fused Project/OrderBy, grouping, deduplication and + general join kernels require complete input. Their explicit collection + boundaries operate on batches, without converting through Context. The + primary-key join can process each right-side batch independently. +- Mutations retain read/write barriers so a downstream LIMIT cannot skip + writes and writes cannot invalidate data that is still being read. Edge + update/merge retains affected batches to refresh their property pointers. +- BatchInsert consumes the same stream through the storage supplier API and + checks terminal reader errors before reporting success. Sink only sets tags. +- DDL and administration execute once and pass through their stream. + +All mutable cursors and counters belong to the execution's stream, not the +cached operator. `Pipeline::ExecuteStream` connects nested plans without +materializing their results. Its caller must keep the pipeline, storage and +PROFILE timer alive until consumption or destruction. The external +`Pipeline::Execute` query boundary still materializes its public Context result. + +PROFILE charges production during `Next()` and excludes nested upstream time +within a pipeline. Errors retain their producer's operator name while passing +through downstream consumers. + +## Compatibility boundaries and remaining buffering + +Index-scan and export extension callbacks still take Context in their existing +ABI; those operators convert only at that explicit callback boundary. Procedure +and GDS callbacks still return Context, which is exposed as a stream once. +Every registered reader must supply a supplier factory; there is no materialized +reader callback or fallback. Parquet uses RecordBatchReader directly without a +CountRows or ToTable pass. Reader classes expose only supplier creation and schema +inference: their Context-based read methods, ReadLocalState, and the obsolete +batch_read option have been removed. Reader tests consume suppliers directly. + +This does not replace every underlying algorithm with an incremental one. +Graph scan kernels may still produce a large single batch. CSV counts a file +before parsing; JSON array decoding still builds a document. Global sorting, +grouping, deduplication and general join still buffer their inputs. Edge storage +still accumulates endpoints and property batches, and mutation barriers retain +input as described above. The interface is lazy, not asynchronous, and does not +add parallel loading. + +Operators return Stream directly, and ExecuteStream only connects +streams. Fallible execution initialization is deferred until the first Next(); +Next() is the stream error boundary and makes failures terminal. Plan builders +and the materializing public Execute() still return result for their own work. diff --git a/extension/parquet/include/parquet/arrow_reader.h b/extension/parquet/include/parquet/arrow_reader.h index c29eedba5..07b7b1621 100644 --- a/extension/parquet/include/parquet/arrow_reader.h +++ b/extension/parquet/include/parquet/arrow_reader.h @@ -27,6 +27,7 @@ #include "parquet/arrow_options.h" namespace neug { +class IDataChunkSupplier; namespace reader { class DatasetBuilder { @@ -49,8 +50,7 @@ class Reader { fileSystem(std::move(fileSystem)) {} virtual ~Reader() = default; - virtual void read(std::shared_ptr localState, - execution::Context& ctx) = 0; + virtual std::shared_ptr getDataChunkSupplier() = 0; protected: std::shared_ptr sharedState; @@ -74,20 +74,14 @@ class ArrowReader : public Reader { datasetBuilder(std::move(datasetBuilder)) {} ~ArrowReader() override = default; - void read(std::shared_ptr localState, - execution::Context& ctx) override; + std::shared_ptr getDataChunkSupplier() override; arrow::Result> inferSchema(); protected: std::shared_ptr createScanner( std::shared_ptr fs); - void full_read(std::shared_ptr scanner, - execution::Context& output); - void batch_read(std::shared_ptr scanner, - execution::Context& output); - DataChunk finishChunk(DataChunk chunk) const; bool filter_after_read_ = false; std::unique_ptr optionsBuilder; diff --git a/extension/parquet/include/parquet_read_function.h b/extension/parquet/include/parquet_read_function.h index 43177547b..892d5790a 100644 --- a/extension/parquet/include/parquet_read_function.h +++ b/extension/parquet/include/parquet_read_function.h @@ -37,15 +37,15 @@ struct ParquetReadFunction { static function_set getFunctionSet() { auto typeIDs = std::vector<::neug::DataTypeId>{::neug::DataTypeId::kVarchar}; - auto readFunction = std::make_unique(name, typeIDs); - readFunction->execFunc = execFunc; + auto readFunction = + std::make_unique(name, typeIDs, supplierFunc); readFunction->sniffFunc = sniffFunc; function_set functionSet; functionSet.push_back(std::move(readFunction)); return functionSet; } - static execution::Context execFunc( + static std::shared_ptr supplierFunc( std::shared_ptr state) { const auto& vfs = neug::main::MetadataRegistry::getVFS(); const auto& fs = vfs->Provide(state->schema.file); @@ -64,10 +64,7 @@ struct ParquetReadFunction { auto reader = std::make_unique( state, std::move(optionsBuilder), std::move(arrowFs)); - execution::Context ctx; - auto localState = std::make_shared(); - reader->read(localState, ctx); - return ctx; + return reader->getDataChunkSupplier(); } static std::shared_ptr sniffFunc( diff --git a/extension/parquet/src/arrow_reader.cc b/extension/parquet/src/arrow_reader.cc index c6bd6076a..79864055b 100644 --- a/extension/parquet/src/arrow_reader.cc +++ b/extension/parquet/src/arrow_reader.cc @@ -103,7 +103,7 @@ std::vector fallbackProjection(const ReadSharedState& state) { } } - // Keep a stable scan order; finishChunk restores the requested output order. + // Keep a stable scan order; the supplier restores the requested output order. std::vector columns; for (const auto& name : all_columns) { if (required.erase(name)) { @@ -119,27 +119,48 @@ std::vector fallbackProjection(const ReadSharedState& state) { } // namespace -void ArrowReader::read(std::shared_ptr localState, - execution::Context& ctx) { - if (!sharedState) { - THROW_INVALID_ARGUMENT_EXCEPTION("SharedState is null"); - } - - if (!fileSystem) { - THROW_INVALID_ARGUMENT_EXCEPTION("FileSystem is null"); - } - +std::shared_ptr ArrowReader::getDataChunkSupplier() { auto scanner = createScanner(fileSystem); - NEUG_ASSERT(scanner != nullptr); + auto batches = scanner->ToRecordBatchReader(); + if (!batches.ok()) { + THROW_IO_EXCEPTION("Failed to create RecordBatchReader: " + + batches.status().message()); + } + // Row count is unknown until consumption. Never scan the dataset merely to + // count rows before producing its first batch. + auto supplier = + std::make_shared(batches.ValueOrDie(), -1); + if (!filter_after_read_) { + return supplier; + } + class FilteredSupplier final : public IDataChunkSupplier { + public: + FilteredSupplier(std::shared_ptr input, + ReadSharedState state, std::vector columns) + : input_(std::move(input)), + state_(std::move(state)), + columns_(std::move(columns)) {} + + int64_t RowNum() const override { return -1; } + + std::shared_ptr GetNextChunk() override { + auto chunk = input_->GetNextChunk(); + if (!chunk || chunk->col_num() == 0) { + return chunk; + } + auto filtered = + filter_chunk(*chunk, state_.skipRows, columns_, state_.parameters); + return std::make_shared( + project_chunk(filtered, columns_, state_.projectColumns)); + } - // Choose read mode: batch_read streams data, full_read loads entire dataset - const auto& fileSchema = sharedState->schema.file; - ReadOptions options; - if (options.batch_read.get(fileSchema.options)) { - batch_read(scanner, ctx); - } else { - full_read(scanner, ctx); - } + private: + std::shared_ptr input_; + ReadSharedState state_; + std::vector columns_; + }; + return std::make_shared(supplier, *sharedState, + fallback_columns_); } std::shared_ptr ArrowReader::createScanner( @@ -272,90 +293,6 @@ std::shared_ptr ArrowReader::createScanner( return scanner_result.ValueOrDie(); } -void ArrowReader::full_read(std::shared_ptr scanner, - execution::Context& output) { - if (!sharedState) { - THROW_INVALID_ARGUMENT_EXCEPTION("SharedState is null"); - } - if (!scanner) { - THROW_INVALID_ARGUMENT_EXCEPTION("Scanner is null"); - } - - auto table_result = scanner->ToTable(); - if (!table_result.ok()) { - LOG(ERROR) << "Failed to read table via scanner: " - << table_result.status().message(); - THROW_IO_EXCEPTION("Failed to read table via scanner: " + - table_result.status().message()); - } - auto table = table_result.ValueOrDie(); - - int num_cols = - filter_after_read_ ? fallback_columns_.size() : sharedState->columnNum(); - if (num_cols != table->num_columns()) { - THROW_IO_EXCEPTION( - "Column number mismatch between schema and table, schema: " + - std::to_string(num_cols) + - ", table: " + std::to_string(table->num_columns())); - } - - output.clear(); - DataChunk chunk; - for (int i = 0; i < num_cols; ++i) { - auto table_column = table->column(i); - chunk.set(i, arrow_arrays_to_value_column(table_column->chunks())); - } - output.append_chunk(finishChunk(std::move(chunk))); -} - -void ArrowReader::batch_read(std::shared_ptr scanner, - execution::Context& output) { - if (!sharedState) { - THROW_INVALID_ARGUMENT_EXCEPTION("SharedState is null"); - } - if (!scanner) { - THROW_INVALID_ARGUMENT_EXCEPTION("Scanner is null"); - } - auto row_num_result = scanner->CountRows(); - int64_t row_num = 0; - if (!row_num_result.ok()) { - LOG(WARNING) << "Failed to count rows via scanner: " - << row_num_result.status().message(); - THROW_IO_EXCEPTION("Failed to count rows via scanner: " + - row_num_result.status().message()); - } else { - VLOG(10) << "Row count from scanner: " << row_num_result.ValueOrDie(); - row_num = row_num_result.ValueOrDie(); - } - - auto batch_reader_result = scanner->ToRecordBatchReader(); - if (!batch_reader_result.ok()) { - LOG(ERROR) << "Failed to create RecordBatchReader from scanner: " - << batch_reader_result.status().message(); - THROW_IO_EXCEPTION("Failed to create RecordBatchReader from scanner: " + - batch_reader_result.status().message()); - } - auto batch_reader = batch_reader_result.ValueOrDie(); - - auto batch_supplier = - std::make_shared(batch_reader, row_num); - - output.clear(); - while (auto chunk = batch_supplier->GetNextChunk()) { - output.append_chunk(finishChunk(std::move(*chunk))); - } -} - -DataChunk ArrowReader::finishChunk(DataChunk chunk) const { - if (!filter_after_read_ || chunk.col_num() == 0) { - return chunk; - } - auto filtered = filter_chunk(chunk, sharedState->skipRows, fallback_columns_, - sharedState->parameters); - return project_chunk(filtered, fallback_columns_, - sharedState->projectColumns); -} - arrow::Result> ArrowReader::inferSchema() { if (!sharedState) { return arrow::Status::Invalid(neug::StatusCode::ERR_INVALID_ARGUMENT, diff --git a/extension/parquet/src/parquet_options.cc b/extension/parquet/src/parquet_options.cc index 4518de80d..10334e675 100644 --- a/extension/parquet/src/parquet_options.cc +++ b/extension/parquet/src/parquet_options.cc @@ -44,6 +44,9 @@ ArrowOptions ArrowParquetOptionsBuilder::build() const { // this the scanner always runs single-threaded, regardless of // ArrowReaderProperties::set_use_threads() below. ReadOptions readOpts; + ParquetParseOptions parquetOpts; + scanOptions->batch_size = + parquetOpts.row_batch_size.get(state->schema.file.options); scanOptions->use_threads = readOpts.use_threads.get(state->schema.file.options); // The scanner needs an IOContext with an executor for parallel scans. diff --git a/extension/parquet/tests/parquet_test.cc b/extension/parquet/tests/parquet_test.cc index 2e01e03a6..b761c934e 100644 --- a/extension/parquet/tests/parquet_test.cc +++ b/extension/parquet/tests/parquet_test.cc @@ -31,6 +31,7 @@ #include "neug/compiler/common/case_insensitive_map.h" #include "neug/execution/common/context.h" #include "neug/generated/proto/plan/basic_type.pb.h" +#include "neug/storages/loader/loader_utils.h" #include "neug/utils/exception/exception.h" #include "neug/utils/io/read/common/options.h" #include "neug/utils/io/read/common/schema.h" @@ -226,11 +227,13 @@ TEST_F(ParquetTest, ReaderPreservesNativeArithmeticFiltering) { ASSERT_TRUE(builder.skipRows(options)); EXPECT_FALSE( options.scanOptions->filter.Equals(arrow::compute::literal(true))); - for (const auto* batch : {"false", "true"}) { - state->schema.file.options = {{"batch_read", batch}, - {"row_batch_size", "1"}}; + for (const auto* batch_size : {"1", "2"}) { + state->schema.file.options = {{"PARQUET_BATCH_ROWS", batch_size}}; execution::Context output; - createParquetReader(state)->read(nullptr, output); + auto supplier = createParquetReader(state)->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + output.append_chunk(std::move(*chunk)); + } ASSERT_EQ(output.row_num(), 2); ASSERT_EQ(output.col_num(), 1); std::vector names; @@ -264,9 +267,8 @@ TEST_F(ParquetTest, ReaderFallsBackBeforeProjectionAndRebindsParameters) { ASSERT_FALSE(builder.skipRows(options)); EXPECT_TRUE( options.scanOptions->filter.Equals(arrow::compute::literal(true))); - for (const auto* batch : {"false", "true"}) { - state->schema.file.options = {{"batch_read", batch}, - {"row_batch_size", "1"}}; + for (const auto* batch_size : {"1", "2"}) { + state->schema.file.options = {{"PARQUET_BATCH_ROWS", batch_size}}; auto reader = createParquetReader(state); auto scanner = reader->createScanner(std::make_shared()); @@ -277,7 +279,10 @@ TEST_F(ParquetTest, ReaderFallsBackBeforeProjectionAndRebindsParameters) { for (int64_t minimum : {2, 3, 0}) { state->parameters = {{"minimum", Value::INT64(minimum)}}; execution::Context output; - reader->read(nullptr, output); + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + output.append_chunk(std::move(*chunk)); + } EXPECT_EQ(output.row_num(), 3 - minimum); ASSERT_EQ(output.col_num(), 1); for (const auto& chunk : output.chunks()) { @@ -337,9 +342,8 @@ TEST_F(ParquetTest, ReaderPrunesColumnsReferencedInsideNestedPredicates) { state->parameters = {{"fallback", Value::STRING("no match")}}; const auto original = predicate->SerializeAsString(); - for (const auto* batch : {"false", "true"}) { - state->schema.file.options = {{"batch_read", batch}, - {"row_batch_size", "1"}}; + for (const auto* batch_size : {"1", "2"}) { + state->schema.file.options = {{"PARQUET_BATCH_ROWS", batch_size}}; auto reader = createParquetReader(state); // Reuse the reader with reordered, duplicated and all-column outputs. for (const auto& projection : std::vector>{ @@ -359,7 +363,10 @@ TEST_F(ParquetTest, ReaderPrunesColumnsReferencedInsideNestedPredicates) { EXPECT_EQ((*scanned)->schema()->field_names(), expected_columns); execution::Context output; - reader->read(nullptr, output); + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + output.append_chunk(std::move(*chunk)); + } ASSERT_EQ(output.row_num(), 2); const auto& output_columns = projection.empty() ? state->schema.entry->columnNames : projection; @@ -403,9 +410,8 @@ TEST_F(ParquetTest, ReaderPrunesConstantFiltersAndValidatesHiddenColumns) { param->mutable_data_type()->mutable_data_type()->set_primitive_type( ::common::PrimitiveType::DT_BOOL); state->skipRows = predicate; - for (const auto* batch : {"false", "true"}) { - state->schema.file.options = {{"batch_read", batch}, - {"row_batch_size", "1"}}; + for (const auto* batch_size : {"1", "2"}) { + state->schema.file.options = {{"PARQUET_BATCH_ROWS", batch_size}}; auto reader = createParquetReader(state); auto scanner = reader->createScanner(std::make_shared()); @@ -414,7 +420,10 @@ TEST_F(ParquetTest, ReaderPrunesConstantFiltersAndValidatesHiddenColumns) { for (const bool keep : {false, true}) { state->parameters = {{"keep", Value::BOOLEAN(keep)}}; execution::Context output; - reader->read(nullptr, output); + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + output.append_chunk(std::move(*chunk)); + } EXPECT_EQ(output.row_num(), keep ? 3 : 0); EXPECT_EQ(output.col_num(), 1); } @@ -745,13 +754,17 @@ TEST_F(ParquetTest, TestTypeMapping_StringToLargeUtf8) { // Extension should convert to large_utf8 for consistency auto sharedState = createSharedState( "test_string_type.parquet", {"id", "name", "value"}, - {createInt64Type(), createStringType(), createDoubleType()}, - {{"batch_read", "false"}}); + {createInt64Type(), createStringType(), createDoubleType()}, {}); auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Verify string column type auto col1 = ctx.chunk(0).columns()[1]; @@ -798,12 +811,17 @@ TEST_F(ParquetTest, TestTypeMapping_PreserveNumericTypes) { {"int32_col", "int64_col", "double_col", "bool_col"}, {createInt32Type(), createInt64Type(), createDoubleType(), createBoolType()}, - {{"batch_read", "false"}}); + {}); auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 4); EXPECT_EQ(ctx.row_num(), 1); @@ -862,7 +880,7 @@ TEST_F(ParquetTest, TestIntegration_ColumnPruning) { reader::FileSchema fileSchema; fileSchema.paths = {filepath}; fileSchema.format = "parquet"; - fileSchema.options = {{"batch_read", "false"}}; + fileSchema.options = {}; reader::ExternalSchema externalSchema; externalSchema.entry = entrySchema; @@ -873,9 +891,14 @@ TEST_F(ParquetTest, TestIntegration_ColumnPruning) { sharedState->projectColumns = {"id", "score", "grade"}; auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Verify extension translates projectColumns to Arrow projection // Should have 3 columns (id, score, grade - "name" is excluded) @@ -935,7 +958,7 @@ TEST_F(ParquetTest, TestIntegration_FilterPushdown) { reader::FileSchema fileSchema; fileSchema.paths = {filepath}; fileSchema.format = "parquet"; - fileSchema.options = {{"batch_read", "false"}}; + fileSchema.options = {}; reader::ExternalSchema externalSchema; externalSchema.entry = entrySchema; @@ -944,9 +967,14 @@ TEST_F(ParquetTest, TestIntegration_FilterPushdown) { sharedState->skipRows = filterExpr; // Neug's filter expression auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Verify extension translates Neug filter to Arrow filter EXPECT_EQ(ctx.col_num(), 2); @@ -967,38 +995,87 @@ TEST_F(ParquetTest, TestIntegration_FilterPushdown) { } } -TEST_F(ParquetTest, TestIntegration_BatchReadMode) { - createSimpleParquetFile("test_batch_mode.parquet"); +TEST_F(ParquetTest, FilteredSuppliersOwnIndependentExecutionState) { + createSimpleParquetFile("independent_filtered.parquet"); + auto state = createSharedState( + "independent_filtered.parquet", {"id", "name", "value"}, + {createInt64Type(), createStringType(), createDoubleType()}); + state->schema.file.options = {{"PARQUET_BATCH_ROWS", "1"}}; + state->projectColumns = {"name"}; + auto predicate = std::make_shared<::common::Expression>(); + predicate->add_operators()->mutable_var()->mutable_tag()->set_name("id"); + predicate->add_operators()->set_logical(::common::Logical::GT); + auto* parameter = predicate->add_operators()->mutable_param(); + parameter->set_name("minimum"); + *parameter->mutable_data_type()->mutable_data_type() = *createInt64Type(); + state->skipRows = predicate; + auto reader = createParquetReader(state); + state->parameters = {{"minimum", Value::INT64(1)}}; + auto first = reader->getDataChunkSupplier(); + state->parameters = {{"minimum", Value::INT64(2)}}; + auto second = reader->getDataChunkSupplier(); + reader.reset(); + state.reset(); + std::vector first_names; + std::vector second_names; + for (size_t batch = 0; batch < 3; ++batch) { + for (auto entry : {std::make_pair(first, &first_names), + std::make_pair(second, &second_names)}) { + auto chunk = entry.first->GetNextChunk(); + ASSERT_NE(chunk, nullptr); + for (size_t row = 0; row < chunk->row_num(); ++row) { + entry.second->push_back( + chunk->get(0)->get_elem(row).GetValue()); + } + } + } + EXPECT_EQ(first->GetNextChunk(), nullptr); + EXPECT_EQ(second->GetNextChunk(), nullptr); + EXPECT_EQ(first_names, (std::vector{"Bob", "Charlie"})); + EXPECT_EQ(second_names, (std::vector{"Charlie"})); +} - // Test with batch_read=true (streaming mode) - auto sharedState = createSharedState( - "test_batch_mode.parquet", {"id", "name", "value"}, +TEST_F(ParquetTest, SupplierOwnsReaderAndProducesIndividualBatches) { + createSimpleParquetFile("supplier.parquet"); + auto state = createSharedState( + "supplier.parquet", {"id", "name", "value"}, {createInt64Type(), createStringType(), createDoubleType()}, - {{"batch_read", "true"}}); - - auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); - execution::Context ctx; - reader->read(localState, ctx); - - EXPECT_GT(ctx.chunk_num(), 0); // batch mode: data materialized into chunks - EXPECT_GT(ctx.col_num(), 0) << "Extension should materialize data into " - "Context chunks when batch_read=true"; + {{"PARQUET_BATCH_ROWS", "1"}}); + auto reader = createParquetReader(state); + auto supplier = reader->getDataChunkSupplier(); + EXPECT_EQ(supplier->RowNum(), -1); + reader.reset(); + size_t rows = 0; + size_t batches = 0; + while (auto chunk = supplier->GetNextChunk()) { + EXPECT_LE(chunk->row_num(), 1); + rows += chunk->row_num(); + ++batches; + } + EXPECT_EQ(rows, 3); + EXPECT_EQ(batches, 3); +} - // Test with batch_read=false (full read mode) - auto sharedState2 = createSharedState( - "test_batch_mode.parquet", {"id", "name", "value"}, +TEST_F(ParquetTest, TestIntegration_IndependentSuppliers) { + createSimpleParquetFile("independent_suppliers.parquet"); + auto state = createSharedState( + "independent_suppliers.parquet", {"id", "name", "value"}, {createInt64Type(), createStringType(), createDoubleType()}, - {{"batch_read", "false"}}); - - auto reader2 = createParquetReader(sharedState2); - auto localState2 = std::make_shared(); - execution::Context ctx2; - reader2->read(localState2, ctx2); - - auto col0_2 = ctx2.chunk(0).columns()[0]; - EXPECT_EQ(col0_2->column_type(), ContextColumnType::kValue) - << "Extension should use Value column type when batch_read=false"; + {{"PARQUET_BATCH_ROWS", "1"}}); + auto reader = createParquetReader(state); + auto first = reader->getDataChunkSupplier(); + auto second = reader->getDataChunkSupplier(); + for (int64_t id = 1; id <= 3; ++id) { + for (const auto& supplier : {first, second}) { + auto chunk = supplier->GetNextChunk(); + ASSERT_NE(chunk, nullptr); + ASSERT_EQ(chunk->row_num(), 1); + ASSERT_EQ(chunk->col_num(), 3); + EXPECT_EQ(chunk->columns[0]->get_elem(0).GetValue(), id); + } + } + EXPECT_EQ(first->GetNextChunk(), nullptr); + EXPECT_EQ(second->GetNextChunk(), nullptr); } TEST_F(ParquetTest, TestIntegration_ParallelReadMultiRowGroup) { @@ -1029,15 +1106,19 @@ TEST_F(ParquetTest, TestIntegration_ParallelReadMultiRowGroup) { // Read with parallel=true: fragments are split by row group and scanned // with the thread pool; all rows must still come back intact. - auto sharedState = - createSharedState("test_parallel_multi_rg.parquet", {"id", "value"}, - {createInt64Type(), createDoubleType()}, - {{"parallel", "true"}, {"batch_read", "true"}}); + auto sharedState = createSharedState( + "test_parallel_multi_rg.parquet", {"id", "value"}, + {createInt64Type(), createDoubleType()}, {{"parallel", "true"}}); auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } int64_t totalRows = 0; for (size_t i = 0; i < ctx.chunk_num(); ++i) { @@ -1048,14 +1129,18 @@ TEST_F(ParquetTest, TestIntegration_ParallelReadMultiRowGroup) { // A row count alone cannot catch duplicated or dropped rows. Read again in // full (non-batch) mode and verify the aggregate of each value column. - auto sharedState2 = - createSharedState("test_parallel_multi_rg.parquet", {"id", "value"}, - {createInt64Type(), createDoubleType()}, - {{"parallel", "true"}, {"batch_read", "false"}}); + auto sharedState2 = createSharedState( + "test_parallel_multi_rg.parquet", {"id", "value"}, + {createInt64Type(), createDoubleType()}, {{"parallel", "true"}}); auto reader2 = createParquetReader(sharedState2); - auto localState2 = std::make_shared(); + execution::Context ctx2; - reader2->read(localState2, ctx2); + { + auto supplier = reader2->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx2.append_chunk(std::move(*chunk)); + } + } int64_t id_sum = 0; double value_sum = 0.0; @@ -1114,7 +1199,7 @@ TEST_F(ParquetTest, TestIntegration_BatchReadWithFilter) { auto const_opr = filterExpr->add_operators(); const_opr->mutable_const_()->set_f64(90.0); - // batch_read=true + filter + // supplier batches + filter auto sharedState = std::make_shared(); auto entrySchema = std::make_shared(); entrySchema->columnNames = {"id", "score"}; @@ -1123,7 +1208,7 @@ TEST_F(ParquetTest, TestIntegration_BatchReadWithFilter) { reader::FileSchema fileSchema; fileSchema.paths = {filepath}; fileSchema.format = "parquet"; - fileSchema.options = {{"batch_read", "true"}}; + fileSchema.options = {}; reader::ExternalSchema externalSchema; externalSchema.entry = entrySchema; @@ -1132,9 +1217,14 @@ TEST_F(ParquetTest, TestIntegration_BatchReadWithFilter) { sharedState->skipRows = filterExpr; auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Arrow scanner applies filter in both batch and full modes EXPECT_EQ(ctx.col_num(), 2); @@ -1143,8 +1233,8 @@ TEST_F(ParquetTest, TestIntegration_BatchReadWithFilter) { for (size_t i = 0; i < ctx.chunk_num(); ++i) { totalRows += static_cast(ctx.chunk(i).chunk().row_num()); } - EXPECT_EQ(totalRows, 3) - << "batch_read=true with filter should still apply Arrow filter pushdown"; + EXPECT_EQ(totalRows, 3) << "supplier batches with filter should still apply " + "Arrow filter pushdown"; } TEST_F(ParquetTest, TestIntegration_BatchReadWithFilterAndProjection) { @@ -1196,7 +1286,7 @@ TEST_F(ParquetTest, TestIntegration_BatchReadWithFilterAndProjection) { auto const_opr = filterExpr->add_operators(); const_opr->mutable_const_()->set_f64(90.0); - // batch_read=true + filter + projection + // supplier batches + filter + projection auto sharedState = std::make_shared(); auto entrySchema = std::make_shared(); entrySchema->columnNames = {"id", "name", "score", "grade"}; @@ -1206,7 +1296,7 @@ TEST_F(ParquetTest, TestIntegration_BatchReadWithFilterAndProjection) { reader::FileSchema fileSchema; fileSchema.paths = {filepath}; fileSchema.format = "parquet"; - fileSchema.options = {{"batch_read", "true"}}; + fileSchema.options = {}; reader::ExternalSchema externalSchema; externalSchema.entry = entrySchema; @@ -1216,20 +1306,25 @@ TEST_F(ParquetTest, TestIntegration_BatchReadWithFilterAndProjection) { sharedState->skipRows = filterExpr; auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Arrow scanner applies both filter and projection in batch mode EXPECT_EQ(ctx.col_num(), 2) - << "batch_read=true with projection should still apply column pruning"; + << "supplier batches with projection should still apply column pruning"; int64_t totalRows = 0; for (size_t i = 0; i < ctx.chunk_num(); ++i) { totalRows += static_cast(ctx.chunk(i).chunk().row_num()); } EXPECT_EQ(totalRows, 2) - << "batch_read=true with filter+projection should filter to 2 rows"; + << "supplier batches with filter+projection should filter to 2 rows"; EXPECT_EQ(sharedState->columnNum(), 2); } @@ -1292,7 +1387,7 @@ TEST_F(ParquetTest, TestIntegration_CombinedFilterAndProjection) { reader::FileSchema fileSchema; fileSchema.paths = {filepath}; fileSchema.format = "parquet"; - fileSchema.options = {{"batch_read", "false"}}; + fileSchema.options = {}; reader::ExternalSchema externalSchema; externalSchema.entry = entrySchema; @@ -1302,9 +1397,14 @@ TEST_F(ParquetTest, TestIntegration_CombinedFilterAndProjection) { sharedState->skipRows = filterExpr; // Filter score > 90.0 auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Verify extension correctly combines filter and projection EXPECT_EQ(ctx.col_num(), 3) @@ -1354,7 +1454,7 @@ TEST_F(ParquetTest, TestMultiFile_ExplicitPaths) { std::string(PARQUET_TEST_DIR) + "/test_multi_1.parquet", std::string(PARQUET_TEST_DIR) + "/test_multi_2.parquet"}; fileSchema.format = "parquet"; - fileSchema.options = {{"batch_read", "false"}}; + fileSchema.options = {}; reader::ExternalSchema externalSchema; externalSchema.entry = entrySchema; @@ -1362,9 +1462,14 @@ TEST_F(ParquetTest, TestMultiFile_ExplicitPaths) { sharedState->schema = std::move(externalSchema); auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 1); EXPECT_EQ(ctx.row_num(), 30) << "Extension should correctly read and " @@ -1436,13 +1541,17 @@ TEST_F(ParquetTest, TestParquetExportWriter) { // Read it back and verify auto sharedState = createSharedState( "export_writer_test.parquet", {"id", "name", "value"}, - {createInt64Type(), createStringType(), createDoubleType()}, - {{"batch_read", "false"}}); + {createInt64Type(), createStringType(), createDoubleType()}, {}); auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 3); @@ -1495,14 +1604,19 @@ TEST_F(ParquetTest, TestParquetExportWithNulls) { ASSERT_TRUE(std::filesystem::exists(export_path)); // Read back and verify - auto sharedState = createSharedState( - "export_nulls_test.parquet", {"id", "name"}, - {createInt64Type(), createStringType()}, {{"batch_read", "false"}}); + auto sharedState = + createSharedState("export_nulls_test.parquet", {"id", "name"}, + {createInt64Type(), createStringType()}, {}); auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 2); EXPECT_EQ(ctx.row_num(), 3); @@ -1700,24 +1814,34 @@ TEST_F(ParquetTest, TestParquetExportWithCompressionOptions) { << "ZSTD compressed file should be smaller than uncompressed"; // Verify both files are readable - auto sharedState_zstd = createSharedState( - "export_zstd.parquet", {"id", "name"}, - {createInt64Type(), createStringType()}, {{"batch_read", "false"}}); + auto sharedState_zstd = + createSharedState("export_zstd.parquet", {"id", "name"}, + {createInt64Type(), createStringType()}, {}); auto reader_zstd = createParquetReader(sharedState_zstd); - auto localState_zstd = std::make_shared(); + execution::Context ctx_zstd; - reader_zstd->read(localState_zstd, ctx_zstd); + { + auto supplier = reader_zstd->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx_zstd.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx_zstd.row_num(), 100); - auto sharedState_none = createSharedState( - "export_none.parquet", {"id", "name"}, - {createInt64Type(), createStringType()}, {{"batch_read", "false"}}); + auto sharedState_none = + createSharedState("export_none.parquet", {"id", "name"}, + {createInt64Type(), createStringType()}, {}); auto reader_none = createParquetReader(sharedState_none); - auto localState_none = std::make_shared(); + execution::Context ctx_none; - reader_none->read(localState_none, ctx_none); + { + auto supplier = reader_none->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx_none.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx_none.row_num(), 100); } @@ -1793,14 +1917,18 @@ TEST_F(ParquetTest, TestParquetExportWithRowGroupSize) { ASSERT_TRUE(std::filesystem::exists(export_path)); // Verify file is readable - auto sharedState = - createSharedState("export_rowgroup.parquet", {"id"}, {createInt64Type()}, - {{"batch_read", "false"}}); + auto sharedState = createSharedState("export_rowgroup.parquet", {"id"}, + {createInt64Type()}, {}); auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.row_num(), 100); } @@ -1891,24 +2019,34 @@ TEST_F(ParquetTest, TestParquetExportWithDictionaryEncoding) { "smaller for low-cardinality strings"; // Verify both files are readable - auto sharedState_dict = createSharedState( - "export_dict_enabled.parquet", {"id", "category"}, - {createInt64Type(), createStringType()}, {{"batch_read", "false"}}); + auto sharedState_dict = + createSharedState("export_dict_enabled.parquet", {"id", "category"}, + {createInt64Type(), createStringType()}, {}); auto reader_dict = createParquetReader(sharedState_dict); - auto localState_dict = std::make_shared(); + execution::Context ctx_dict; - reader_dict->read(localState_dict, ctx_dict); + { + auto supplier = reader_dict->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx_dict.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx_dict.row_num(), num_rows); - auto sharedState_nodict = createSharedState( - "export_dict_disabled.parquet", {"id", "category"}, - {createInt64Type(), createStringType()}, {{"batch_read", "false"}}); + auto sharedState_nodict = + createSharedState("export_dict_disabled.parquet", {"id", "category"}, + {createInt64Type(), createStringType()}, {}); auto reader_nodict = createParquetReader(sharedState_nodict); - auto localState_nodict = std::make_shared(); + execution::Context ctx_nodict; - reader_nodict->read(localState_nodict, ctx_nodict); + { + auto supplier = reader_nodict->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx_nodict.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx_nodict.row_num(), num_rows); } @@ -1974,13 +2112,17 @@ TEST_F(ParquetTest, TestParquetExportWithDateAndTimestamp) { // Verify file is readable auto sharedState = createSharedState( "export_datetime.parquet", {"id", "created_date", "updated_timestamp"}, - {createInt64Type(), createDateType(), createTimestampType()}, - {{"batch_read", "false"}}); + {createInt64Type(), createDateType(), createTimestampType()}, {}); auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.row_num(), num_rows); } @@ -2370,15 +2512,20 @@ TEST_F(ParquetTest, TestParquetNonExistentColumnThrows) { std::vector> columnTypes = { createInt64Type(), createStringType(), createDoubleType()}; - auto sharedState = createSharedState("test_nonexist.parquet", columnNames, - columnTypes, {{"batch_read", "false"}}); + auto sharedState = + createSharedState("test_nonexist.parquet", columnNames, columnTypes, {}); auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - EXPECT_THROW(reader->read(localState, ctx), - exception::SchemaMismatchException); + EXPECT_THROW( + [&] { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + }(), + exception::SchemaMismatchException); } // ============================================================================= @@ -2477,13 +2624,17 @@ TEST_F(ParquetTest, TestIntegration_ReadZstdWithSmallBufferedStream) { auto sharedState = createSharedState( "test_zstd_small_buffer.parquet", {"id", "name", "value"}, {createInt64Type(), createStringType(), createDoubleType()}, - {{"batch_read", "true"}, {"batch_size", "4096"}}); + {{"batch_size", "4096"}}); auto reader = createParquetReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - ASSERT_NO_THROW(reader->read(localState, ctx)) - << "Reading ZSTD parquet with small buffered stream must not crash"; + ASSERT_NO_THROW([&] { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + }()) << "Reading ZSTD parquet with small buffered stream must not crash"; int64_t totalRows = 0; for (size_t i = 0; i < ctx.chunk_num(); ++i) { diff --git a/include/neug/common/columns/array_columns.h b/include/neug/common/columns/array_columns.h index bbd4921ea..606907a02 100644 --- a/include/neug/common/columns/array_columns.h +++ b/include/neug/common/columns/array_columns.h @@ -56,6 +56,9 @@ class ContextArrayColumn : public IContextColumn { return ContextColumnType::kValue; } + std::shared_ptr union_col( + std::shared_ptr other) const override; + const DataType& elem_type() const override { return type_; } std::shared_ptr shuffle( diff --git a/include/neug/common/columns/list_columns.h b/include/neug/common/columns/list_columns.h index f56a37a37..c1d31bfe3 100644 --- a/include/neug/common/columns/list_columns.h +++ b/include/neug/common/columns/list_columns.h @@ -50,6 +50,9 @@ class ListColumn : public IContextColumn { std::shared_ptr optional_shuffle( const sel_vec_t& offsets) const override; + std::shared_ptr union_col( + std::shared_ptr other) const override; + const DataType& elem_type() const override { return type_; } Value get_elem(size_t idx) const override { if (is_optional_ && !valids_[idx]) { diff --git a/include/neug/compiler/function/import/csv_read_function.h b/include/neug/compiler/function/import/csv_read_function.h index 9cbc25743..a469c8689 100644 --- a/include/neug/compiler/function/import/csv_read_function.h +++ b/include/neug/compiler/function/import/csv_read_function.h @@ -35,8 +35,8 @@ struct CSVReadFunction { static function_set getFunctionSet() { auto typeIDs = std::vector{common::DataTypeId::kVarchar}; - auto readFunction = std::make_unique(name, typeIDs); - readFunction->execFunc = execFunc; + auto readFunction = + std::make_unique(name, typeIDs, supplierFunc); readFunction->sniffFunc = sniffFunc; function_set functionSet; functionSet.push_back(std::move(readFunction)); @@ -112,7 +112,7 @@ struct CSVReadFunction { } } - static execution::Context execFunc( + static std::unique_ptr createReader( std::shared_ptr state) { validateAndConvertExecOptions(state); const auto& vfs = neug::main::MetadataRegistry::getVFS(); @@ -128,10 +128,12 @@ struct CSVReadFunction { auto optionsBuilder = std::make_unique(state); auto reader = std::make_unique(state, std::move(optionsBuilder)); - execution::Context ctx; - auto localState = std::make_shared(); - reader->read(localState, ctx); - return ctx; + return reader; + } + + static std::shared_ptr supplierFunc( + std::shared_ptr state) { + return createReader(std::move(state))->getDataChunkSupplier(); } static std::shared_ptr sniffFunc( diff --git a/include/neug/compiler/function/import/json_read_function.h b/include/neug/compiler/function/import/json_read_function.h index 95f10fed7..9cc56d263 100644 --- a/include/neug/compiler/function/import/json_read_function.h +++ b/include/neug/compiler/function/import/json_read_function.h @@ -37,15 +37,15 @@ struct JsonReadFunction { static function_set getFunctionSet() { auto typeIDs = std::vector{common::DataTypeId::kVarchar}; - auto readFunction = std::make_unique(name, typeIDs); - readFunction->execFunc = jsonExecFunc; + auto readFunction = + std::make_unique(name, typeIDs, supplierFunc); readFunction->sniffFunc = jsonSniffFunc; function_set functionSet; functionSet.push_back(std::move(readFunction)); return functionSet; } - static execution::Context jsonExecFunc( + static std::unique_ptr createReader( std::shared_ptr state) { const auto& vfs = neug::main::MetadataRegistry::getVFS(); const auto& fs = vfs->Provide(state->schema.file); @@ -61,10 +61,12 @@ struct JsonReadFunction { std::make_unique(state, true); auto reader = std::make_unique(state, std::move(optionsBuilder)); - execution::Context ctx; - auto localState = std::make_shared(); - reader->read(localState, ctx); - return ctx; + return reader; + } + + static std::shared_ptr supplierFunc( + std::shared_ptr state) { + return createReader(std::move(state))->getDataChunkSupplier(); } static std::shared_ptr jsonSniffFunc( @@ -105,15 +107,15 @@ struct JsonLReadFunction { static function_set getFunctionSet() { auto typeIDs = std::vector{common::DataTypeId::kVarchar}; - auto readFunction = std::make_unique(name, typeIDs); - readFunction->execFunc = jsonLExecFunc; + auto readFunction = + std::make_unique(name, typeIDs, supplierFunc); readFunction->sniffFunc = jsonLSniffFunc; function_set functionSet; functionSet.push_back(std::move(readFunction)); return functionSet; } - static execution::Context jsonLExecFunc( + static std::unique_ptr createReader( std::shared_ptr state) { const auto& vfs = neug::main::MetadataRegistry::getVFS(); const auto& fs = vfs->Provide(state->schema.file); @@ -129,10 +131,12 @@ struct JsonLReadFunction { std::make_unique(state, false); auto reader = std::make_unique(state, std::move(optionsBuilder)); - execution::Context ctx; - auto localState = std::make_shared(); - reader->read(localState, ctx); - return ctx; + return reader; + } + + static std::shared_ptr supplierFunc( + std::shared_ptr state) { + return createReader(std::move(state))->getDataChunkSupplier(); } static std::shared_ptr jsonLSniffFunc( diff --git a/include/neug/compiler/function/read_function.h b/include/neug/compiler/function/read_function.h index 61bcbb380..c3e99dcfa 100644 --- a/include/neug/compiler/function/read_function.h +++ b/include/neug/compiler/function/read_function.h @@ -23,18 +23,16 @@ #include #include #include "neug/compiler/function/table/table_function.h" -#include "neug/execution/common/context.h" -#include "neug/execution/execute/ops/batch/batch_update_utils.h" #include "neug/utils/io/read/common/schema.h" #include "neug/utils/io/reader.h" namespace neug { +class IDataChunkSupplier; namespace function { -// The exec function invoked by data source operators to load data from external -// data sources. -using read_exec_func_t = std::function state)>; +// Required factory for a single execution's incremental reader. +using read_supplier_func_t = std::function( + std::shared_ptr)>; // The function used to sniff/infer file column names and their types from // external data sources. @@ -42,11 +40,18 @@ using read_sniff_func_t = std::function( const reader::FileSchema& schema)>; struct ReadFunction : public TableFunction { - read_exec_func_t execFunc = nullptr; + read_supplier_func_t supplierFunc; read_sniff_func_t sniffFunc = nullptr; - ReadFunction(std::string name, std::vector inputTypes) - : TableFunction{std::move(name), std::move(inputTypes)} {} + ReadFunction(std::string name, std::vector inputTypes, + read_supplier_func_t supplier) + : TableFunction{std::move(name), std::move(inputTypes)}, + supplierFunc(std::move(supplier)) { + if (!supplierFunc) { + THROW_INVALID_ARGUMENT_EXCEPTION( + "ReadFunction requires a supplier factory"); + } + } }; } // namespace function -} // namespace neug \ No newline at end of file +} // namespace neug diff --git a/include/neug/execution/common/context.h b/include/neug/execution/common/context.h index 3d3e0047c..7984ed4d8 100644 --- a/include/neug/execution/common/context.h +++ b/include/neug/execution/common/context.h @@ -31,7 +31,9 @@ class StorageReadInterface; namespace execution { /** - * @brief Context is a multi-chunk container passed between operators. + * @brief Context is a materialized multi-chunk container for internal + * algorithms and public query results. Operators exchange Stream + * instead. * * A Context holds one or more ContextChunks (DataChunk + head pairs) that * share the same schema. Operators iterate chunks via `apply_chunks`, whose diff --git a/include/neug/execution/common/stream.h b/include/neug/execution/common/stream.h new file mode 100644 index 000000000..905192396 --- /dev/null +++ b/include/neug/execution/common/stream.h @@ -0,0 +1,253 @@ +/** 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 "neug/execution/common/context.h" + +namespace neug::execution { + +// Result layout, independent of batch contents and execution progress. +// The aliases select output columns in result order, including empty results. +struct StreamMetadata { + std::vector output_columns; +}; + +// A single-consumer, synchronous pull stream. Construction does not read rows. +// The stream yields T directly. Execution uses ContextChunk, which already +// owns the DataChunk and anonymous head. Metadata also describes empty streams. +// EOF and errors are terminal; an empty batch is NOT EOF. +template +class Stream { + public: + using NextResult = result>; + using Pull = std::function; + + Stream() = default; + explicit Stream(Pull pull, StreamMetadata metadata = {}) + : metadata_(std::move(metadata)), pull_(std::move(pull)) {} + Stream(Stream&&) = default; + Stream& operator=(Stream&&) = default; + Stream(const Stream&) = delete; + Stream& operator=(const Stream&) = delete; + + NextResult Next() { + if (error_) { + return tl::unexpected(*error_); + } + if (!pull_) { + return std::optional{}; + } + auto output = PullOne(); + if (!output) { + error_ = output.error(); + pull_ = nullptr; + } else if (!*output) { + pull_ = nullptr; + } + return output; + } + + const StreamMetadata& metadata() const { return metadata_; } + void set_metadata(StreamMetadata metadata) { + metadata_ = std::move(metadata); + } + + private: + NextResult PullOne() { + NextResult output = std::optional{}; + TRY_HANDLE_ALL_WITH_EXCEPTION( + NextResult, [&]() { return pull_(); }, + [&](const Status& status) { output = tl::unexpected(status); }, + [&](NextResult&& batch) { output = std::move(batch); }); + return output; + } + + StreamMetadata metadata_; + Pull pull_; + std::optional error_; +}; + +// An execution error is observed through Next(), just like a read error. +template +Stream error_stream(Status error) { + return Stream( + [error = std::move(error)]() -> + typename Stream::NextResult { return tl::unexpected(error); }); +} + +// Own execution state until first demand. Initialization and its exceptions +// run inside Stream::Next's error boundary, exactly once. +template +Stream defer_stream(Stream input, + Initialize initialize) { + auto metadata = input.metadata(); + struct State { + Stream input; + std::optional> output; + }; + auto state = std::make_shared(State{std::move(input), std::nullopt}); + return Stream( + [state, initialize = std::move( + initialize)]() mutable -> Stream::NextResult { + if (!state->output) { + state->output.emplace(initialize(std::move(state->input))); + } + return state->output->Next(); + }, + std::move(metadata)); +} + +// Put a batch pulled for initialization back in front of its remaining input. +inline Stream prepend_chunk(std::optional first, + Stream input) { + if (!first) { + return std::move(input); + } + auto metadata = input.metadata(); + auto pending = + std::make_shared>(std::move(first)); + auto upstream = std::make_shared>(std::move(input)); + return Stream( + [pending, upstream]() -> Stream::NextResult { + if (*pending) { + auto chunk = std::move(*pending); + pending->reset(); + return chunk; + } + return upstream->Next(); + }, + std::move(metadata)); +} + +// Exactly one upstream pull and one kernel invocation per downstream pull. +template +Stream map_chunks(Stream input, + Transform transform) { + auto metadata = input.metadata(); + auto upstream = std::make_shared>(std::move(input)); + return Stream( + [upstream, transform = std::move( + transform)]() mutable -> Stream::NextResult { + GS_AUTO(next, upstream->Next()); + if (!next) { + return std::optional{}; + } + GS_AUTO(output, transform(std::move(*next))); + return std::optional(std::move(output)); + }, + std::move(metadata)); +} + +// Invoke a producer on first demand, without an intermediate Context. +template +Stream generate_chunk(Producer producer, + StreamMetadata metadata = {}) { + return Stream( + [producer = std::move(producer), + done = false]() mutable -> Stream::NextResult { + if (done) { + return std::optional{}; + } + done = true; + GS_AUTO(chunk, producer()); + return std::optional(std::move(chunk)); + }, + std::move(metadata)); +} + +// Explicit global-input boundary. Row-local operators never collect input. +inline result collect_chunk(Stream input) { + std::optional accumulated; + while (true) { + GS_AUTO(next, input.Next()); + if (!next) { + return accumulated ? std::move(*accumulated) : ContextChunk{}; + } + ContextChunk chunk = std::move(*next); + if (accumulated) { + *accumulated = accumulated->union_with(chunk); + } else { + accumulated = std::move(chunk); + } + } +} + +template +Stream reduce_stream(Stream input, Reduce reduce) { + auto metadata = input.metadata(); + auto upstream = std::make_shared>(std::move(input)); + return generate_chunk( + [upstream, reduce = std::move(reduce)]() mutable -> result { + GS_AUTO(chunk, collect_chunk(std::move(*upstream))); + return reduce(std::move(chunk)); + }, + std::move(metadata)); +} + +// Buffer only when an operator must replay its input or stabilize it before +// mutations. Batches retain their boundaries and share column ownership. +inline result> collect_batches( + Stream input) { + std::vector chunks; + while (true) { + GS_AUTO(next, input.Next()); + if (!next) { + return chunks; + } + chunks.push_back(std::move(*next)); + } +} + +inline Stream stream_from_batches( + std::vector chunks, StreamMetadata metadata = {}) { + auto batches = std::make_shared>(std::move(chunks)); + return Stream( + [batches, + index = size_t{0}]() mutable -> Stream::NextResult { + if (index == batches->size()) { + return std::optional{}; + } + return std::optional(std::move((*batches)[index++])); + }, + std::move(metadata)); +} + +inline Stream stream_from_context(Context ctx) { + return stream_from_batches(std::move(ctx.chunks()), + StreamMetadata{std::move(ctx.tag_ids)}); +} + +inline result materialize(Stream stream) { + Context ctx; + ctx.tag_ids = stream.metadata().output_columns; + while (true) { + auto next = stream.Next(); + if (!next) { + return tl::unexpected(next.error()); + } + if (!*next) { + return ctx; + } + auto& batch = **next; + ctx.append_chunk(std::move(batch)); + } +} + +} // namespace neug::execution diff --git a/include/neug/execution/execute/operator.h b/include/neug/execution/execute/operator.h index ed99bec63..79a009c0e 100644 --- a/include/neug/execution/execute/operator.h +++ b/include/neug/execution/execute/operator.h @@ -17,8 +17,8 @@ #include #include -#include "neug/execution/common/context.h" #include "neug/execution/common/params_map.h" +#include "neug/execution/common/stream.h" #include "neug/execution/utils/opr_timer.h" #include "neug/generated/proto/plan/physical.pb.h" #include "neug/storages/graph/graph_interface.h" @@ -34,9 +34,10 @@ class IOperator { virtual std::string get_operator_name() const = 0; - virtual neug::result Eval(IStorageInterface& graph, - const ParamsMap& params, Context&& ctx, - OprTimer* timer) = 0; + virtual Stream Eval(IStorageInterface& graph, + const ParamsMap& params, + Stream&& input, + OprTimer* timer) = 0; virtual void build_explain_children(OprTimer* parent_timer, const ParamsMap& params, 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 58ee9686e..f3a2a3f66 100644 --- a/include/neug/execution/execute/ops/batch/batch_update_utils.h +++ b/include/neug/execution/execute/ops/batch/batch_update_utils.h @@ -17,7 +17,8 @@ #include #include "neug/common/types/graph_types.h" -#include "neug/execution/common/context.h" +#include "neug/execution/common/stream.h" +#include "neug/storages/loader/loader_utils.h" #include "neug/utils/property/types.h" namespace physical { @@ -46,9 +47,24 @@ std::string edge_to_json_string(const EdgeRecord& edge, std::string path_to_json_string(Path& path, const StorageReadInterface& graph); -std::shared_ptr create_data_chunk_supplier( - const Context& ctx, - const std::vector>& prop_mappings); +class StreamChunkSupplier final : public IDataChunkSupplier { + public: + StreamChunkSupplier(Stream stream, + std::vector> mappings); + std::shared_ptr GetNextChunk() override; + int64_t RowNum() const override { return -1; } + const Status& status() const { return status_; } + size_t rows_read() const { return rows_read_; } + + private: + Stream stream_; + std::vector> mappings_; + Status status_ = Status::OK(); + size_t rows_read_ = 0; +}; + +// Preserve COPY result cardinality without retaining its input payload. +Stream batch_insert_result(size_t rows); std::vector match_files_with_pattern(const std::string& file_path); diff --git a/include/neug/execution/execute/ops/edge_column_rebuild.h b/include/neug/execution/execute/ops/edge_column_rebuild.h index c5858d7a4..602e42b81 100644 --- a/include/neug/execution/execute/ops/edge_column_rebuild.h +++ b/include/neug/execution/execute/ops/edge_column_rebuild.h @@ -43,7 +43,7 @@ struct EdgeColumnSnapshots { }; EdgeColumnSnapshots CaptureEdgeColumnsForRefresh( - StorageUpdateInterface& graph, Context& ctx, + StorageUpdateInterface& graph, std::vector& chunks, const std::set& affected_labels); void RefreshEdgeColumns(StorageUpdateInterface& graph, diff --git a/include/neug/execution/execute/ops/retrieve/gds_algo.h b/include/neug/execution/execute/ops/retrieve/gds_algo.h index 1b483f3ef..791d9d3a3 100644 --- a/include/neug/execution/execute/ops/retrieve/gds_algo.h +++ b/include/neug/execution/execute/ops/retrieve/gds_algo.h @@ -32,10 +32,9 @@ class GDSAlgoOpr : public IOperator { std::string get_operator_name() const override { return "GDSAlgoOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override; private: std::unique_ptr algo_input_; diff --git a/include/neug/execution/execute/pipeline.h b/include/neug/execution/execute/pipeline.h index f65093a26..91acffe08 100644 --- a/include/neug/execution/execute/pipeline.h +++ b/include/neug/execution/execute/pipeline.h @@ -31,6 +31,12 @@ class Pipeline { : operators_(std::move(operators)) {} ~Pipeline() = default; + // Caller keeps this pipeline, storage and timer alive until the stream is + // consumed or destroyed. Params are captured by value by lazy operators. + Stream ExecuteStream(IStorageInterface& graph, + Stream input, + const ParamsMap& params, OprTimer* timer); + neug::result Execute(IStorageInterface& graph, Context&& ctx, const ParamsMap& params, OprTimer* timer); diff --git a/include/neug/execution/utils/opr_timer.h b/include/neug/execution/utils/opr_timer.h index 3631adde4..daa8297dd 100644 --- a/include/neug/execution/utils/opr_timer.h +++ b/include/neug/execution/utils/opr_timer.h @@ -65,6 +65,8 @@ class OprTimer { void record(const TimerUnit& tu) { time_ += tu.elapsed(); } + void add_elapsed(double seconds) { time_ += seconds; } + void add_num_tuples(uint64_t num) { numTuples_ += num; } ~OprTimer() = default; diff --git a/include/neug/utils/io/read/common/options.h b/include/neug/utils/io/read/common/options.h index be885d11d..bd133d4dd 100644 --- a/include/neug/utils/io/read/common/options.h +++ b/include/neug/utils/io/read/common/options.h @@ -145,7 +145,6 @@ struct CSVParseOptions { struct ReadOptions { Option use_threads = Option::BoolOption("parallel", false); - Option batch_read = Option::BoolOption("batch_read", true); Option batch_size = Option::Int64Option("batch_size", 1 << 20); Option autogenerate_column_names = diff --git a/include/neug/utils/io/read/common/read_state.h b/include/neug/utils/io/read/common/read_state.h index 7af2dc153..63cbe22b2 100644 --- a/include/neug/utils/io/read/common/read_state.h +++ b/include/neug/utils/io/read/common/read_state.h @@ -31,30 +31,6 @@ class Expression; namespace neug { namespace reader { -struct ReadLocalState { - virtual ~ReadLocalState() = default; - - template - TARGET& cast() { - return common::neug_dynamic_cast(*this); - } - - template - TARGET* ptrCast() { - return common::neug_dynamic_cast(this); - } - - template - const TARGET& constCast() const { - return common::neug_dynamic_cast(*this); - } - - template - const TARGET* constPtrCast() const { - return common::neug_dynamic_cast(this); - } -}; - struct ReadSharedState { ExternalSchema schema; std::vector projectColumns; diff --git a/include/neug/utils/io/read/common/row_expression_filter.h b/include/neug/utils/io/read/common/row_expression_filter.h index 4a4edda9a..e1ecbadea 100644 --- a/include/neug/utils/io/read/common/row_expression_filter.h +++ b/include/neug/utils/io/read/common/row_expression_filter.h @@ -25,7 +25,6 @@ #include "neug/generated/proto/plan/expr.pb.h" namespace neug { -class IDataChunkSupplier; namespace reader { /// Evaluates a file predicate using the same expressions as query execution. @@ -55,9 +54,6 @@ DataChunk project_chunk(const DataChunk& input, const std::vector& column_names, const std::vector& project_columns); -DataChunk read_all_chunks( - const std::vector>& suppliers); - // Merge dense reader chunks in order, preserving column types and NULLs. // Ignore null/columnless chunks, retain empty schemas, and reject mismatched // column counts, types, or lengths before allocating output columns. diff --git a/include/neug/utils/io/read/csv/csv_reader.h b/include/neug/utils/io/read/csv/csv_reader.h index bf4f66d74..5886187d6 100644 --- a/include/neug/utils/io/read/csv/csv_reader.h +++ b/include/neug/utils/io/read/csv/csv_reader.h @@ -18,7 +18,6 @@ #include #include -#include "neug/execution/common/context.h" #include "neug/utils/io/read/common/options.h" #include "neug/utils/io/read/common/read_state.h" #include "neug/utils/io/read/csv/csv_read_config.h" @@ -28,10 +27,6 @@ namespace neug { class IDataChunkSupplier; -namespace execution { -class Context; -} - namespace reader { class CsvReader { @@ -40,19 +35,11 @@ class CsvReader { std::unique_ptr optionsBuilder); ~CsvReader(); - void read(std::shared_ptr localState, - execution::Context& ctx); + std::shared_ptr getDataChunkSupplier(); result> inferSchema(); private: - void full_read( - const std::vector>& suppliers, - execution::Context& output, const CsvReadConfig& output_config); - void batch_read( - const std::vector>& suppliers, - execution::Context& output); - std::shared_ptr sharedState_; std::unique_ptr optionsBuilder_; }; diff --git a/include/neug/utils/io/read/json/json_reader.h b/include/neug/utils/io/read/json/json_reader.h index b88011352..266757f41 100644 --- a/include/neug/utils/io/read/json/json_reader.h +++ b/include/neug/utils/io/read/json/json_reader.h @@ -18,7 +18,6 @@ #include #include -#include "neug/execution/common/context.h" #include "neug/utils/io/read/common/options.h" #include "neug/utils/io/read/common/read_state.h" #include "neug/utils/io/read/json/json_read_config.h" @@ -28,10 +27,6 @@ namespace neug { class IDataChunkSupplier; -namespace execution { -class Context; -} - namespace reader { class JsonReader { @@ -40,19 +35,11 @@ class JsonReader { std::unique_ptr optionsBuilder); ~JsonReader(); - void read(std::shared_ptr localState, - execution::Context& ctx); + std::shared_ptr getDataChunkSupplier(); result> inferSchema(); private: - void full_read( - const std::vector>& suppliers, - execution::Context& output, const JsonReadConfig& output_config); - void batch_read( - const std::vector>& suppliers, - execution::Context& output); - std::shared_ptr sharedState_; std::unique_ptr optionsBuilder_; }; diff --git a/src/common/columns/array_columns.cc b/src/common/columns/array_columns.cc index 3ea97fa20..d5043e7b0 100644 --- a/src/common/columns/array_columns.cc +++ b/src/common/columns/array_columns.cc @@ -20,6 +20,23 @@ namespace neug { +std::shared_ptr ContextArrayColumn::union_col( + std::shared_ptr other) const { + if (!other || other->elem_type() != type_) { + THROW_INVALID_ARGUMENT_EXCEPTION( + "Cannot merge columns with different types"); + } + ContextArrayColumnBuilder builder(type_); + builder.reserve(size() + other->size()); + for (size_t row = 0; row < size(); ++row) { + builder.push_back_elem(get_elem(row)); + } + for (size_t row = 0; row < other->size(); ++row) { + builder.push_back_elem(other->get_elem(row)); + } + return builder.finish(); +} + std::pair, sel_vec_t> ContextArrayColumn::unfold() const { sel_vec_t offsets; diff --git a/src/common/columns/list_columns.cc b/src/common/columns/list_columns.cc index 10a9ce2b7..346204476 100644 --- a/src/common/columns/list_columns.cc +++ b/src/common/columns/list_columns.cc @@ -22,6 +22,23 @@ namespace neug { +std::shared_ptr ListColumn::union_col( + std::shared_ptr other) const { + if (!other || other->elem_type() != type_) { + THROW_INVALID_ARGUMENT_EXCEPTION( + "Cannot merge columns with different types"); + } + ListColumnBuilder builder(elem_type_); + builder.reserve(size() + other->size()); + for (size_t row = 0; row < size(); ++row) { + builder.push_back_elem(get_elem(row)); + } + for (size_t row = 0; row < other->size(); ++row) { + builder.push_back_elem(other->get_elem(row)); + } + return builder.finish(); +} + std::pair, sel_vec_t> ListColumn::unfold() const { switch (elem_type_.id()) { diff --git a/src/execution/execute/ops/admin/checkpoint.cc b/src/execution/execute/ops/admin/checkpoint.cc index aab207012..2f9632fb9 100644 --- a/src/execution/execute/ops/admin/checkpoint.cc +++ b/src/execution/execute/ops/admin/checkpoint.cc @@ -26,20 +26,27 @@ class CheckpointOpr : public IOperator { CheckpointOpr() = default; ~CheckpointOpr() override = default; std::string get_operator_name() const override { return "CheckpointOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override; }; -neug::result CheckpointOpr::Eval(IStorageInterface& graph_interface, - const ParamsMap& params, - Context&& ctx, OprTimer* timer) { - (void) graph_interface; - (void) params; - (void) ctx; - (void) timer; - RETURN_ERROR(neug::Status( - neug::StatusCode::ERR_ILLEGAL_OPERATION, - "CHECKPOINT must be executed by the database checkpoint executor")); +Stream CheckpointOpr::Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + (void) graph_interface; + (void) params; + (void) input; + (void) timer; + return error_stream(neug::Status( + neug::StatusCode::ERR_ILLEGAL_OPERATION, + "CHECKPOINT must be executed by the database checkpoint executor")); + }); } neug::result CheckpointOprBuilder::Build( diff --git a/src/execution/execute/ops/admin/extension.cc b/src/execution/execute/ops/admin/extension.cc index 610e172ec..36837b331 100644 --- a/src/execution/execute/ops/admin/extension.cc +++ b/src/execution/execute/ops/admin/extension.cc @@ -83,8 +83,9 @@ class ExtensionInstallOpr : public IOperator { std::string get_operator_name() const override { return "ExtensionInstallOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override; private: std::string extension_name_; @@ -96,8 +97,9 @@ class ExtensionLoadOpr : public IOperator { : extension_name_(std::move(extension_name)) {} ~ExtensionLoadOpr() override = default; std::string get_operator_name() const override { return "ExtensionLoadOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override; private: std::string extension_name_; @@ -111,64 +113,82 @@ class ExtensionUninstallOpr : public IOperator { std::string get_operator_name() const override { return "ExtensionUninstallOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override; private: std::string extension_name_; }; -neug::result ExtensionInstallOpr::Eval(IStorageInterface& graph, - const ParamsMap& params, - Context&& ctx, - OprTimer* timer) { - LOG(INFO) << "[Admin Pipeline] Executing ExtensionInstall for: " - << extension_name_; - - checkDeprecatedExtension(extension_name_); - - auto status = neug::extension::install_extension(extension_name_); - if (!status.ok()) { - THROW_EXCEPTION_WITH_FILE_LINE("Install failed: " + status.ToString() + - "; "); - } - return neug::result(std::move(ctx)); +Stream ExtensionInstallOpr::Eval(IStorageInterface& graph, + const ParamsMap& params, + Stream&& input, + OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + LOG(INFO) << "[Admin Pipeline] Executing ExtensionInstall for: " + << extension_name_; + + checkDeprecatedExtension(extension_name_); + + auto status = neug::extension::install_extension(extension_name_); + if (!status.ok()) { + THROW_EXCEPTION_WITH_FILE_LINE( + "Install failed: " + status.ToString() + "; "); + } + return std::move(input); + }); } -neug::result ExtensionLoadOpr::Eval(IStorageInterface& graph, - const ParamsMap& params, - Context&& ctx, OprTimer* timer) { - LOG(INFO) << "[Admin Pipeline] Executing ExtensionLoad for: " - << extension_name_; - - checkDeprecatedExtension(extension_name_); - - auto* index_ddl = dynamic_cast(&graph); - if (index_ddl) { - auto activated = index_ddl->ActivateIndexes(); - if (!activated) { - RETURN_ERROR(activated.error()); - } - } else { - LOG(WARNING) << "[Admin Pipeline] Current storage interface does not " - "support index DDL; skipping pending index activation"; - } - return neug::result(std::move(ctx)); +Stream ExtensionLoadOpr::Eval(IStorageInterface& graph, + const ParamsMap& params, + Stream&& input, + OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + LOG(INFO) << "[Admin Pipeline] Executing ExtensionLoad for: " + << extension_name_; + + checkDeprecatedExtension(extension_name_); + + auto* index_ddl = dynamic_cast(&graph); + if (index_ddl) { + auto activated = index_ddl->ActivateIndexes(); + if (!activated) { + return error_stream(activated.error()); + } + } else { + LOG(WARNING) + << "[Admin Pipeline] Current storage interface does not " + "support index DDL; skipping pending index activation"; + } + return std::move(input); + }); } -neug::result ExtensionUninstallOpr::Eval(IStorageInterface& graph, - const ParamsMap& params, - Context&& ctx, - OprTimer* timer) { - LOG(INFO) << "[Admin Pipeline] Executing ExtensionUninstall for: " - << extension_name_; - - auto status = neug::extension::uninstall_extension(extension_name_); - if (!status.ok()) { - THROW_EXCEPTION_WITH_FILE_LINE("Uninstall failed: " + status.ToString() + - "; "); - } - return neug::result(std::move(ctx)); +Stream ExtensionUninstallOpr::Eval(IStorageInterface& graph, + const ParamsMap& params, + Stream&& input, + OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + LOG(INFO) << "[Admin Pipeline] Executing ExtensionUninstall for: " + << extension_name_; + + auto status = neug::extension::uninstall_extension(extension_name_); + if (!status.ok()) { + THROW_EXCEPTION_WITH_FILE_LINE( + "Uninstall failed: " + status.ToString() + "; "); + } + return std::move(input); + }); } // Builders diff --git a/src/execution/execute/ops/batch/batch_delete_edge.cc b/src/execution/execute/ops/batch/batch_delete_edge.cc index 6555f65f1..bce82f9ce 100644 --- a/src/execution/execute/ops/batch/batch_delete_edge.cc +++ b/src/execution/execute/ops/batch/batch_delete_edge.cc @@ -34,8 +34,9 @@ class BatchDeleteEdgeOpr : public IOperator { return "BatchDeleteEdgeOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override; private: std::vector>> @@ -43,88 +44,106 @@ class BatchDeleteEdgeOpr : public IOperator { std::vector edge_bindings_; }; -neug::result BatchDeleteEdgeOpr::Eval( - IStorageInterface& graph_interface, const ParamsMap& params, Context&& ctx, - OprTimer* timer) { - auto& graph = dynamic_cast(graph_interface); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - size_t binding_size = edge_bindings_.size(); - for (size_t i = 0; i < binding_size; i++) { - int32_t alias = edge_bindings_[i]; - auto& edge_triplets = edge_triplets_[i]; - auto edge_column = - std::dynamic_pointer_cast(chunk.get(alias)); - if (edge_triplets.size() == 1) { - label_t src_v_label = std::get<0>(edge_triplets[0]); - label_t dst_v_label = std::get<1>(edge_triplets[0]); - label_t edge_label = std::get<2>(edge_triplets[0]); - LabelTriplet request_triplet = - LabelTriplet(src_v_label, dst_v_label, edge_label); - size_t edge_size = edge_column->size(); - auto oe_view = graph.GetGenericOutgoingGraphView( - src_v_label, dst_v_label, edge_label); - auto ie_view = graph.GetGenericIncomingGraphView( - dst_v_label, src_v_label, edge_label); - auto edge_prop_types = graph.schema().get_edge_properties( - src_v_label, dst_v_label, edge_label); - std::vector> oe_to_delete, ie_to_delete; - oe_to_delete.reserve(edge_size); - ie_to_delete.reserve(edge_size); - for (size_t j = 0; j < edge_size; j++) { - auto record = edge_column->get_edge(j); - if (record.label == request_triplet) { - auto offset_pair = record_to_csr_offset_pair( - oe_view, ie_view, record, edge_prop_types); - oe_to_delete.emplace_back(record.src, offset_pair.first); - ie_to_delete.emplace_back(record.dst, offset_pair.second); - } - } - RETURN_STATUS_ERROR_IF_NOT_OK( - graph.BatchDeleteEdges(src_v_label, dst_v_label, edge_label, - oe_to_delete, ie_to_delete)); - } else { - flat_hash_map> edges_map; - for (size_t j = 0; j < edge_column->size(); j++) { - auto edge = edge_column->get_edge(j); - uint32_t index = graph.schema().generate_edge_label( - edge.label.src_label, edge.label.dst_label, - edge.label.edge_label); - if (edges_map.find(index) != edges_map.end()) { - edges_map[index].emplace_back(edge); - } else { - edges_map[index] = {edge}; - } - } +Stream BatchDeleteEdgeOpr::Eval( + IStorageInterface& graph_interface, const ParamsMap& params, + Stream&& input, OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Finish reading before mutation; downstream cancellation must not skip + // writes. + return reduce_stream( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + auto& graph = + dynamic_cast(graph_interface); + { + size_t binding_size = edge_bindings_.size(); + for (size_t i = 0; i < binding_size; i++) { + int32_t alias = edge_bindings_[i]; + auto& edge_triplets = edge_triplets_[i]; + auto edge_column = + std::dynamic_pointer_cast(chunk.get(alias)); + if (edge_triplets.size() == 1) { + label_t src_v_label = std::get<0>(edge_triplets[0]); + label_t dst_v_label = std::get<1>(edge_triplets[0]); + label_t edge_label = std::get<2>(edge_triplets[0]); + LabelTriplet request_triplet = + LabelTriplet(src_v_label, dst_v_label, edge_label); + size_t edge_size = edge_column->size(); + auto oe_view = graph.GetGenericOutgoingGraphView( + src_v_label, dst_v_label, edge_label); + auto ie_view = graph.GetGenericIncomingGraphView( + dst_v_label, src_v_label, edge_label); + auto edge_prop_types = graph.schema().get_edge_properties( + src_v_label, dst_v_label, edge_label); + std::vector> oe_to_delete, + ie_to_delete; + oe_to_delete.reserve(edge_size); + ie_to_delete.reserve(edge_size); + for (size_t j = 0; j < edge_size; j++) { + auto record = edge_column->get_edge(j); + if (record.label == request_triplet) { + auto offset_pair = record_to_csr_offset_pair( + oe_view, ie_view, record, edge_prop_types); + oe_to_delete.emplace_back(record.src, + offset_pair.first); + ie_to_delete.emplace_back(record.dst, + offset_pair.second); + } + } + RETURN_STATUS_ERROR_IF_NOT_OK(graph.BatchDeleteEdges( + src_v_label, dst_v_label, edge_label, oe_to_delete, + ie_to_delete)); + } else { + flat_hash_map> edges_map; + for (size_t j = 0; j < edge_column->size(); j++) { + auto edge = edge_column->get_edge(j); + uint32_t index = graph.schema().generate_edge_label( + edge.label.src_label, edge.label.dst_label, + edge.label.edge_label); + if (edges_map.find(index) != edges_map.end()) { + edges_map[index].emplace_back(edge); + } else { + edges_map[index] = {edge}; + } + } - for (auto& [index, edges] : edges_map) { - auto [src_v_label, dst_v_label, edge_label] = - graph.schema().parse_edge_label(index); - auto oe_view = graph.GetGenericOutgoingGraphView( - src_v_label, dst_v_label, edge_label); - auto ie_view = graph.GetGenericIncomingGraphView( - dst_v_label, src_v_label, edge_label); - std::vector> oe_to_delete, ie_to_delete; - oe_to_delete.reserve(edges.size()); - ie_to_delete.reserve(edges.size()); - auto edge_prop_types = graph.schema().get_edge_properties( - src_v_label, dst_v_label, edge_label); - for (auto& record : edges) { - auto offset_pair = record_to_csr_offset_pair( - oe_view, ie_view, record, edge_prop_types); - oe_to_delete.emplace_back(record.src, offset_pair.first); - ie_to_delete.emplace_back(record.dst, offset_pair.second); + for (auto& [index, edges] : edges_map) { + auto [src_v_label, dst_v_label, edge_label] = + graph.schema().parse_edge_label(index); + auto oe_view = graph.GetGenericOutgoingGraphView( + src_v_label, dst_v_label, edge_label); + auto ie_view = graph.GetGenericIncomingGraphView( + dst_v_label, src_v_label, edge_label); + std::vector> oe_to_delete, + ie_to_delete; + oe_to_delete.reserve(edges.size()); + ie_to_delete.reserve(edges.size()); + auto edge_prop_types = graph.schema().get_edge_properties( + src_v_label, dst_v_label, edge_label); + for (auto& record : edges) { + auto offset_pair = record_to_csr_offset_pair( + oe_view, ie_view, record, edge_prop_types); + oe_to_delete.emplace_back(record.src, + offset_pair.first); + ie_to_delete.emplace_back(record.dst, + offset_pair.second); + } + RETURN_STATUS_ERROR_IF_NOT_OK(graph.BatchDeleteEdges( + src_v_label, dst_v_label, edge_label, oe_to_delete, + ie_to_delete)); + } + } + sel_vec_t offsets; + chunk.reshuffle(offsets); // reshuffle with empty offsets to + // remove all data + } + return chunk; } - RETURN_STATUS_ERROR_IF_NOT_OK( - graph.BatchDeleteEdges(src_v_label, dst_v_label, edge_label, - oe_to_delete, ie_to_delete)); - } - } - sel_vec_t offsets; - chunk.reshuffle( - offsets); // reshuffle with empty offsets to remove all data - } - return chunk; + }); }); } diff --git a/src/execution/execute/ops/batch/batch_delete_vertex.cc b/src/execution/execute/ops/batch/batch_delete_vertex.cc index cd3c0fa52..e32fdb382 100644 --- a/src/execution/execute/ops/batch/batch_delete_vertex.cc +++ b/src/execution/execute/ops/batch/batch_delete_vertex.cc @@ -31,63 +31,78 @@ class BatchDeleteVertexOpr : public IOperator { return "BatchDeleteVertexOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override; private: std::vector> vertex_labels_; std::vector vertex_bindings_; }; -neug::result BatchDeleteVertexOpr::Eval( - IStorageInterface& graph_interface, const ParamsMap& params, Context&& ctx, - OprTimer* timer) { - auto& graph = dynamic_cast(graph_interface); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - size_t binding_size = vertex_bindings_.size(); - for (size_t i = 0; i < binding_size; i++) { - int32_t alias = vertex_bindings_[i]; - auto vertex_column = - std::dynamic_pointer_cast(chunk.get(alias)); - if (vertex_column->vertex_column_type() == - VertexColumnType::kSingle) { - auto sl_vertex_column = - std::dynamic_pointer_cast(vertex_column); - std::vector vids; - for (auto v : sl_vertex_column->vertices()) { - vids.emplace_back(v); - } - RETURN_STATUS_ERROR_IF_NOT_OK( - graph.BatchDeleteVertices(sl_vertex_column->label(), vids)); - } else if (vertex_column->vertex_column_type() == - VertexColumnType::kMultiple || - vertex_column->vertex_column_type() == - VertexColumnType::kMultiSegment) { - std::unordered_map> vids_map; - for (auto label : vertex_column->get_labels_set()) { - std::vector vids; - vids_map.insert({label, vids}); - } - size_t vertex_size = vertex_column->size(); - for (size_t j = 0; j < vertex_size; j++) { - auto vertex = vertex_column->get_vertex(j); - vids_map.at(vertex.label_).emplace_back(vertex.vid_); - } - for (auto& vids_pair : vids_map) { - RETURN_STATUS_ERROR_IF_NOT_OK( - graph.BatchDeleteVertices(vids_pair.first, vids_pair.second)); - } - } else { - THROW_RUNTIME_ERROR( - "Unsupported vertex column type for batch delete vertex " - "operation."); - } - sel_vec_t offsets; - chunk.reshuffle( - offsets); // reshuffle with empty offsets to remove all data - } - return chunk; +Stream BatchDeleteVertexOpr::Eval( + IStorageInterface& graph_interface, const ParamsMap& params, + Stream&& input, OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Finish reading before mutation; downstream cancellation must not skip + // writes. + return reduce_stream( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + auto& graph = + dynamic_cast(graph_interface); + { + size_t binding_size = vertex_bindings_.size(); + for (size_t i = 0; i < binding_size; i++) { + int32_t alias = vertex_bindings_[i]; + auto vertex_column = std::dynamic_pointer_cast( + chunk.get(alias)); + if (vertex_column->vertex_column_type() == + VertexColumnType::kSingle) { + auto sl_vertex_column = + std::dynamic_pointer_cast( + vertex_column); + std::vector vids; + for (auto v : sl_vertex_column->vertices()) { + vids.emplace_back(v); + } + RETURN_STATUS_ERROR_IF_NOT_OK(graph.BatchDeleteVertices( + sl_vertex_column->label(), vids)); + } else if (vertex_column->vertex_column_type() == + VertexColumnType::kMultiple || + vertex_column->vertex_column_type() == + VertexColumnType::kMultiSegment) { + std::unordered_map> vids_map; + for (auto label : vertex_column->get_labels_set()) { + std::vector vids; + vids_map.insert({label, vids}); + } + size_t vertex_size = vertex_column->size(); + for (size_t j = 0; j < vertex_size; j++) { + auto vertex = vertex_column->get_vertex(j); + vids_map.at(vertex.label_).emplace_back(vertex.vid_); + } + for (auto& vids_pair : vids_map) { + RETURN_STATUS_ERROR_IF_NOT_OK(graph.BatchDeleteVertices( + vids_pair.first, vids_pair.second)); + } + } else { + THROW_RUNTIME_ERROR( + "Unsupported vertex column type for batch delete " + "vertex " + "operation."); + } + sel_vec_t offsets; + chunk.reshuffle(offsets); // reshuffle with empty offsets to + // remove all data + } + return chunk; + } + }); }); } diff --git a/src/execution/execute/ops/batch/batch_insert_edge.cc b/src/execution/execute/ops/batch/batch_insert_edge.cc index f4fa3c25b..95f90fbb8 100644 --- a/src/execution/execute/ops/batch/batch_insert_edge.cc +++ b/src/execution/execute/ops/batch/batch_insert_edge.cc @@ -106,8 +106,9 @@ class BatchInsertEdgeOpr : public IOperator { return "BatchInsertEdgeOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override; private: physical::EdgeType edge_type_; @@ -115,39 +116,67 @@ class BatchInsertEdgeOpr : public IOperator { src_vertex_bindings_, dst_vertex_bindings_; }; -neug::result BatchInsertEdgeOpr::Eval( - IStorageInterface& graph_interface, const ParamsMap& params, Context&& ctx, - OprTimer* timer) { - (void) params; - (void) timer; - auto& graph = dynamic_cast(graph_interface); - label_t edge_label_id = 0; - label_t src_label_id = 0; - label_t dst_label_id = 0; - if (!resolve_edge_triplet(graph.schema(), edge_type_, edge_label_id, - src_label_id, dst_label_id)) { - RETURN_STATUS_ERROR(StatusCode::ERR_INVALID_ARGUMENT, - "Failed to resolve edge type or vertex endpoints for " - "BatchInsertEdge"); - } - - std::vector> total_mappings; - total_mappings.reserve(src_vertex_bindings_.size() + - dst_vertex_bindings_.size() + prop_mappings_.size()); - for (const auto& mapping : src_vertex_bindings_) { - total_mappings.emplace_back(mapping); - } - for (const auto& mapping : dst_vertex_bindings_) { - total_mappings.emplace_back(mapping); - } - for (const auto& mapping : prop_mappings_) { - total_mappings.emplace_back(mapping); - } - auto 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)); +Stream BatchInsertEdgeOpr::Eval( + IStorageInterface& graph_interface, const ParamsMap& params, + Stream&& input, OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Pull once so upstream schema creation completes before resolving + // the target label; retain the batch for the storage supplier. + auto first = input.Next(); + if (!first) { + return error_stream(first.error()); + } + input = prepend_chunk(std::move(*first), std::move(input)); + + (void) params; + (void) timer; + auto& graph = dynamic_cast(graph_interface); + label_t edge_label_id = 0; + label_t src_label_id = 0; + label_t dst_label_id = 0; + if (!resolve_edge_triplet(graph.schema(), edge_type_, edge_label_id, + src_label_id, dst_label_id)) { + return error_stream( + Status(StatusCode::ERR_INVALID_ARGUMENT, + "Failed to resolve edge type or vertex endpoints for " + "BatchInsertEdge")); + } + + std::vector> total_mappings; + total_mappings.reserve(src_vertex_bindings_.size() + + dst_vertex_bindings_.size() + + prop_mappings_.size()); + for (const auto& mapping : src_vertex_bindings_) { + total_mappings.emplace_back(mapping); + } + for (const auto& mapping : dst_vertex_bindings_) { + total_mappings.emplace_back(mapping); + } + for (const auto& mapping : prop_mappings_) { + total_mappings.emplace_back(mapping); + } + auto supplier = std::make_shared(std::move(input), + total_mappings); + + auto insert_status = graph.BatchAddEdges(src_label_id, dst_label_id, + edge_label_id, supplier); + { + auto status = supplier->status(); + if (!status.ok()) { + return error_stream(status); + } + }; + { + auto status = insert_status; + if (!status.ok()) { + return error_stream(status); + } + }; + return batch_insert_result(supplier->rows_read()); + }); } neug::result BatchInsertEdgeOprBuilder::Build( diff --git a/src/execution/execute/ops/batch/batch_insert_vertex.cc b/src/execution/execute/ops/batch/batch_insert_vertex.cc index 6c48cff78..ffde20252 100644 --- a/src/execution/execute/ops/batch/batch_insert_vertex.cc +++ b/src/execution/execute/ops/batch/batch_insert_vertex.cc @@ -43,45 +43,68 @@ class BatchInsertVertexOpr : public IOperator { return "BatchInsertVertexOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override; private: common::NameOrId vertex_type_; std::vector> prop_mappings_; }; -neug::result BatchInsertVertexOpr::Eval( - IStorageInterface& graph_interface, const ParamsMap& params, Context&& ctx, - OprTimer* timer) { - (void) params; - (void) timer; - auto& graph = dynamic_cast(graph_interface); - label_t vertex_label_id = 0; - switch (vertex_type_.item_case()) { - case common::NameOrId::kId: - vertex_label_id = vertex_type_.id(); - break; - case common::NameOrId::kName: { - const auto& name = vertex_type_.name(); - if (!graph.schema().is_vertex_label_valid(name)) { - LOG(ERROR) << "Unknown vertex type: " << vertex_type_.DebugString(); - RETURN_STATUS_ERROR(StatusCode::ERR_INVALID_ARGUMENT, - "Unknown vertex type: " + name); - } - vertex_label_id = graph.schema().get_vertex_label_id(name); - break; - } - default: - THROW_INVALID_ARGUMENT_EXCEPTION( - "BatchInsertVertexOpr: invalid vertex_type: " + - vertex_type_.DebugString()); - } - auto supplier = create_data_chunk_supplier(ctx, prop_mappings_); - GS_AUTO(inserted_vids, - graph.BatchAddVertices(vertex_label_id, std::move(supplier))); - (void) inserted_vids; - return neug::result(std::move(ctx)); +Stream BatchInsertVertexOpr::Eval( + IStorageInterface& graph_interface, const ParamsMap& params, + Stream&& input, OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Pull once so upstream schema creation completes before resolving + // the target label; retain the batch for the storage supplier. + auto first = input.Next(); + if (!first) { + return error_stream(first.error()); + } + input = prepend_chunk(std::move(*first), std::move(input)); + + (void) params; + (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 error_stream( + Status(StatusCode::ERR_INVALID_ARGUMENT, + "Unknown vertex type: " + name)); + } + vertex_label_id = graph.schema().get_vertex_label_id(name); + break; + } + default: + THROW_INVALID_ARGUMENT_EXCEPTION( + "BatchInsertVertexOpr: invalid vertex_type: " + + vertex_type_.DebugString()); + } + auto supplier = std::make_shared(std::move(input), + prop_mappings_); + auto inserted_vids = graph.BatchAddVertices(vertex_label_id, supplier); + { + auto status = supplier->status(); + if (!status.ok()) { + return error_stream(status); + } + }; + if (!inserted_vids) { + return error_stream(inserted_vids.error()); + } + return batch_insert_result(supplier->rows_read()); + }); } neug::result BatchInsertVertexOprBuilder::Build( diff --git a/src/execution/execute/ops/batch/batch_update_edge.cc b/src/execution/execute/ops/batch/batch_update_edge.cc index d30eb5aec..348e4fabb 100644 --- a/src/execution/execute/ops/batch/batch_update_edge.cc +++ b/src/execution/execute/ops/batch/batch_update_edge.cc @@ -44,111 +44,137 @@ class UpdateEdgeOpr : public IOperator { std::string get_operator_name() const override { return "UpdateEdgeOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override; private: edge_data_vec_t edge_data_; }; -neug::result UpdateEdgeOpr::Eval(IStorageInterface& graph_interface, - const ParamsMap& params, - Context&& ctx, OprTimer* timer) { - auto& graph = dynamic_cast(graph_interface); - VLOG(10) << "Executing UpdateEdgeOpr with " << edge_data_.size() - << " entries."; - - struct PendingUpdate { - EdgeRecord record; - std::pair offsets; - int32_t property_id; - Value value; - }; - - std::set edge_tags; - for (const auto& entry : edge_data_) { - edge_tags.insert(std::get<0>(entry)); - } - - std::set affected_labels; - for (const auto tag_id : edge_tags) { - for (auto& chunk : ctx.chunks()) { - auto column = chunk.get(tag_id); - auto edge_column = std::dynamic_pointer_cast(column); - if (!edge_column) { - continue; - } - for (const auto& label : edge_column->get_labels()) { - const auto edge_schema = graph.schema().get_edge_schema( - label.src_label, label.dst_label, label.edge_label); - if (edge_schema->is_bundled()) { - affected_labels.insert(label); +Stream UpdateEdgeOpr::Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Edge property writes can invalidate pointers in later batches. Retain + // and refresh every affected column before exposing the result + // downstream. + auto metadata = input.metadata(); + auto chunks_result = collect_batches(std::move(input)); + if (!chunks_result) { + return error_stream(chunks_result.error()); } - } - } - } - - auto snapshots = CaptureEdgeColumnsForRefresh(graph, ctx, affected_labels); - for (const auto& [tag_id, property_name, expression] : edge_data_) { - auto bound_expression = expression->bind(&graph, params); - const auto& record_expression = bound_expression->Cast(); - std::vector updates; - bool refresh_columns = false; - for (auto& chunk : ctx.chunks()) { - auto column = chunk.get(tag_id); - if (!column) { - THROW_RUNTIME_ERROR("Column " + std::to_string(tag_id) + - " not found in context."); - } - auto edge_column = std::dynamic_pointer_cast(column); - if (!edge_column) { - THROW_RUNTIME_ERROR("Column " + std::to_string(tag_id) + - " is not an edge column."); - } - - for (size_t row = 0; row < edge_column->size(); ++row) { - if (!edge_column->has_value(row)) { - continue; + auto chunks = std::move(*chunks_result); + + auto& graph = dynamic_cast(graph_interface); + VLOG(10) << "Executing UpdateEdgeOpr with " << edge_data_.size() + << " entries."; + + struct PendingUpdate { + EdgeRecord record; + std::pair offsets; + int32_t property_id; + Value value; + }; + + std::set edge_tags; + for (const auto& entry : edge_data_) { + edge_tags.insert(std::get<0>(entry)); } - const auto record = edge_column->get_edge(row); - const auto edge_schema = graph.schema().get_edge_schema( - record.label.src_label, record.label.dst_label, - record.label.edge_label); - const auto property_id = edge_schema->get_property_index(property_name); - if (property_id < 0) { - THROW_RUNTIME_ERROR( - "Property " + property_name + " does not exist for edge label: " + - std::to_string(static_cast(record.label.edge_label))); + std::set affected_labels; + for (const auto tag_id : edge_tags) { + for (auto& chunk : chunks) { + auto column = chunk.get(tag_id); + auto edge_column = std::dynamic_pointer_cast(column); + if (!edge_column) { + continue; + } + for (const auto& label : edge_column->get_labels()) { + const auto edge_schema = graph.schema().get_edge_schema( + label.src_label, label.dst_label, label.edge_label); + if (edge_schema->is_bundled()) { + affected_labels.insert(label); + } + } + } } - auto value = record_expression.eval_record(chunk.chunk(), row); - if (value.IsNull()) { - THROW_NOT_SUPPORTED_EXCEPTION("Setting NULL for property " + - property_name); + auto snapshots = + CaptureEdgeColumnsForRefresh(graph, chunks, affected_labels); + for (const auto& [tag_id, property_name, expression] : edge_data_) { + auto bound_expression = expression->bind(&graph, params); + const auto& record_expression = + bound_expression->Cast(); + std::vector updates; + bool refresh_columns = false; + for (auto& chunk : chunks) { + auto column = chunk.get(tag_id); + if (!column) { + THROW_RUNTIME_ERROR("Column " + std::to_string(tag_id) + + " not found in context."); + } + auto edge_column = std::dynamic_pointer_cast(column); + if (!edge_column) { + THROW_RUNTIME_ERROR("Column " + std::to_string(tag_id) + + " is not an edge column."); + } + + for (size_t row = 0; row < edge_column->size(); ++row) { + if (!edge_column->has_value(row)) { + continue; + } + + const auto record = edge_column->get_edge(row); + const auto edge_schema = graph.schema().get_edge_schema( + record.label.src_label, record.label.dst_label, + record.label.edge_label); + const auto property_id = + edge_schema->get_property_index(property_name); + if (property_id < 0) { + THROW_RUNTIME_ERROR( + "Property " + property_name + + " does not exist for edge label: " + + std::to_string(static_cast(record.label.edge_label))); + } + + auto value = record_expression.eval_record(chunk.chunk(), row); + if (value.IsNull()) { + THROW_NOT_SUPPORTED_EXCEPTION("Setting NULL for property " + + property_name); + } + if (edge_schema->properties[property_id] != value.type()) { + THROW_RUNTIME_ERROR("Property type mismatch for property " + + property_name); + } + refresh_columns = refresh_columns || edge_schema->is_bundled(); + updates.push_back(PendingUpdate{record, + ResolveEdgeOffsets(graph, record), + property_id, std::move(value)}); + } + } + for (const auto& update : updates) { + { + auto status = graph.UpdateEdgeProperty( + update.record.label.src_label, update.record.src, + update.record.label.dst_label, update.record.dst, + update.record.label.edge_label, update.offsets.first, + update.offsets.second, update.property_id, update.value); + if (!status.ok()) { + return error_stream(status); + } + }; + } + if (refresh_columns) { + RefreshEdgeColumns(graph, snapshots); + } } - if (edge_schema->properties[property_id] != value.type()) { - THROW_RUNTIME_ERROR("Property type mismatch for property " + - property_name); - } - refresh_columns = refresh_columns || edge_schema->is_bundled(); - updates.push_back(PendingUpdate{record, - ResolveEdgeOffsets(graph, record), - property_id, std::move(value)}); - } - } - for (const auto& update : updates) { - RETURN_STATUS_ERROR_IF_NOT_OK(graph.UpdateEdgeProperty( - update.record.label.src_label, update.record.src, - update.record.label.dst_label, update.record.dst, - update.record.label.edge_label, update.offsets.first, - update.offsets.second, update.property_id, update.value)); - } - if (refresh_columns) { - RefreshEdgeColumns(graph, snapshots); - } - } - return std::move(ctx); + return stream_from_batches(std::move(chunks), std::move(metadata)); + }); } neug::result UpdateEdgeOprBuilder::Build( diff --git a/src/execution/execute/ops/batch/batch_update_utils.cc b/src/execution/execute/ops/batch/batch_update_utils.cc index bac2c05bc..6a6674e67 100644 --- a/src/execution/execute/ops/batch/batch_update_utils.cc +++ b/src/execution/execute/ops/batch/batch_update_utils.cc @@ -273,51 +273,57 @@ std::string path_to_json_string(Path& path, const StorageReadInterface& graph) { return buffer.GetString(); } -/// A supplier that yields pre-projected DataChunks one by one. -class MultiChunkSupplier : public IDataChunkSupplier { - public: - explicit MultiChunkSupplier(std::vector> chunks) - : chunks_(std::move(chunks)), index_(0) {} +StreamChunkSupplier::StreamChunkSupplier( + Stream stream, + std::vector> mappings) + : stream_(std::move(stream)), mappings_(std::move(mappings)) {} - std::shared_ptr GetNextChunk() override { - if (index_ >= chunks_.size()) - return nullptr; - return chunks_[index_++]; +std::shared_ptr StreamChunkSupplier::GetNextChunk() { + auto next = stream_.Next(); + if (!next) { + status_ = next.error(); + return nullptr; } - - int64_t RowNum() const override { - int64_t total = 0; - for (const auto& chunk : chunks_) { - total += static_cast(chunk->row_num()); + if (!*next) { + return nullptr; + } + rows_read_ += (**next).row_num(); + auto output = std::make_shared(); + for (size_t i = 0; i < mappings_.size(); ++i) { + auto column = (**next).get(mappings_[i].first); + if (!column) { + THROW_INTERNAL_EXCEPTION("Column not found for tag id: " + + std::to_string(mappings_[i].first)); } - return total; + output->set(static_cast(i), std::move(column)); } + return output; +} - private: - std::vector> chunks_; - size_t index_; -}; - -std::shared_ptr create_data_chunk_supplier( - const Context& ctx, - const std::vector>& prop_mappings) { - std::vector> projected_chunks; - projected_chunks.reserve(ctx.chunk_num()); - for (size_t i = 0; i < ctx.chunk_num(); ++i) { - const auto& chunk = ctx.chunk(i).chunk(); - auto out_chunk = std::make_shared(); - for (size_t j = 0; j < prop_mappings.size(); ++j) { - auto tag_id = prop_mappings[j].first; - auto column = chunk.get(tag_id); - if (column == nullptr) { - THROW_INTERNAL_EXCEPTION("Column not found for tag id: " + - std::to_string(tag_id)); - } - out_chunk->set(static_cast(j), column); +Stream batch_insert_result(size_t rows) { + // COPY's sink has no output tags, but QueryResponse still reports the + // consumed row count. A head-only constant column preserves that contract + // in O(1) space rather than keeping every input property column alive. + class CardinalityColumn final : public IContextColumn { + public: + explicit CardinalityColumn(size_t rows) : rows_(rows) {} + size_t size() const override { return rows_; } + std::string column_info() const override { return "COPY cardinality"; } + ContextColumnType column_type() const override { + return ContextColumnType::kValue; } - projected_chunks.push_back(std::move(out_chunk)); - } - return std::make_shared(std::move(projected_chunks)); + const DataType& elem_type() const override { return DataType::BOOLEAN; } + Value get_elem(size_t) const override { return Value::BOOLEAN(true); } + bool is_optional() const override { return false; } + + private: + size_t rows_; + }; + return generate_chunk([rows]() -> result { + ContextChunk output; + output.set(-1, std::make_shared(rows)); + return output; + }); } std::vector match_files_with_pattern( diff --git a/src/execution/execute/ops/batch/batch_update_vertex.cc b/src/execution/execute/ops/batch/batch_update_vertex.cc index 759c60399..a4a6468f3 100644 --- a/src/execution/execute/ops/batch/batch_update_vertex.cc +++ b/src/execution/execute/ops/batch/batch_update_vertex.cc @@ -40,8 +40,9 @@ class UpdateVertexOpr : public IOperator { const ParamsMap& params, ContextChunk&& chunk, OprTimer* timer); - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override; private: // No alias is produced in this operator. @@ -114,13 +115,25 @@ neug::result UpdateVertexOpr::eval_impl( return chunk; } -neug::result UpdateVertexOpr::Eval(IStorageInterface& graph_interface, - const ParamsMap& params, - Context&& ctx, OprTimer* timer) { - auto& graph = dynamic_cast(graph_interface); - return ctx.apply_chunks([&](ContextChunk&& chunk) { - return eval_impl(graph, params, std::move(chunk), timer); - }); +Stream UpdateVertexOpr::Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Finish reading before mutation; downstream cancellation must not skip + // writes. + return reduce_stream( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + auto& graph = + dynamic_cast(graph_interface); + { return eval_impl(graph, params, std::move(chunk), timer); } + }); + }); } neug::result UpdateVertexOprBuilder::Build( diff --git a/src/execution/execute/ops/batch/data_export.cc b/src/execution/execute/ops/batch/data_export.cc index 9a8e6c878..f5df1a69b 100644 --- a/src/execution/execute/ops/batch/data_export.cc +++ b/src/execution/execute/ops/batch/data_export.cc @@ -36,10 +36,9 @@ class DataExportOpr : public IOperator { std::string get_operator_name() const override { return "DataExportOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override; private: reader::FileSchema schema_; @@ -47,19 +46,36 @@ class DataExportOpr : public IOperator { function::ExportFunction* exportFunction_; }; -neug::result DataExportOpr::Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, neug::execution::OprTimer* timer) { - const auto& graph = - dynamic_cast(graph_interface); - if (!exportFunction_) { - THROW_IO_EXCEPTION("DataExportOpr: export function is nullptr"); - } - if (!exportFunction_->execFunc) { - THROW_IO_EXCEPTION( - "DataExportOpr: write function in export function is nullptr"); - } - return exportFunction_->execFunc(ctx, schema_, entry_schema_, graph); +Stream DataExportOpr::Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Legacy extension ABI: Context conversion is confined to this + // boundary. + + auto ctx_result = materialize(std::move(input)); + if (!ctx_result) { + return error_stream(ctx_result.error()); + } + auto ctx = std::move(*ctx_result); + + const auto& graph = + dynamic_cast(graph_interface); + if (!exportFunction_) { + THROW_IO_EXCEPTION("DataExportOpr: export function is nullptr"); + } + if (!exportFunction_->execFunc) { + THROW_IO_EXCEPTION( + "DataExportOpr: write function in export function is nullptr"); + } + auto output = + exportFunction_->execFunc(ctx, schema_, entry_schema_, graph); + return stream_from_context(std::move(output)); + }); } neug::result DataExportOprBuilder::Build( diff --git a/src/execution/execute/ops/batch/data_source.cc b/src/execution/execute/ops/batch/data_source.cc index 432776cc1..a853a064c 100644 --- a/src/execution/execute/ops/batch/data_source.cc +++ b/src/execution/execute/ops/batch/data_source.cc @@ -24,6 +24,8 @@ #include "neug/compiler/main/metadata_registry.h" #include "neug/execution/common/context.h" #include "neug/execution/execute/ops/batch/data_source.h" +#include "neug/execution/utils/params.h" +#include "neug/storages/loader/loader_utils.h" #include "neug/utils/io/read/common/schema.h" #include "neug/utils/io/reader.h" #include "neug/utils/result.h" @@ -49,15 +51,53 @@ class DataSourceOpr : public IOperator { std::string get_operator_name() const override { return "DataSourceOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - NEUG_ASSERT(readFunction != nullptr); - // Parameters belong to this evaluation, not to the cached physical plan. - auto state = std::make_shared(*sharedState); - state->parameters = params; - return readFunction->execFunc(state); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + while (true) { + auto before = input.Next(); + if (!before) { + return error_stream(before.error()); + } + if (!*before) { + break; + } + } + + NEUG_ASSERT(readFunction != nullptr); + // Reader initialization may expand globs and normalize options. Never + // mutate the state captured by the cached operator. + auto state = std::make_shared(*sharedState); + state->parameters = params; + auto function = readFunction; + struct Cursor { + bool initialized = false; + std::shared_ptr supplier; + }; + auto cursor = std::make_shared(); + auto raw = Stream([state, function, cursor]() mutable + -> Stream::NextResult { + if (!cursor->initialized) { + cursor->initialized = true; + cursor->supplier = function->supplierFunc(state); + if (!cursor->supplier) { + return tl::unexpected( + Status::InternalError("Reader returned a null supplier")); + } + } + auto chunk = cursor->supplier->GetNextChunk(); + if (!chunk) { + return std::optional{}; + } + return std::optional(std::in_place, + std::move(*chunk)); + }); + return std::move(raw); + }); } }; diff --git a/src/execution/execute/ops/ddl/add_edge_property.cc b/src/execution/execute/ops/ddl/add_edge_property.cc index 6224cacad..c71523a27 100644 --- a/src/execution/execute/ops/ddl/add_edge_property.cc +++ b/src/execution/execute/ops/ddl/add_edge_property.cc @@ -37,36 +37,50 @@ class AddEdgePropertySchemaOpr : public IOperator { std::string get_operator_name() const override { return "AddEdgePropertySchemaOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - StorageUpdateInterface& storage = - dynamic_cast(graph); - label_t src, dst, edge; - auto resolve = ResolveEdgeTriplet(storage.schema(), src_type_, dst_type_, - edge_type_, src, dst, edge); - if (!resolve.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(resolve)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to add edge property to type: " << edge_type_ - << ", reason: " << resolve.ToString(); - RETURN_ERROR(resolve); - } - AddEdgePropertiesParamBuilder builder; - for (const auto& [prop_name, prop_value] : properties_) { - builder.AddProperty(prop_name, prop_value); - } - auto config = builder.Build(); - auto res = storage.AddEdgeProperties(src, dst, edge, config); - if (!res.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(res)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to add edge property to type: " << edge_type_ - << ", reason: " << res.ToString(); - RETURN_ERROR(res); - } - return neug::result(std::move(ctx)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); + + StorageUpdateInterface& storage = + dynamic_cast(graph); + label_t src, dst, edge; + auto resolve = + ResolveEdgeTriplet(storage.schema(), src_type_, dst_type_, + edge_type_, src, dst, edge); + if (!resolve.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(resolve)) { + return std::move(input); + } + LOG(ERROR) << "Fail to add edge property to type: " << edge_type_ + << ", reason: " << resolve.ToString(); + return error_stream(resolve); + } + AddEdgePropertiesParamBuilder builder; + for (const auto& [prop_name, prop_value] : properties_) { + builder.AddProperty(prop_name, prop_value); + } + auto config = builder.Build(); + auto res = storage.AddEdgeProperties(src, dst, edge, config); + if (!res.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(res)) { + return std::move(input); + } + LOG(ERROR) << "Fail to add edge property to type: " << edge_type_ + << ", reason: " << res.ToString(); + return error_stream(res); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/ddl/add_vertex_property.cc b/src/execution/execute/ops/ddl/add_vertex_property.cc index f2a1ee2d2..483d460a0 100644 --- a/src/execution/execute/ops/ddl/add_vertex_property.cc +++ b/src/execution/execute/ops/ddl/add_vertex_property.cc @@ -35,35 +35,49 @@ class AddVertexPropertySchemaOpr : public IOperator { std::string get_operator_name() const override { return "AddVertexPropertySchemaOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - StorageUpdateInterface& storage = - dynamic_cast(graph); - label_t label; - auto resolve = ResolveVertexLabel(storage.schema(), vertex_type_, label); - if (!resolve.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(resolve)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to add vertex property to type: " << vertex_type_ - << ", reason: " << resolve.ToString(); - RETURN_ERROR(resolve); - } - AddVertexPropertiesParamBuilder builder; - for (const auto& [prop_name, prop_value] : properties_) { - builder.AddProperty(prop_name, prop_value); - } - auto config = builder.Build(); - auto res = storage.AddVertexProperties(label, config); - if (!res.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(res)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to add vertex property to type: " << vertex_type_ - << ", reason: " << res.ToString(); - RETURN_ERROR(res); - } - return neug::result(std::move(ctx)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); + + StorageUpdateInterface& storage = + dynamic_cast(graph); + label_t label; + auto resolve = + ResolveVertexLabel(storage.schema(), vertex_type_, label); + if (!resolve.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(resolve)) { + return std::move(input); + } + LOG(ERROR) << "Fail to add vertex property to type: " + << vertex_type_ << ", reason: " << resolve.ToString(); + return error_stream(resolve); + } + AddVertexPropertiesParamBuilder builder; + for (const auto& [prop_name, prop_value] : properties_) { + builder.AddProperty(prop_name, prop_value); + } + auto config = builder.Build(); + auto res = storage.AddVertexProperties(label, config); + if (!res.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(res)) { + return std::move(input); + } + LOG(ERROR) << "Fail to add vertex property to type: " + << vertex_type_ << ", reason: " << res.ToString(); + return error_stream(res); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/ddl/create_edge_type.cc b/src/execution/execute/ops/ddl/create_edge_type.cc index 45ac40397..d8d32e3ed 100644 --- a/src/execution/execute/ops/ddl/create_edge_type.cc +++ b/src/execution/execute/ops/ddl/create_edge_type.cc @@ -35,66 +35,82 @@ class CreateEdgeTypeOpr : public IOperator { std::string get_operator_name() const override { return "CreateEdgeTypeOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - StorageUpdateInterface& storage = - dynamic_cast(graph); - int32_t defs_size = create_edge_types_.size(); - // Track indices of edge types actually created by this operator, - // so rollback only reverts what we created (not pre-existing types). - std::vector created_indices; - Status status; - bool failed = false; - for (int32_t i = 0; i < defs_size; ++i) { - const auto& create_edge_def = create_edge_types_[i]; - CreateEdgeTypeParamBuilder config_builder; - for (const auto& [prop_name, prop_value] : std::get<3>(create_edge_def)) { - config_builder.AddProperty(prop_name, prop_value); - } - config_builder.SrcLabel(std::get<0>(create_edge_def)) - .DstLabel(std::get<1>(create_edge_def)) - .EdgeLabel(std::get<2>(create_edge_def)) - .OEEdgeStrategy(std::get<4>(create_edge_def)) - .IEEdgeStrategy(std::get<5>(create_edge_def)) - .SortKeyForNbr(std::get<6>(create_edge_def)) - .Temporary(std::get<7>(create_edge_def)); - status = storage.CreateEdgeType(config_builder.Build()); - if (!status.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(status)) { - continue; - } - LOG(ERROR) << "Fail to insert edge triplet: " - << std::get<0>(create_edge_def) << ", " - << std::get<1>(create_edge_def) << ", " - << std::get<2>(create_edge_def) - << ", reason: " << status.ToString(); - failed = true; - break; - } - created_indices.push_back(i); - } - if (failed) { - // Rollback only the edge types we actually created. - for (auto it = created_indices.rbegin(); it != created_indices.rend(); - ++it) { - const auto& create_edge_def = create_edge_types_[*it]; - label_t src, dst, edge; - auto resolve = - ResolveEdgeTriplet(storage.schema(), std::get<0>(create_edge_def), - std::get<1>(create_edge_def), - std::get<2>(create_edge_def), src, dst, edge); - // Resolve may fail if this entry left nothing to revert; skip then. - if (!resolve.ok()) { - continue; - } - if (!storage.DeleteEdgeType(src, dst, edge).ok()) { - LOG(ERROR) << "Fail to revert created edge type in CreateEdgeSchema " - "request"; - } - } - RETURN_ERROR(status); - } - return neug::result(std::move(ctx)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); + + StorageUpdateInterface& storage = + dynamic_cast(graph); + int32_t defs_size = create_edge_types_.size(); + // Track indices of edge types actually created by this operator, + // so rollback only reverts what we created (not pre-existing types). + std::vector created_indices; + Status status; + bool failed = false; + for (int32_t i = 0; i < defs_size; ++i) { + const auto& create_edge_def = create_edge_types_[i]; + CreateEdgeTypeParamBuilder config_builder; + for (const auto& [prop_name, prop_value] : + std::get<3>(create_edge_def)) { + config_builder.AddProperty(prop_name, prop_value); + } + config_builder.SrcLabel(std::get<0>(create_edge_def)) + .DstLabel(std::get<1>(create_edge_def)) + .EdgeLabel(std::get<2>(create_edge_def)) + .OEEdgeStrategy(std::get<4>(create_edge_def)) + .IEEdgeStrategy(std::get<5>(create_edge_def)) + .SortKeyForNbr(std::get<6>(create_edge_def)) + .Temporary(std::get<7>(create_edge_def)); + status = storage.CreateEdgeType(config_builder.Build()); + if (!status.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(status)) { + continue; + } + LOG(ERROR) << "Fail to insert edge triplet: " + << std::get<0>(create_edge_def) << ", " + << std::get<1>(create_edge_def) << ", " + << std::get<2>(create_edge_def) + << ", reason: " << status.ToString(); + failed = true; + break; + } + created_indices.push_back(i); + } + if (failed) { + // Rollback only the edge types we actually created. + for (auto it = created_indices.rbegin(); + it != created_indices.rend(); ++it) { + const auto& create_edge_def = create_edge_types_[*it]; + label_t src, dst, edge; + auto resolve = ResolveEdgeTriplet( + storage.schema(), std::get<0>(create_edge_def), + std::get<1>(create_edge_def), std::get<2>(create_edge_def), + src, dst, edge); + // Resolve may fail if this entry left nothing to revert; skip + // then. + if (!resolve.ok()) { + continue; + } + if (!storage.DeleteEdgeType(src, dst, edge).ok()) { + LOG(ERROR) + << "Fail to revert created edge type in CreateEdgeSchema " + "request"; + } + } + return error_stream(status); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/ddl/create_index.cc b/src/execution/execute/ops/ddl/create_index.cc index 9bb39b1d0..c24f02405 100644 --- a/src/execution/execute/ops/ddl/create_index.cc +++ b/src/execution/execute/ops/ddl/create_index.cc @@ -66,27 +66,42 @@ class CreateIndexOpr : public IOperator { std::string get_operator_name() const override { return "CreateIndexOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - auto* index_interface = dynamic_cast(&graph); - if (!index_interface) { - RETURN_STATUS_ERROR( - StatusCode::ERR_NOT_SUPPORTED, - "Current storage interface does not support index DDL"); - } - - auto index_meta = CreateIndexMeta(graph.schema(), create_index_); - auto index = index_interface->CreateIndex(std::move(index_meta)); - if (!index) { - // The storage layer reports ERR_ILLEGAL_OPERATION when an index with - // the same name already exists; honor IF NOT EXISTS in that case. - if (ignore_conflict_ && - index.error().error_code() == StatusCode::ERR_ILLEGAL_OPERATION) { - return std::move(ctx); - } - RETURN_ERROR(index.error()); - } - return std::move(ctx); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); + + auto* index_interface = + dynamic_cast(&graph); + if (!index_interface) { + return error_stream( + Status(StatusCode::ERR_NOT_SUPPORTED, + "Current storage interface does not support index DDL")); + } + + auto index_meta = CreateIndexMeta(graph.schema(), create_index_); + auto index = index_interface->CreateIndex(std::move(index_meta)); + if (!index) { + // The storage layer reports ERR_ILLEGAL_OPERATION when an index + // with the same name already exists; honor IF NOT EXISTS in that + // case. + if (ignore_conflict_ && index.error().error_code() == + StatusCode::ERR_ILLEGAL_OPERATION) { + return std::move(input); + } + return error_stream(index.error()); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/ddl/create_vertex_type.cc b/src/execution/execute/ops/ddl/create_vertex_type.cc index bacb1ef6a..864860cb3 100644 --- a/src/execution/execute/ops/ddl/create_vertex_type.cc +++ b/src/execution/execute/ops/ddl/create_vertex_type.cc @@ -38,27 +38,40 @@ class CreateVertexTypeOpr : public IOperator { return "CreateVertexTypeOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - StorageUpdateInterface& storage = - dynamic_cast(graph); - CreateVertexTypeParamBuilder builder; - builder.VertexLabel(type_name_) - .PrimaryKeyNames(pks_) - .Temporary(is_temporary_); - for (const auto& [prop_name, prop_value] : properties_) { - builder.AddProperty(prop_name, prop_value); - } - auto res = storage.CreateVertexType(builder.Build()); - if (!res.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(res)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to create vertex type: " << type_name_ - << ", reason: " << res.ToString(); - RETURN_ERROR(res); - } - return neug::result(std::move(ctx)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); + + StorageUpdateInterface& storage = + dynamic_cast(graph); + CreateVertexTypeParamBuilder builder; + builder.VertexLabel(type_name_) + .PrimaryKeyNames(pks_) + .Temporary(is_temporary_); + for (const auto& [prop_name, prop_value] : properties_) { + builder.AddProperty(prop_name, prop_value); + } + auto res = storage.CreateVertexType(builder.Build()); + if (!res.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(res)) { + return std::move(input); + } + LOG(ERROR) << "Fail to create vertex type: " << type_name_ + << ", reason: " << res.ToString(); + return error_stream(res); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/ddl/drop_edge_property.cc b/src/execution/execute/ops/ddl/drop_edge_property.cc index 0185016dc..f93f8e6b7 100644 --- a/src/execution/execute/ops/ddl/drop_edge_property.cc +++ b/src/execution/execute/ops/ddl/drop_edge_property.cc @@ -37,33 +37,47 @@ class DropEdgePropertySchemaOpr : public IOperator { std::string get_operator_name() const override { return "DropEdgePropertySchemaOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - StorageUpdateInterface& storage = - dynamic_cast(graph); - label_t src, dst, edge; - auto resolve = ResolveEdgeTriplet(storage.schema(), src_type_, dst_type_, - edge_type_, src, dst, edge); - if (!resolve.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(resolve)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to drop edge property from type: " << edge_type_ - << ", reason: " << resolve.ToString(); - RETURN_ERROR(resolve); - } - DeleteEdgePropertiesParamBuilder builder; - auto config = builder.DeleteProperties(property_names_).Build(); - auto res = storage.DeleteEdgeProperties(src, dst, edge, config); - if (!res.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(res)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to drop edge property from type: " << edge_type_ - << ", reason: " << res.ToString(); - RETURN_ERROR(res); - } - return neug::result(std::move(ctx)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); + + StorageUpdateInterface& storage = + dynamic_cast(graph); + label_t src, dst, edge; + auto resolve = + ResolveEdgeTriplet(storage.schema(), src_type_, dst_type_, + edge_type_, src, dst, edge); + if (!resolve.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(resolve)) { + return std::move(input); + } + LOG(ERROR) << "Fail to drop edge property from type: " << edge_type_ + << ", reason: " << resolve.ToString(); + return error_stream(resolve); + } + DeleteEdgePropertiesParamBuilder builder; + auto config = builder.DeleteProperties(property_names_).Build(); + auto res = storage.DeleteEdgeProperties(src, dst, edge, config); + if (!res.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(res)) { + return std::move(input); + } + LOG(ERROR) << "Fail to drop edge property from type: " << edge_type_ + << ", reason: " << res.ToString(); + return error_stream(res); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/ddl/drop_edge_type.cc b/src/execution/execute/ops/ddl/drop_edge_type.cc index 72d780a7a..f27b05222 100644 --- a/src/execution/execute/ops/ddl/drop_edge_type.cc +++ b/src/execution/execute/ops/ddl/drop_edge_type.cc @@ -31,31 +31,45 @@ class DropEdgeTypeOpr : public IOperator { ignore_conflict_(ignore_conflict) {} std::string get_operator_name() const override { return "DropEdgeTypeOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - StorageUpdateInterface& storage = - dynamic_cast(graph); - label_t src, dst, edge; - auto resolve = ResolveEdgeTriplet(storage.schema(), src_type_, dst_type_, - edge_type_, src, dst, edge); - if (!resolve.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(resolve)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to drop edge type: " << edge_type_ - << ", reason: " << resolve.ToString(); - RETURN_ERROR(resolve); - } - auto res = storage.DeleteEdgeType(src, dst, edge); - if (!res.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(res)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to drop edge type: " << edge_type_ - << ", reason: " << res.ToString(); - RETURN_ERROR(res); - } - return neug::result(std::move(ctx)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); + + StorageUpdateInterface& storage = + dynamic_cast(graph); + label_t src, dst, edge; + auto resolve = + ResolveEdgeTriplet(storage.schema(), src_type_, dst_type_, + edge_type_, src, dst, edge); + if (!resolve.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(resolve)) { + return std::move(input); + } + LOG(ERROR) << "Fail to drop edge type: " << edge_type_ + << ", reason: " << resolve.ToString(); + return error_stream(resolve); + } + auto res = storage.DeleteEdgeType(src, dst, edge); + if (!res.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(res)) { + return std::move(input); + } + LOG(ERROR) << "Fail to drop edge type: " << edge_type_ + << ", reason: " << res.ToString(); + return error_stream(res); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/ddl/drop_index.cc b/src/execution/execute/ops/ddl/drop_index.cc index d52f02c51..477ecae28 100644 --- a/src/execution/execute/ops/ddl/drop_index.cc +++ b/src/execution/execute/ops/ddl/drop_index.cc @@ -30,26 +30,40 @@ class DropIndexOpr : public IOperator { std::string get_operator_name() const override { return "DropIndexOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap&, - Context&& ctx, OprTimer*) override { - auto* indexInterface = dynamic_cast(&graph); - if (!indexInterface) { - RETURN_STATUS_ERROR( - StatusCode::ERR_NOT_SUPPORTED, - "Current storage interface does not support index DDL"); - } + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); - auto status = indexInterface->DropIndex(indexName_); - if (!status.ok()) { - // The storage layer reports ERR_NOT_FOUND when the target index does - // not exist; honor IF EXISTS in that case. - if (ignore_conflict_ && - status.error_code() == StatusCode::ERR_NOT_FOUND) { - return std::move(ctx); - } - RETURN_ERROR(status); - } - return std::move(ctx); + auto* indexInterface = + dynamic_cast(&graph); + if (!indexInterface) { + return error_stream( + Status(StatusCode::ERR_NOT_SUPPORTED, + "Current storage interface does not support index DDL")); + } + + auto status = indexInterface->DropIndex(indexName_); + if (!status.ok()) { + // The storage layer reports ERR_NOT_FOUND when the target index + // does not exist; honor IF EXISTS in that case. + if (ignore_conflict_ && + status.error_code() == StatusCode::ERR_NOT_FOUND) { + return std::move(input); + } + return error_stream(status); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/ddl/drop_vertex_property.cc b/src/execution/execute/ops/ddl/drop_vertex_property.cc index 4a42f8791..ad1d826b4 100644 --- a/src/execution/execute/ops/ddl/drop_vertex_property.cc +++ b/src/execution/execute/ops/ddl/drop_vertex_property.cc @@ -34,32 +34,46 @@ class DropVertexPropertySchemaOpr : public IOperator { return "DropVertexPropertySchemaOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - StorageUpdateInterface& storage = - dynamic_cast(graph); - label_t label; - auto resolve = ResolveVertexLabel(storage.schema(), vertex_type_, label); - if (!resolve.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(resolve)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to drop vertex property from type: " << vertex_type_ - << ", reason: " << resolve.ToString(); - RETURN_ERROR(resolve); - } - DeleteVertexPropertiesParamBuilder builder; - auto config = builder.DeleteProperties(property_names_).Build(); - auto res = storage.DeleteVertexProperties(label, config); - if (!res.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(res)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to drop vertex property from type: " << vertex_type_ - << ", reason: " << res.ToString(); - RETURN_ERROR(res); - } - return neug::result(std::move(ctx)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); + + StorageUpdateInterface& storage = + dynamic_cast(graph); + label_t label; + auto resolve = + ResolveVertexLabel(storage.schema(), vertex_type_, label); + if (!resolve.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(resolve)) { + return std::move(input); + } + LOG(ERROR) << "Fail to drop vertex property from type: " + << vertex_type_ << ", reason: " << resolve.ToString(); + return error_stream(resolve); + } + DeleteVertexPropertiesParamBuilder builder; + auto config = builder.DeleteProperties(property_names_).Build(); + auto res = storage.DeleteVertexProperties(label, config); + if (!res.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(res)) { + return std::move(input); + } + LOG(ERROR) << "Fail to drop vertex property from type: " + << vertex_type_ << ", reason: " << res.ToString(); + return error_stream(res); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/ddl/drop_vertex_type.cc b/src/execution/execute/ops/ddl/drop_vertex_type.cc index dbe83e134..f2942b63d 100644 --- a/src/execution/execute/ops/ddl/drop_vertex_type.cc +++ b/src/execution/execute/ops/ddl/drop_vertex_type.cc @@ -27,30 +27,44 @@ class DropVertexTypeOpr : public IOperator { : vertex_type_(vertex_type), ignore_conflict_(ignore_conflict) {} std::string get_operator_name() const override { return "DropVertexTypeOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - StorageUpdateInterface& storage = - dynamic_cast(graph); - label_t label; - auto resolve = ResolveVertexLabel(storage.schema(), vertex_type_, label); - if (!resolve.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(resolve)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to drop vertex type: " << vertex_type_ - << ", reason: " << resolve.ToString(); - RETURN_ERROR(resolve); - } - auto res = storage.DeleteVertexType(label); - if (!res.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(res)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to drop vertex type: " << vertex_type_ - << ", reason: " << res.ToString(); - RETURN_ERROR(res); - } - return neug::result(std::move(ctx)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); + + StorageUpdateInterface& storage = + dynamic_cast(graph); + label_t label; + auto resolve = + ResolveVertexLabel(storage.schema(), vertex_type_, label); + if (!resolve.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(resolve)) { + return std::move(input); + } + LOG(ERROR) << "Fail to drop vertex type: " << vertex_type_ + << ", reason: " << resolve.ToString(); + return error_stream(resolve); + } + auto res = storage.DeleteVertexType(label); + if (!res.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(res)) { + return std::move(input); + } + LOG(ERROR) << "Fail to drop vertex type: " << vertex_type_ + << ", reason: " << res.ToString(); + return error_stream(res); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/ddl/rename_edge_property.cc b/src/execution/execute/ops/ddl/rename_edge_property.cc index 3e10eec06..bb60fa59e 100644 --- a/src/execution/execute/ops/ddl/rename_edge_property.cc +++ b/src/execution/execute/ops/ddl/rename_edge_property.cc @@ -36,33 +36,47 @@ class RenameEdgePropertySchemaOpr : public IOperator { std::string get_operator_name() const override { return "RenameEdgePropertySchemaOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - StorageUpdateInterface& storage = - dynamic_cast(graph); - label_t src, dst, edge; - auto resolve = ResolveEdgeTriplet(storage.schema(), src_type_, dst_type_, - edge_type_, src, dst, edge); - if (!resolve.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(resolve)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to rename edge property in type: " << edge_type_ - << ", reason: " << resolve.ToString(); - RETURN_ERROR(resolve); - } - RenameEdgePropertiesParamBuilder builder; - auto config = builder.RenameProperties(rename_properties_).Build(); - auto res = storage.RenameEdgeProperties(src, dst, edge, config); - if (!res.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(res)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to rename edge property in type: " << edge_type_ - << ", reason: " << res.ToString(); - RETURN_ERROR(res); - } - return neug::result(std::move(ctx)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); + + StorageUpdateInterface& storage = + dynamic_cast(graph); + label_t src, dst, edge; + auto resolve = + ResolveEdgeTriplet(storage.schema(), src_type_, dst_type_, + edge_type_, src, dst, edge); + if (!resolve.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(resolve)) { + return std::move(input); + } + LOG(ERROR) << "Fail to rename edge property in type: " << edge_type_ + << ", reason: " << resolve.ToString(); + return error_stream(resolve); + } + RenameEdgePropertiesParamBuilder builder; + auto config = builder.RenameProperties(rename_properties_).Build(); + auto res = storage.RenameEdgeProperties(src, dst, edge, config); + if (!res.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(res)) { + return std::move(input); + } + LOG(ERROR) << "Fail to rename edge property in type: " << edge_type_ + << ", reason: " << res.ToString(); + return error_stream(res); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/ddl/rename_vertex_property.cc b/src/execution/execute/ops/ddl/rename_vertex_property.cc index e39e6a451..bf5d35a46 100644 --- a/src/execution/execute/ops/ddl/rename_vertex_property.cc +++ b/src/execution/execute/ops/ddl/rename_vertex_property.cc @@ -35,32 +35,46 @@ class RenameVertexPropertyOpr : public IOperator { return "RenameVertexPropertyOpr"; } - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - StorageUpdateInterface& storage = - dynamic_cast(graph); - label_t label; - auto resolve = ResolveVertexLabel(storage.schema(), vertex_type_, label); - if (!resolve.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(resolve)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to rename vertex property in type: " << vertex_type_ - << ", reason: " << resolve.ToString(); - RETURN_ERROR(resolve); - } - RenameVertexPropertiesParamBuilder builder; - auto config = builder.RenameProperties(rename_properties_).Build(); - auto res = storage.RenameVertexProperties(label, config); - if (!res.ok()) { - if (ignore_conflict_ && IsSchemaConflictError(res)) { - return neug::result(std::move(ctx)); - } - LOG(ERROR) << "Fail to rename vertex property in type: " << vertex_type_ - << ", reason: " << res.ToString(); - RETURN_ERROR(res); - } - return neug::result(std::move(ctx)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto before = collect_batches(std::move(input)); + if (!before) { + return error_stream(before.error()); + } + input = stream_from_batches(std::move(*before), std::move(metadata)); + + StorageUpdateInterface& storage = + dynamic_cast(graph); + label_t label; + auto resolve = + ResolveVertexLabel(storage.schema(), vertex_type_, label); + if (!resolve.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(resolve)) { + return std::move(input); + } + LOG(ERROR) << "Fail to rename vertex property in type: " + << vertex_type_ << ", reason: " << resolve.ToString(); + return error_stream(resolve); + } + RenameVertexPropertiesParamBuilder builder; + auto config = builder.RenameProperties(rename_properties_).Build(); + auto res = storage.RenameVertexProperties(label, config); + if (!res.ok()) { + if (ignore_conflict_ && IsSchemaConflictError(res)) { + return std::move(input); + } + LOG(ERROR) << "Fail to rename vertex property in type: " + << vertex_type_ << ", reason: " << res.ToString(); + return error_stream(res); + } + return std::move(input); + }); } private: diff --git a/src/execution/execute/ops/edge_column_rebuild.cc b/src/execution/execute/ops/edge_column_rebuild.cc index 47eafdd87..c591b9c93 100644 --- a/src/execution/execute/ops/edge_column_rebuild.cc +++ b/src/execution/execute/ops/edge_column_rebuild.cc @@ -81,7 +81,7 @@ void RefreshEdgeRecord(StorageUpdateInterface& graph, EdgeRecord& record, } EdgeColumnSnapshots CaptureEdgeColumnsForRefresh( - StorageUpdateInterface& graph, Context& ctx, + StorageUpdateInterface& graph, std::vector& chunks, const std::set& affected_labels) { EdgeColumnSnapshots snapshots; if (affected_labels.empty()) { @@ -101,7 +101,7 @@ EdgeColumnSnapshots CaptureEdgeColumnsForRefresh( snapshots.columns[it->second].aliases.push_back(&column); }; - for (auto& chunk : ctx.chunks()) { + for (auto& chunk : chunks) { for (auto& column : chunk.columns()) { capture(column); } diff --git a/src/execution/execute/ops/insert/create_edge.cc b/src/execution/execute/ops/insert/create_edge.cc index 84c2f5a8c..382e075f6 100644 --- a/src/execution/execute/ops/insert/create_edge.cc +++ b/src/execution/execute/ops/insert/create_edge.cc @@ -36,37 +36,51 @@ class CreateEdgeOpr : public IOperator { src_dst_tags_(src_dst_tags), properties_(std::move(properties)) {} - neug::result Eval(IStorageInterface& graph_interface, - const ParamsMap& params, Context&& ctx, - OprTimer* timer) override { - const StorageReadInterface* graph_ptr = nullptr; - if (graph_interface.readable()) { - graph_ptr = dynamic_cast(&graph_interface); - } - std::vector< - std::vector>>> - expr_properties; - for (size_t i = 0; i < labels_.size(); ++i) { - const auto& props = properties_[i]; - std::vector>> - expr_props; - for (const auto& [prop, prop_value] : props) { - auto expr = prop_value->bind(graph_ptr, params); - expr_props.emplace_back(prop, std::move(expr)); - } - expr_properties.emplace_back(std::move(expr_props)); - } - // TODO(liulx20,zhanglei): CREATE on bundled edges may detach or grow CSR - // storage, - // leaving edge-property pointers in other chunks stale. Preserve the - // chunk-oriented apply_chunks path for now; track a compatible fix at - // https://github.com/alibaba/neug/issues/927. - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return CreateEdge::insert_edge( - dynamic_cast(graph_interface), - std::move(chunk), labels_, src_dst_tags_, - std::move(expr_properties), alias_); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Finish reading before mutation; downstream cancellation must not + // skip writes. + return reduce_stream( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const StorageReadInterface* graph_ptr = nullptr; + if (graph_interface.readable()) { + graph_ptr = dynamic_cast( + &graph_interface); + } + std::vector>>> + expr_properties; + for (size_t i = 0; i < labels_.size(); ++i) { + const auto& props = properties_[i]; + std::vector< + std::pair>> + expr_props; + for (const auto& [prop, prop_value] : props) { + auto expr = prop_value->bind(graph_ptr, params); + expr_props.emplace_back(prop, std::move(expr)); + } + expr_properties.emplace_back(std::move(expr_props)); + } + // TODO(liulx20,zhanglei): CREATE on bundled edges may detach or + // grow CSR storage, leaving edge-property pointers in other + // chunks stale. Preserve the chunk-oriented apply_chunks path + // for now; track a compatible fix at + // https://github.com/alibaba/neug/issues/927. + { + return CreateEdge::insert_edge( + dynamic_cast(graph_interface), + std::move(chunk), labels_, src_dst_tags_, + std::move(expr_properties), alias_); + } + }); }); } std::string get_operator_name() const override { return "CreateEdgeOpr"; } diff --git a/src/execution/execute/ops/insert/create_vertex.cc b/src/execution/execute/ops/insert/create_vertex.cc index 209034ba1..aa03c3616 100644 --- a/src/execution/execute/ops/insert/create_vertex.cc +++ b/src/execution/execute/ops/insert/create_vertex.cc @@ -33,33 +33,48 @@ class CreateVertexOpr : public IOperator { properties) : labels_(labels), alias_(alias), properties_(std::move(properties)) {} - neug::result Eval(IStorageInterface& graph_interface, - const ParamsMap& params, Context&& ctx, - OprTimer* timer) override { - // Implementation of vertex creation logic goes here. + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Finish reading before mutation; downstream cancellation must not + // skip writes. + return reduce_stream( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + // Implementation of vertex creation logic goes here. - const StorageReadInterface* graph_ptr = nullptr; - if (graph_interface.readable()) { - graph_ptr = dynamic_cast(&graph_interface); - } - std::vector< - std::vector>>> - expr_properties; - for (size_t i = 0; i < labels_.size(); ++i) { - const auto& props = properties_[i]; - std::vector>> - expr_props; - for (auto& [prop, prop_value] : props) { - auto expr = prop_value->bind(graph_ptr, params); - expr_props.emplace_back(prop, std::move(expr)); - } - expr_properties.emplace_back(std::move(expr_props)); - } - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return CreateVertex::insert_vertex( - dynamic_cast(graph_interface), - std::move(chunk), labels_, std::move(expr_properties), alias_); + const StorageReadInterface* graph_ptr = nullptr; + if (graph_interface.readable()) { + graph_ptr = dynamic_cast( + &graph_interface); + } + std::vector>>> + expr_properties; + for (size_t i = 0; i < labels_.size(); ++i) { + const auto& props = properties_[i]; + std::vector< + std::pair>> + expr_props; + for (auto& [prop, prop_value] : props) { + auto expr = prop_value->bind(graph_ptr, params); + expr_props.emplace_back(prop, std::move(expr)); + } + expr_properties.emplace_back(std::move(expr_props)); + } + { + return CreateVertex::insert_vertex( + dynamic_cast(graph_interface), + std::move(chunk), labels_, std::move(expr_properties), + alias_); + } + }); }); } std::string get_operator_name() const override { return "CreateVertexOpr"; } diff --git a/src/execution/execute/ops/insert/merge_edge.cc b/src/execution/execute/ops/insert/merge_edge.cc index 7171c190e..3c1e800e6 100644 --- a/src/execution/execute/ops/insert/merge_edge.cc +++ b/src/execution/execute/ops/insert/merge_edge.cc @@ -186,211 +186,236 @@ class MergeEdgeOpr : public IOperator { std::string get_operator_name() const override { return "MergeEdgeOpr"; } - neug::result Eval(IStorageInterface& graph_interface, - const ParamsMap& params, Context&& ctx, - OprTimer* timer) override { - (void) timer; - auto& graph = dynamic_cast(graph_interface); - const StorageReadInterface* graph_read = nullptr; - if (graph_interface.readable()) { - graph_read = dynamic_cast(&graph_interface); - } - - for (const auto& plan : entries_) { - std::vector>> - pattern_binded; - std::vector>> - on_create_binded; - std::vector>> - on_match_binded; - for (const auto& [n, e] : plan.pattern_props) { - pattern_binded.emplace_back(n, e->bind(graph_read, params)); - } - for (const auto& [n, e] : plan.on_create_props) { - on_create_binded.emplace_back(n, e->bind(graph_read, params)); - } - for (const auto& [n, e] : plan.on_match_props) { - on_match_binded.emplace_back(n, e->bind(graph_read, params)); - } - auto merged_binded = merge_pattern_and_on_create( - std::move(pattern_binded), std::move(on_create_binded)); - - const auto edge_schema = graph.schema().get_edge_schema( - plan.labels.src_label, plan.labels.dst_label, plan.labels.edge_label); - const bool bundled = edge_schema->is_bundled(); - bool has_unmatched_rows = false; - if (bundled && on_match_binded.empty()) { - for (auto& chunk : ctx.chunks()) { - if (chunk.row_num() == 0) { - continue; + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Edge property writes can invalidate pointers in later batches. + // Retain and refresh every affected column before exposing the result + // downstream. + auto metadata = input.metadata(); + auto chunks_result = collect_batches(std::move(input)); + if (!chunks_result) { + return error_stream(chunks_result.error()); } - auto edge_column = - std::dynamic_pointer_cast(chunk.get(plan.alias_id)); - if (!edge_column) { - continue; - } - for (size_t row = 0; row < edge_column->size(); ++row) { - if (!edge_column->has_value(row) || - edge_column->get_edge(row).label != plan.labels) { - has_unmatched_rows = true; - break; - } - } - if (has_unmatched_rows) { - break; - } - } - } - std::set affected_labels; - if (bundled && (!on_match_binded.empty() || has_unmatched_rows)) { - affected_labels.insert(plan.labels); - } - auto snapshots = - CaptureEdgeColumnsForRefresh(graph, ctx, affected_labels); + auto chunks = std::move(*chunks_result); - struct PendingRow { - EdgeRecord record; - std::optional insert; - }; - struct MatchChunk { - ContextChunk* chunk; - const EdgeColumnSnapshot* alias_snapshot; - std::vector rows; - }; - std::vector matched_chunks; - for (auto& chunk : ctx.chunks()) { - const auto nrows = chunk.row_num(); - MatchChunk matched_chunk{&chunk, nullptr, - std::vector(nrows)}; - if (nrows == 0) { - matched_chunks.push_back(std::move(matched_chunk)); - continue; - } - if (!chunk.exist(plan.alias_id)) { - THROW_RUNTIME_ERROR( - "MERGE edge requires the pattern edge alias in context " - "(missing column for alias id " + - std::to_string(plan.alias_id) + ")"); - } - auto alias_column = chunk.get(plan.alias_id); - if (!alias_column || alias_column->size() != nrows) { - THROW_RUNTIME_ERROR( - "MERGE edge alias column size does not match " - "context row count"); - } - auto edge_column = std::dynamic_pointer_cast(alias_column); - if (!edge_column) { - THROW_RUNTIME_ERROR( - "MERGE edge pattern alias must refer to an edge column (alias " - "id " + - std::to_string(plan.alias_id) + ")"); - } - matched_chunk.alias_snapshot = snapshots.Find(edge_column.get()); - const auto& src_vertex_col = dynamic_cast( - *chunk.get(plan.src_dst_tags.first).get()); - const auto& dst_vertex_col = dynamic_cast( - *chunk.get(plan.src_dst_tags.second).get()); - for (size_t row = 0; row < nrows; ++row) { - auto& pending = matched_chunk.rows[row]; - bool matched = false; - if (edge_column->has_value(row)) { - pending.record = matched_chunk.alias_snapshot == nullptr - ? edge_column->get_edge(row) - : matched_chunk.alias_snapshot->records[row]; - matched = pending.record.label == plan.labels; + (void) timer; + auto& graph = dynamic_cast(graph_interface); + const StorageReadInterface* graph_read = nullptr; + if (graph_interface.readable()) { + graph_read = + dynamic_cast(&graph_interface); } - if (!matched) { - pending.insert = prepare_edge_insert( - graph, chunk.chunk(), row, plan.labels.src_label, - plan.labels.dst_label, plan.labels.edge_label, src_vertex_col, - dst_vertex_col, merged_binded); - } - } - matched_chunks.push_back(std::move(matched_chunk)); - } - for (const auto& [prop_name, expression] : on_match_binded) { - const auto property_id = edge_schema->get_property_index(prop_name); - if (property_id < 0) { - THROW_RUNTIME_ERROR( - "Property " + prop_name + " does not exist for edge label " + - std::to_string(static_cast(plan.labels.edge_label))); - } - std::vector mutations; - for (auto& matched_chunk : matched_chunks) { - auto& chunk = *matched_chunk.chunk; - for (size_t row = 0; row < matched_chunk.rows.size(); ++row) { - auto& pending = matched_chunk.rows[row]; - if (pending.insert) { - continue; + for (const auto& plan : entries_) { + std::vector>> + pattern_binded; + std::vector>> + on_create_binded; + std::vector>> + on_match_binded; + for (const auto& [n, e] : plan.pattern_props) { + pattern_binded.emplace_back(n, e->bind(graph_read, params)); } - std::pair offsets; - if (matched_chunk.alias_snapshot != nullptr && - matched_chunk.alias_snapshot->refresh_rows[row]) { - pending.record = matched_chunk.alias_snapshot->records[row]; - offsets = matched_chunk.alias_snapshot->offsets[row]; - } else { - offsets = ResolveEdgeOffsets(graph, pending.record); + for (const auto& [n, e] : plan.on_create_props) { + on_create_binded.emplace_back(n, e->bind(graph_read, params)); } - auto value = expression->Cast().eval_record( - chunk.chunk(), row); - if (edge_schema->properties[property_id] != value.type()) { - THROW_RUNTIME_ERROR("Property type mismatch for property " + - prop_name); + for (const auto& [n, e] : plan.on_match_props) { + on_match_binded.emplace_back(n, e->bind(graph_read, params)); } - mutations.push_back(EdgePropertyMutation{ - pending.record, offsets, property_id, std::move(value)}); - } - } - for (const auto& mutation : mutations) { - auto status = graph.UpdateEdgeProperty( - mutation.record.label.src_label, mutation.record.src, - mutation.record.label.dst_label, mutation.record.dst, - mutation.record.label.edge_label, mutation.offsets.first, - mutation.offsets.second, mutation.property_id, mutation.value); - if (!status.ok()) { - THROW_RUNTIME_ERROR(status.ToString()); - } - } - if (bundled && !mutations.empty()) { - RefreshEdgeColumns(graph, snapshots); - } - } + auto merged_binded = merge_pattern_and_on_create( + std::move(pattern_binded), std::move(on_create_binded)); - for (auto& matched_chunk : matched_chunks) { - for (auto& pending : matched_chunk.rows) { - if (!pending.insert) { - continue; - } - pending.record = - apply_edge_insert(graph, plan.labels, *pending.insert); - } - } + const auto edge_schema = graph.schema().get_edge_schema( + plan.labels.src_label, plan.labels.dst_label, + plan.labels.edge_label); + const bool bundled = edge_schema->is_bundled(); + bool has_unmatched_rows = false; + if (bundled && on_match_binded.empty()) { + for (auto& chunk : chunks) { + if (chunk.row_num() == 0) { + continue; + } + auto edge_column = std::dynamic_pointer_cast( + chunk.get(plan.alias_id)); + if (!edge_column) { + continue; + } + for (size_t row = 0; row < edge_column->size(); ++row) { + if (!edge_column->has_value(row) || + edge_column->get_edge(row).label != plan.labels) { + has_unmatched_rows = true; + break; + } + } + if (has_unmatched_rows) { + break; + } + } + } + std::set affected_labels; + if (bundled && (!on_match_binded.empty() || has_unmatched_rows)) { + affected_labels.insert(plan.labels); + } + auto snapshots = + CaptureEdgeColumnsForRefresh(graph, chunks, affected_labels); - RefreshEdgeColumns(graph, snapshots); + struct PendingRow { + EdgeRecord record; + std::optional insert; + }; + struct MatchChunk { + ContextChunk* chunk; + const EdgeColumnSnapshot* alias_snapshot; + std::vector rows; + }; + std::vector matched_chunks; + for (auto& chunk : chunks) { + const auto nrows = chunk.row_num(); + MatchChunk matched_chunk{&chunk, nullptr, + std::vector(nrows)}; + if (nrows == 0) { + matched_chunks.push_back(std::move(matched_chunk)); + continue; + } + if (!chunk.exist(plan.alias_id)) { + THROW_RUNTIME_ERROR( + "MERGE edge requires the pattern edge alias in context " + "(missing column for alias id " + + std::to_string(plan.alias_id) + ")"); + } + auto alias_column = chunk.get(plan.alias_id); + if (!alias_column || alias_column->size() != nrows) { + THROW_RUNTIME_ERROR( + "MERGE edge alias column size does not match " + "context row count"); + } + auto edge_column = + std::dynamic_pointer_cast(alias_column); + if (!edge_column) { + THROW_RUNTIME_ERROR( + "MERGE edge pattern alias must refer to an edge column " + "(alias " + "id " + + std::to_string(plan.alias_id) + ")"); + } + matched_chunk.alias_snapshot = snapshots.Find(edge_column.get()); + const auto& src_vertex_col = dynamic_cast( + *chunk.get(plan.src_dst_tags.first).get()); + const auto& dst_vertex_col = dynamic_cast( + *chunk.get(plan.src_dst_tags.second).get()); + for (size_t row = 0; row < nrows; ++row) { + auto& pending = matched_chunk.rows[row]; + bool matched = false; + if (edge_column->has_value(row)) { + pending.record = + matched_chunk.alias_snapshot == nullptr + ? edge_column->get_edge(row) + : matched_chunk.alias_snapshot->records[row]; + matched = pending.record.label == plan.labels; + } + if (!matched) { + pending.insert = prepare_edge_insert( + graph, chunk.chunk(), row, plan.labels.src_label, + plan.labels.dst_label, plan.labels.edge_label, + src_vertex_col, dst_vertex_col, merged_binded); + } + } + matched_chunks.push_back(std::move(matched_chunk)); + } - for (auto& matched_chunk : matched_chunks) { - auto& chunk = *matched_chunk.chunk; - SDSLEdgeColumnBuilder builder(Direction::kOut, plan.labels); - for (size_t row = 0; row < matched_chunk.rows.size(); ++row) { - auto& pending = matched_chunk.rows[row]; - if (!pending.insert && matched_chunk.alias_snapshot != nullptr && - matched_chunk.alias_snapshot->refresh_rows[row]) { - pending.record = matched_chunk.alias_snapshot->records[row]; - } else if (pending.insert) { - RefreshEdgeRecord(graph, pending.record, - ResolveEdgeOffsets(graph, pending.record)); + for (const auto& [prop_name, expression] : on_match_binded) { + const auto property_id = + edge_schema->get_property_index(prop_name); + if (property_id < 0) { + THROW_RUNTIME_ERROR( + "Property " + prop_name + + " does not exist for edge label " + + std::to_string(static_cast(plan.labels.edge_label))); + } + std::vector mutations; + for (auto& matched_chunk : matched_chunks) { + auto& chunk = *matched_chunk.chunk; + for (size_t row = 0; row < matched_chunk.rows.size(); ++row) { + auto& pending = matched_chunk.rows[row]; + if (pending.insert) { + continue; + } + std::pair offsets; + if (matched_chunk.alias_snapshot != nullptr && + matched_chunk.alias_snapshot->refresh_rows[row]) { + pending.record = matched_chunk.alias_snapshot->records[row]; + offsets = matched_chunk.alias_snapshot->offsets[row]; + } else { + offsets = ResolveEdgeOffsets(graph, pending.record); + } + auto value = expression->Cast().eval_record( + chunk.chunk(), row); + if (edge_schema->properties[property_id] != value.type()) { + THROW_RUNTIME_ERROR("Property type mismatch for property " + + prop_name); + } + mutations.push_back(EdgePropertyMutation{ + pending.record, offsets, property_id, std::move(value)}); + } + } + for (const auto& mutation : mutations) { + auto status = graph.UpdateEdgeProperty( + mutation.record.label.src_label, mutation.record.src, + mutation.record.label.dst_label, mutation.record.dst, + mutation.record.label.edge_label, mutation.offsets.first, + mutation.offsets.second, mutation.property_id, + mutation.value); + if (!status.ok()) { + THROW_RUNTIME_ERROR(status.ToString()); + } + } + if (bundled && !mutations.empty()) { + RefreshEdgeColumns(graph, snapshots); + } + } + + for (auto& matched_chunk : matched_chunks) { + for (auto& pending : matched_chunk.rows) { + if (!pending.insert) { + continue; + } + pending.record = + apply_edge_insert(graph, plan.labels, *pending.insert); + } + } + + RefreshEdgeColumns(graph, snapshots); + + for (auto& matched_chunk : matched_chunks) { + auto& chunk = *matched_chunk.chunk; + SDSLEdgeColumnBuilder builder(Direction::kOut, plan.labels); + for (size_t row = 0; row < matched_chunk.rows.size(); ++row) { + auto& pending = matched_chunk.rows[row]; + if (!pending.insert && + matched_chunk.alias_snapshot != nullptr && + matched_chunk.alias_snapshot->refresh_rows[row]) { + pending.record = matched_chunk.alias_snapshot->records[row]; + } else if (pending.insert) { + RefreshEdgeRecord(graph, pending.record, + ResolveEdgeOffsets(graph, pending.record)); + } + builder.push_back_opt(pending.record.src, pending.record.dst, + pending.record.prop); + } + if (chunk.exist(plan.alias_id)) { + chunk.remove(plan.alias_id); + } + chunk.set(plan.alias_id, builder.finish()); + } } - builder.push_back_opt(pending.record.src, pending.record.dst, - pending.record.prop); - } - if (chunk.exist(plan.alias_id)) { - chunk.remove(plan.alias_id); - } - chunk.set(plan.alias_id, builder.finish()); - } - } - return std::move(ctx); + return stream_from_batches(std::move(chunks), std::move(metadata)); + }); } private: diff --git a/src/execution/execute/ops/insert/merge_vertex.cc b/src/execution/execute/ops/insert/merge_vertex.cc index 13cf66c80..bd364319d 100644 --- a/src/execution/execute/ops/insert/merge_vertex.cc +++ b/src/execution/execute/ops/insert/merge_vertex.cc @@ -198,84 +198,108 @@ class MergeVertexOpr : public IOperator { std::string get_operator_name() const override { return "MergeVertexOpr"; } - neug::result Eval(IStorageInterface& graph_interface, - const ParamsMap& params, Context&& ctx, - OprTimer* timer) override { - (void) timer; - auto& graph = dynamic_cast(graph_interface); - const StorageReadInterface* graph_read = nullptr; - if (graph_interface.readable()) { - graph_read = dynamic_cast(&graph_interface); - } + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Finish reading before mutation; downstream cancellation must not + // skip writes. + return reduce_stream( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + (void) timer; + auto& graph = + dynamic_cast(graph_interface); + const StorageReadInterface* graph_read = nullptr; + if (graph_interface.readable()) { + graph_read = dynamic_cast( + &graph_interface); + } - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - for (const auto& plan : entries_) { - std::vector>> - pattern_binded; - std::vector>> - on_create_binded; - std::vector>> - on_match_binded; - for (const auto& [n, e] : plan.pattern_props) { - pattern_binded.emplace_back(n, e->bind(graph_read, params)); - } - for (const auto& [n, e] : plan.on_create_props) { - on_create_binded.emplace_back(n, e->bind(graph_read, params)); - } - for (const auto& [n, e] : plan.on_match_props) { - on_match_binded.emplace_back(n, e->bind(graph_read, params)); - } - auto merged_binded = merge_pattern_and_on_create( - std::move(pattern_binded), std::move(on_create_binded)); + { + for (const auto& plan : entries_) { + std::vector< + std::pair>> + pattern_binded; + std::vector< + std::pair>> + on_create_binded; + std::vector< + std::pair>> + on_match_binded; + for (const auto& [n, e] : plan.pattern_props) { + pattern_binded.emplace_back(n, + e->bind(graph_read, params)); + } + for (const auto& [n, e] : plan.on_create_props) { + on_create_binded.emplace_back( + n, e->bind(graph_read, params)); + } + for (const auto& [n, e] : plan.on_match_props) { + on_match_binded.emplace_back(n, + e->bind(graph_read, params)); + } + auto merged_binded = merge_pattern_and_on_create( + std::move(pattern_binded), std::move(on_create_binded)); - MSVertexColumnBuilder builder(plan.label); + MSVertexColumnBuilder builder(plan.label); - // Standalone MERGE after OPTIONAL MATCH can yield row_num() == 0 - // when the inner scan finds no row. MERGE write semantics still - // need exactly one logical row (CREATE/MATCH branch once). - std::shared_ptr alias_col; - if (chunk.exist(plan.alias_id)) { - auto c = chunk.get(plan.alias_id); - if (c != nullptr && c->size() > 0) { - alias_col = std::move(c); - } - } - size_t num_rows = chunk.row_num(); - if (num_rows == 0) { - num_rows = 1; - } + // Standalone MERGE after OPTIONAL MATCH can yield row_num() + // == 0 when the inner scan finds no row. MERGE write + // semantics still need exactly one logical row + // (CREATE/MATCH branch once). + std::shared_ptr alias_col; + if (chunk.exist(plan.alias_id)) { + auto c = chunk.get(plan.alias_id); + if (c != nullptr && c->size() > 0) { + alias_col = std::move(c); + } + } + size_t num_rows = chunk.row_num(); + if (num_rows == 0) { + num_rows = 1; + } - for (size_t row = 0; row < num_rows; ++row) { - bool matched = false; - vid_t matched_vid = 0; - if (alias_col) { - auto vc = std::dynamic_pointer_cast(alias_col); - if (vc && row < vc->size() && vc->has_value(row)) { - auto vr = vc->get_vertex(row); - if (vr.label_ == plan.label) { - matched = true; - matched_vid = vr.vid_; + for (size_t row = 0; row < num_rows; ++row) { + bool matched = false; + vid_t matched_vid = 0; + if (alias_col) { + auto vc = + std::dynamic_pointer_cast(alias_col); + if (vc && row < vc->size() && vc->has_value(row)) { + auto vr = vc->get_vertex(row); + if (vr.label_ == plan.label) { + matched = true; + matched_vid = vr.vid_; + } + } + } + if (matched) { + apply_on_match_vertex(graph, chunk.chunk(), row, + plan.label, matched_vid, + on_match_binded); + builder.push_back_opt(matched_vid); + } else { + vid_t vid; + GS_ASSIGN(vid, + insert_vertex_row(graph, chunk.chunk(), row, + plan.label, merged_binded)); + builder.push_back_opt(vid); + } + } + if (chunk.exist(plan.alias_id)) { + chunk.remove(plan.alias_id); + } + chunk.set(plan.alias_id, builder.finish()); } + return chunk; } - } - if (matched) { - apply_on_match_vertex(graph, chunk.chunk(), row, plan.label, - matched_vid, on_match_binded); - builder.push_back_opt(matched_vid); - } else { - vid_t vid; - GS_ASSIGN(vid, insert_vertex_row(graph, chunk.chunk(), row, - plan.label, merged_binded)); - builder.push_back_opt(vid); - } - } - if (chunk.exist(plan.alias_id)) { - chunk.remove(plan.alias_id); - } - chunk.set(plan.alias_id, builder.finish()); - } - return chunk; + }); }); } diff --git a/src/execution/execute/ops/retrieve/dedup.cc b/src/execution/execute/ops/retrieve/dedup.cc index da649ddcd..b67868e17 100644 --- a/src/execution/execute/ops/retrieve/dedup.cc +++ b/src/execution/execute/ops/retrieve/dedup.cc @@ -38,15 +38,14 @@ class DedupOpr : public IOperator { explicit DedupOpr(const std::vector& tag_ids) : tag_ids_(tag_ids) {} std::string get_operator_name() const override { return "DedupOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - ctx.ensure_single_chunk("DedupOpr"); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Dedup::dedup(std::move(chunk), tag_ids_); - }); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return reduce_stream(std::move(input), + [this, &graph, params, + timer](ContextChunk&& chunk) -> result { + { return Dedup::dedup(std::move(chunk), tag_ids_); } + }); } std::vector tag_ids_; diff --git a/src/execution/execute/ops/retrieve/edge.cc b/src/execution/execute/ops/retrieve/edge.cc index 358264e36..27c041c09 100644 --- a/src/execution/execute/ops/retrieve/edge.cc +++ b/src/execution/execute/ops/retrieve/edge.cc @@ -102,33 +102,35 @@ class EdgeExpandVWithEPCmpOpr : public IOperator { return "EdgeExpandVWithEPCmpOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - if ((!eep_.is_optional) && - (config_.ptype == SPPredicateType::kPropertyLT || - config_.ptype == SPPredicateType::kPropertyGT)) { - const auto& param_value = params.at(config_.param_names[0]); - auto ret = ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::expand_vertex_ep_cmp( - graph, std::move(chunk), eep_, param_value, config_.ptype); - }); - if (ret) { - return ret.value(); - } - } + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + if ((!eep_.is_optional) && + (config_.ptype == SPPredicateType::kPropertyLT || + config_.ptype == SPPredicateType::kPropertyGT)) { + const auto& param_value = params.at(config_.param_names[0]); + auto candidate = chunk; + auto ret = EdgeExpand::expand_vertex_ep_cmp( + graph, std::move(candidate), eep_, param_value, config_.ptype); + if (ret) { + return ret.value(); + } + } - auto expr = pred_->bind(&graph, params); - GeneralPred expr_wrapper(std::move(expr)); - EdgePredicate pred(expr_wrapper); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::expand_vertex( - graph, std::move(chunk), eep_, pred); + auto expr = pred_->bind(&graph, params); + GeneralPred expr_wrapper(std::move(expr)); + EdgePredicate pred(expr_wrapper); + { + return EdgeExpand::expand_vertex( + graph, std::move(chunk), eep_, pred); + } }); } @@ -145,29 +147,32 @@ class EdgeExpandVOpr : public IOperator { std::string get_operator_name() const override { return "EdgeExpandVOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - - if (pred_ != nullptr) { - auto expr = pred_->bind(&graph, params); - GeneralPred expr_wrapper(std::move(expr)); - EdgePredicate pred(expr_wrapper); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::expand_vertex( - graph, std::move(chunk), eep_, pred); - }); - } else { - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::expand_vertex(graph, std::move(chunk), - eep_, DummyPred()); - }); - } + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + + if (pred_ != nullptr) { + auto expr = pred_->bind(&graph, params); + GeneralPred expr_wrapper(std::move(expr)); + EdgePredicate pred(expr_wrapper); + { + return EdgeExpand::expand_vertex( + graph, std::move(chunk), eep_, pred); + } + } else { + { + return EdgeExpand::expand_vertex( + graph, std::move(chunk), eep_, DummyPred()); + } + } + }); } private: @@ -185,17 +190,21 @@ class EdgeExpandEWithSPredOpr : public IOperator { return "EdgeExpandEWithSPredOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::expand_edge_with_special_edge_predicate( - graph, std::move(chunk), eep_, config_, - params.at(config_.param_names[0])); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + { + return EdgeExpand::expand_edge_with_special_edge_predicate( + graph, std::move(chunk), eep_, config_, + params.at(config_.param_names[0])); + } }); } @@ -211,28 +220,31 @@ class EdgeExpandEOpr : public IOperator { std::string get_operator_name() const override { return "EdgeExpandEOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - if (pred_ == nullptr) { - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::expand_edge(graph, std::move(chunk), eep_, - DummyPred()); - }); - } else { - auto expr = pred_->bind(&graph, params); - GeneralPred expr_wrapper(std::move(expr)); - EdgePredicate pred(expr_wrapper); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::expand_edge( - graph, std::move(chunk), eep_, pred); - }); - } + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + if (pred_ == nullptr) { + { + return EdgeExpand::expand_edge(graph, std::move(chunk), eep_, + DummyPred()); + } + } else { + auto expr = pred_->bind(&graph, params); + GeneralPred expr_wrapper(std::move(expr)); + EdgePredicate pred(expr_wrapper); + { + return EdgeExpand::expand_edge( + graph, std::move(chunk), eep_, pred); + } + } + }); } private: @@ -250,16 +262,20 @@ class EdgeExpandVWithSPVertexPredOpr : public IOperator { return "EdgeExpandVWithSPVertexPredOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::expand_vertex_with_special_vertex_predicate( - graph, std::move(chunk), eep_, config_, params); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + { + return EdgeExpand::expand_vertex_with_special_vertex_predicate( + graph, std::move(chunk), eep_, config_, params); + } }); } @@ -277,19 +293,23 @@ class EdgeExpandVWithGPVertexPredOpr : public IOperator { return "EdgeExpandVWithGPVertexPredOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - auto expr = pred_->bind(&graph, params); - GeneralPred expr_wrapper(std::move(expr)); - EdgeNbrPredicate vpred(expr_wrapper); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::expand_vertex>( - graph, std::move(chunk), eep_, vpred); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + auto expr = pred_->bind(&graph, params); + GeneralPred expr_wrapper(std::move(expr)); + EdgeNbrPredicate vpred(expr_wrapper); + { + return EdgeExpand::expand_vertex>( + graph, std::move(chunk), eep_, vpred); + } }); } @@ -302,15 +322,17 @@ class EdgeExpandDegreeOpr : public IOperator { public: EdgeExpandDegreeOpr(const EdgeExpandParams& eep) : eep_(eep) {} - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::expand_degree(graph, std::move(chunk), eep_); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + { return EdgeExpand::expand_degree(graph, std::move(chunk), eep_); } }); } @@ -507,15 +529,17 @@ class ExpandCountOpr : public IOperator { public: ExpandCountOpr(const EdgeExpandParams& eep) : eep_(eep) {} - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::expand_count(graph, std::move(chunk), eep_); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + { return EdgeExpand::expand_count(graph, std::move(chunk), eep_); } }); } diff --git a/src/execution/execute/ops/retrieve/gds_algo.cc b/src/execution/execute/ops/retrieve/gds_algo.cc index 73a224738..b36cafab9 100644 --- a/src/execution/execute/ops/retrieve/gds_algo.cc +++ b/src/execution/execute/ops/retrieve/gds_algo.cc @@ -28,24 +28,44 @@ GDSAlgoOpr::GDSAlgoOpr(std::unique_ptr algo_input, function::GDSAlgoFunction* algo_func) : algo_input_(std::move(algo_input)), algo_func_(algo_func) {} -neug::result GDSAlgoOpr::Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, neug::execution::OprTimer* timer) { - (void) ctx; - (void) timer; - if (algo_func_ == nullptr) { - THROW_RUNTIME_ERROR("GDSAlgoOpr: GDSAlgoFunction pointer is null"); - } - if (algo_func_->execFunc == nullptr) { - THROW_RUNTIME_ERROR( - "GDSAlgoOpr: algoExec not registered for GDS algorithm"); - } - if (algo_input_ == nullptr) { - THROW_RUNTIME_ERROR("GDSAlgoOpr: algo input is null"); - } - auto bound_input = algo_input_->bindParams(params); - const auto& input = bound_input ? *bound_input : *algo_input_; - return algo_func_->execFunc(input, graph_interface); +Stream GDSAlgoOpr::Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + // Legacy extension ABI: Context conversion is confined to this + // boundary. + + while (true) { + auto next_result = input.Next(); + if (!next_result) { + return error_stream(next_result.error()); + } + auto next = std::move(*next_result); + if (!next) { + break; + } + } + + (void) timer; + if (algo_func_ == nullptr) { + THROW_RUNTIME_ERROR("GDSAlgoOpr: GDSAlgoFunction pointer is null"); + } + if (algo_func_->execFunc == nullptr) { + THROW_RUNTIME_ERROR( + "GDSAlgoOpr: algoExec not registered for GDS algorithm"); + } + if (algo_input_ == nullptr) { + THROW_RUNTIME_ERROR("GDSAlgoOpr: algo input is null"); + } + auto bound_input = algo_input_->bindParams(params); + const auto& bound = bound_input ? *bound_input : *algo_input_; + auto output = algo_func_->execFunc(bound, graph_interface); + return stream_from_context(std::move(output)); + }); } neug::result GDSAlgoOprBuilder::Build( diff --git a/src/execution/execute/ops/retrieve/group_by.cc b/src/execution/execute/ops/retrieve/group_by.cc index e8f1f61ea..92e2377cb 100644 --- a/src/execution/execute/ops/retrieve/group_by.cc +++ b/src/execution/execute/ops/retrieve/group_by.cc @@ -38,20 +38,22 @@ class GroupByOpr : public IOperator { std::string get_operator_name() const override { return "GroupByOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - ctx.ensure_single_chunk("GroupByOpr"); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - auto key = create_key_func(mappings_, graph, chunk.chunk()); - std::vector reducers; - for (auto& aggr : aggrs_) { - reducers.push_back(create_reduce_op(aggr, graph, chunk.chunk())); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return reduce_stream( + std::move(input), + [this, &graph, params, + timer](ContextChunk&& chunk) -> result { + { + auto key = create_key_func(mappings_, graph, chunk.chunk()); + std::vector reducers; + for (auto& aggr : aggrs_) { + reducers.push_back(create_reduce_op(aggr, graph, chunk.chunk())); + } + return GroupBy::group_by(std::move(chunk), std::move(key), + std::move(reducers)); } - return GroupBy::group_by(std::move(chunk), std::move(key), - std::move(reducers)); }); } diff --git a/src/execution/execute/ops/retrieve/index_scan.cc b/src/execution/execute/ops/retrieve/index_scan.cc index d9f0da1a4..97c851102 100644 --- a/src/execution/execute/ops/retrieve/index_scan.cc +++ b/src/execution/execute/ops/retrieve/index_scan.cc @@ -28,26 +28,44 @@ class IndexScanOpr final : public IOperator { function::NeugCallFunction* function) : input{std::move(input)}, function{function} {} - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer*) override { - if (input == nullptr) { - THROW_RUNTIME_ERROR("IndexScanOpr: index scan input is null"); - } - if (function == nullptr || function->execFunc == nullptr) { - THROW_RUNTIME_ERROR( - "IndexScanOpr: index scan function is not executable"); - } - auto bound_input = input->bindParams(params); - if (bound_input == nullptr) { - THROW_RUNTIME_ERROR( - "IndexScanOpr: index scan input did not create a per-Eval instance"); - } - auto context_bound_input = bound_input->bindContext(std::move(ctx)); - if (context_bound_input == nullptr) { - THROW_RUNTIME_ERROR( - "IndexScanOpr: index scan input did not bind the input context"); - } - return function->execFunc(*context_bound_input, graph); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& upstream, + OprTimer* timer) override { + return defer_stream( + std::move(upstream), + [this, &graph, params, timer]( + Stream&& upstream) mutable -> Stream { + // Legacy extension ABI: Context conversion is confined to this + // boundary. + + auto ctx_result = materialize(std::move(upstream)); + if (!ctx_result) { + return error_stream(ctx_result.error()); + } + auto ctx = std::move(*ctx_result); + + if (input == nullptr) { + THROW_RUNTIME_ERROR("IndexScanOpr: index scan input is null"); + } + if (function == nullptr || function->execFunc == nullptr) { + THROW_RUNTIME_ERROR( + "IndexScanOpr: index scan function is not executable"); + } + auto bound_input = input->bindParams(params); + if (bound_input == nullptr) { + THROW_RUNTIME_ERROR( + "IndexScanOpr: index scan input did not create a per-Eval " + "instance"); + } + auto context_bound_input = bound_input->bindContext(std::move(ctx)); + if (context_bound_input == nullptr) { + THROW_RUNTIME_ERROR( + "IndexScanOpr: index scan input did not bind the input " + "context"); + } + auto output = function->execFunc(*context_bound_input, graph); + return stream_from_context(std::move(output)); + }); } std::string get_operator_name() const override { return "IndexScanOpr"; } diff --git a/src/execution/execute/ops/retrieve/intersect.cc b/src/execution/execute/ops/retrieve/intersect.cc index cc9b53b87..d5c09db59 100644 --- a/src/execution/execute/ops/retrieve/intersect.cc +++ b/src/execution/execute/ops/retrieve/intersect.cc @@ -41,33 +41,38 @@ class IntersectOprMultip : public IOperator { return "IntersectOprMultip"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - std::vector preds; - for (size_t i = 0; i < edge_preds_.size(); ++i) { - std::unique_ptr v_pred = - vertex_preds_[i] ? vertex_preds_[i]->bind(&graph, params) : nullptr; - std::unique_ptr e_pred = - edge_preds_[i] ? edge_preds_[i]->bind(&graph, params) : nullptr; - preds.emplace_back(std::move(v_pred), std::move(e_pred)); - } - if (eeps_.size() == 2) { - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Intersect::Binary_Intersect( - graph, params, std::move(chunk), std::move(preds[0]), - std::move(preds[1]), eeps_[0], eeps_[1], alias_); - }); - } + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + std::vector preds; + for (size_t i = 0; i < edge_preds_.size(); ++i) { + std::unique_ptr v_pred = + vertex_preds_[i] ? vertex_preds_[i]->bind(&graph, params) + : nullptr; + std::unique_ptr e_pred = + edge_preds_[i] ? edge_preds_[i]->bind(&graph, params) : nullptr; + preds.emplace_back(std::move(v_pred), std::move(e_pred)); + } + if (eeps_.size() == 2) { + { + return Intersect::Binary_Intersect( + graph, params, std::move(chunk), std::move(preds[0]), + std::move(preds[1]), eeps_[0], eeps_[1], alias_); + } + } - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Intersect::Multiple_Intersect(graph, params, std::move(chunk), - std::move(preds), eeps_, alias_); + { + return Intersect::Multiple_Intersect( + graph, params, std::move(chunk), std::move(preds), eeps_, + alias_); + } }); } @@ -93,30 +98,36 @@ class IntersectWithEdgeOpr : public IOperator { return "IntersectWithEdgeOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - std::vector preds; - for (size_t i = 0; i < edge_preds_.size(); ++i) { - std::unique_ptr v_pred = - vertex_preds_[i] ? vertex_preds_[i]->bind(&graph, params) : nullptr; - std::unique_ptr e_pred = - edge_preds_[i] ? edge_preds_[i]->bind(&graph, params) : nullptr; - preds.emplace_back(std::move(v_pred), std::move(e_pred)); - } - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - if (eeps_.size() == 2) { - return Intersect::Binary_Intersect_With_Edge( - graph, params, std::move(chunk), std::move(preds[0]), - std::move(preds[1]), eeps_[0], eeps_[1], v_alias_, edge_alias_); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + std::vector preds; + for (size_t i = 0; i < edge_preds_.size(); ++i) { + std::unique_ptr v_pred = + vertex_preds_[i] ? vertex_preds_[i]->bind(&graph, params) + : nullptr; + std::unique_ptr e_pred = + edge_preds_[i] ? edge_preds_[i]->bind(&graph, params) : nullptr; + preds.emplace_back(std::move(v_pred), std::move(e_pred)); + } + { + if (eeps_.size() == 2) { + return Intersect::Binary_Intersect_With_Edge( + graph, params, std::move(chunk), std::move(preds[0]), + std::move(preds[1]), eeps_[0], eeps_[1], v_alias_, + edge_alias_); + } + return Intersect::Multiple_Intersect_With_Edge( + graph, params, std::move(chunk), std::move(preds), eeps_, + v_alias_, edge_alias_); } - return Intersect::Multiple_Intersect_With_Edge( - graph, params, std::move(chunk), std::move(preds), eeps_, - v_alias_, edge_alias_); }); } diff --git a/src/execution/execute/ops/retrieve/join.cc b/src/execution/execute/ops/retrieve/join.cc index eebda019e..4100b750a 100644 --- a/src/execution/execute/ops/retrieve/join.cc +++ b/src/execution/execute/ops/retrieve/join.cc @@ -46,43 +46,36 @@ class JoinOpr : public IOperator { std::string get_operator_name() const override { return "JoinOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - neug::execution::Context ret_dup(ctx); - - std::unique_ptr left_timer = - (timer != nullptr) ? std::make_unique() - : nullptr; - auto left_ctx = - left_pipeline_.Execute(graph, std::move(ctx), params, left_timer.get()); - if (!left_ctx) { - return left_ctx; - } - std::unique_ptr right_timer = - (timer != nullptr) ? std::make_unique() - : nullptr; - auto right_ctx = right_pipeline_.Execute(graph, std::move(ret_dup), params, - right_timer.get()); - if (!right_ctx) { - return right_ctx; - } - if (NEUG_UNLIKELY(timer != nullptr)) { - timer->add_child(std::move(left_timer)); - timer->add_child(std::move(right_timer)); - } - left_ctx.value().ensure_single_chunk("JoinOpr::left"); - right_ctx.value().ensure_single_chunk("JoinOpr::right"); - auto join_result = - Join::join(std::move(left_ctx.value().chunk(0)), - std::move(right_ctx.value().chunk(0)), params_); - if (!join_result) { - return tl::make_unexpected(join_result.error()); - } - Context out; - out.append_chunk(std::move(join_result.value())); - return out; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto metadata = input.metadata(); + auto upstream = + std::make_shared>(std::move(input)); + return generate_chunk([this, &graph, params, timer, upstream, + metadata]() -> result { + GS_AUTO(seed, collect_batches(std::move(*upstream))); + auto left_timer = timer ? std::make_unique() : nullptr; + auto right_timer = timer ? std::make_unique() : nullptr; + auto left_stream = left_pipeline_.ExecuteStream( + graph, stream_from_batches(seed, metadata), params, + left_timer.get()); + GS_AUTO(left, collect_chunk(std::move(left_stream))); + auto right_stream = right_pipeline_.ExecuteStream( + graph, stream_from_batches(std::move(seed), metadata), params, + right_timer.get()); + GS_AUTO(right, collect_chunk(std::move(right_stream))); + if (timer) { + timer->add_child(std::move(left_timer)); + timer->add_child(std::move(right_timer)); + } + return Join::join(std::move(left), std::move(right), params_); + }); + }); } void build_explain_children(OprTimer* parent_timer, const ParamsMap& params, @@ -198,31 +191,27 @@ class PrimaryKeyJoinOpr : public IOperator { std::string get_operator_name() const override { return "PrimaryJoinOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - neug::execution::Context ret_dup(ctx); - std::unique_ptr right_timer = - (timer != nullptr) ? std::make_unique() - : nullptr; - auto right_ctx = right_pipeline_.Execute(graph, std::move(ret_dup), params, - right_timer.get()); - if (!right_ctx) { - return right_ctx; - } - if (NEUG_UNLIKELY(timer != nullptr)) { - timer->add_child(std::move(right_timer)); - } - right_ctx.value().ensure_single_chunk("PrimaryKeyJoinOpr"); - auto pk_result = Join::pk_join(graph, std::move(right_ctx.value().chunk(0)), - labels_, tag_, alias_); - if (!pk_result) { - return tl::make_unexpected(pk_result.error()); - } - Context out; - out.append_chunk(std::move(pk_result.value())); - return out; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + auto right_timer = timer ? std::make_unique() : nullptr; + auto* child = right_timer.get(); + if (timer) { + timer->add_child(std::move(right_timer)); + } + auto right = right_pipeline_.ExecuteStream(graph, std::move(input), + params, child); + return map_chunks( + std::move(right), + [this, &graph](ContextChunk&& chunk) -> result { + return Join::pk_join(graph, std::move(chunk), labels_, tag_, + alias_); + }); + }); } private: diff --git a/src/execution/execute/ops/retrieve/limit.cc b/src/execution/execute/ops/retrieve/limit.cc index 3c4f58355..b1c2b9b84 100644 --- a/src/execution/execute/ops/retrieve/limit.cc +++ b/src/execution/execute/ops/retrieve/limit.cc @@ -39,15 +39,36 @@ class LimitOpr : public IOperator { std::string get_operator_name() const override { return "LimitOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - ctx.ensure_single_chunk("LimitOpr"); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Limit::limit(std::move(chunk), lower_, upper_); - }); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + auto metadata = input.metadata(); + auto upstream = std::make_shared>(std::move(input)); + return Stream( + [upstream, skip = lower_, + remaining = upper_ > lower_ ? upper_ - lower_ : 0, + done = false]() mutable -> Stream::NextResult { + if (done) { + return std::optional{}; + } + GS_AUTO(next, upstream->Next()); + if (!next) { + return std::optional{}; + } + ContextChunk chunk = std::move(*next); + auto rows = chunk.row_num(); + auto begin = std::min(skip, rows); + skip -= begin; + auto count = std::min(remaining, rows - begin); + remaining -= count; + GS_AUTO(output, Limit::limit(std::move(chunk), begin, begin + count)); + if (remaining == 0) { + done = true; + *upstream = Stream(); + } + return std::optional(std::move(output)); + }, + std::move(metadata)); } private: diff --git a/src/execution/execute/ops/retrieve/order_by.cc b/src/execution/execute/ops/retrieve/order_by.cc index b472cd9e9..d533a345a 100644 --- a/src/execution/execute/ops/retrieve/order_by.cc +++ b/src/execution/execute/ops/retrieve/order_by.cc @@ -33,30 +33,33 @@ class OrderByOpr : public IOperator { std::string get_operator_name() const override { return "OrderByOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - ctx.ensure_single_chunk("OrderByOpr"); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - int keys_num = keys_.size(); - GeneralComparer cmp; - for (int i = 0; i < keys_num; ++i) { - cmp.add_keys(chunk.get(keys_[i].first), keys_[i].second); - } - sel_vec_t indices; - int32_t tag = keys_[0].first; - bool order = keys_[0].second; - if (chunk.get(tag)->order_by_limit(order, upper_, indices)) { - return OrderBy::staged_order_by_with_limit( - graph, std::move(chunk), cmp, lower_, upper_, indices); - } + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return reduce_stream( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + { + int keys_num = keys_.size(); + GeneralComparer cmp; + for (int i = 0; i < keys_num; ++i) { + cmp.add_keys(chunk.get(keys_[i].first), keys_[i].second); + } + sel_vec_t indices; + int32_t tag = keys_[0].first; + bool order = keys_[0].second; + if (chunk.get(tag)->order_by_limit(order, upper_, indices)) { + return OrderBy::staged_order_by_with_limit( + graph, std::move(chunk), cmp, lower_, upper_, indices); + } - return OrderBy::order_by_with_limit( - graph, std::move(chunk), cmp, lower_, upper_); + return OrderBy::order_by_with_limit( + graph, std::move(chunk), cmp, lower_, upper_); + } }); } diff --git a/src/execution/execute/ops/retrieve/path.cc b/src/execution/execute/ops/retrieve/path.cc index 1ba2fefdb..0c07c0710 100644 --- a/src/execution/execute/ops/retrieve/path.cc +++ b/src/execution/execute/ops/retrieve/path.cc @@ -229,22 +229,26 @@ class SPOrderByLimitOpr : public IOperator { std::string get_operator_name() const override { return "SPOrderByLimitOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - std::set expected_labels; - for (auto label : spp_.labels) { - expected_labels.insert(label.src_label); - expected_labels.insert(label.dst_label); - } - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return dispatch_vertex_predicate( - graph, expected_labels, config_, params, graph, std::move(chunk), - spp_, limit_); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + std::set expected_labels; + for (auto label : spp_.labels) { + expected_labels.insert(label.src_label); + expected_labels.insert(label.dst_label); + } + { + return dispatch_vertex_predicate( + graph, expected_labels, config_, params, graph, + std::move(chunk), spp_, limit_); + } }); } @@ -264,32 +268,35 @@ class SPOrderByLimitWithGPredOpr : public IOperator { return "SPOrderByLimitWithGPredOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - if (pred_) { - auto pred = pred_->bind(&graph, params); - - GeneralPred predicate_wrapper(std::move(pred)); - - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return PathExpand:: - single_source_shortest_path_with_order_by_length_limit( - graph, std::move(chunk), spp_, predicate_wrapper, limit_); - }); - } else { - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return PathExpand:: - single_source_shortest_path_with_order_by_length_limit( - graph, std::move(chunk), spp_, - [](label_t, vid_t) { return true; }, limit_); - }); - } + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + if (pred_) { + auto pred = pred_->bind(&graph, params); + + GeneralPred predicate_wrapper(std::move(pred)); + + { + return PathExpand:: + single_source_shortest_path_with_order_by_length_limit( + graph, std::move(chunk), spp_, predicate_wrapper, limit_); + } + } else { + { + return PathExpand:: + single_source_shortest_path_with_order_by_length_limit( + graph, std::move(chunk), spp_, + [](label_t, vid_t) { return true; }, limit_); + } + } + }); } private: @@ -366,17 +373,21 @@ class SPSPredOpr : public IOperator { std::string get_operator_name() const override { return "SPSPredOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return PathExpand:: - single_source_shortest_path_with_special_vertex_predicate( - graph, std::move(chunk), spp_, config_, params); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + { + return PathExpand:: + single_source_shortest_path_with_special_vertex_predicate( + graph, std::move(chunk), spp_, config_, params); + } }); } @@ -392,19 +403,23 @@ class SPGPredOpr : public IOperator { std::string get_operator_name() const override { return "SPGPredOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - auto pred = pred_->bind(&graph, params); - GeneralPred predicate_wrapper(std::move(pred)); - - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return PathExpand::single_source_shortest_path( - graph, std::move(chunk), spp_, predicate_wrapper); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + auto pred = pred_->bind(&graph, params); + GeneralPred predicate_wrapper(std::move(pred)); + + { + return PathExpand::single_source_shortest_path( + graph, std::move(chunk), spp_, predicate_wrapper); + } }); } @@ -418,17 +433,21 @@ class SPWithoutPredOpr : public IOperator { std::string get_operator_name() const override { return "SPWithoutPredOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return PathExpand::single_source_shortest_path( - graph, std::move(chunk), spp_, - [](label_t, vid_t) { return true; }); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + { + return PathExpand::single_source_shortest_path( + graph, std::move(chunk), spp_, + [](label_t, vid_t) { return true; }); + } }); } @@ -461,37 +480,41 @@ class ASPOpr : public IOperator { std::string get_operator_name() const override { return "ASPOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - Value oid; - if (expr_opr_.has_param()) { - auto name = expr_opr_.param().name(); - auto val = params.at(name).GetValue(); - oid = Value::INT64(val); - } else { - const auto& c = expr_opr_.const_(); - oid = Value::INT64(c.i64()); - } - vid_t vid; - if (!graph.GetVertexIndex(aspp_.labels[0].dst_label, oid, vid)) { - LOG(ERROR) << "vertex not found " - << static_cast(aspp_.labels[0].dst_label) << " " - << oid.to_string(); - RETURN_UNSUPPORTED_ERROR( - "vertex not found" + - std::to_string(static_cast(aspp_.labels[0].dst_label)) + " " + - std::string(oid.to_string())); - } - - auto v = std::make_pair(aspp_.labels[0].dst_label, vid); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return PathExpand::all_shortest_paths_with_given_source_and_dest( - graph, std::move(chunk), aspp_, v); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + Value oid; + if (expr_opr_.has_param()) { + auto name = expr_opr_.param().name(); + auto val = params.at(name).GetValue(); + oid = Value::INT64(val); + } else { + const auto& c = expr_opr_.const_(); + oid = Value::INT64(c.i64()); + } + vid_t vid; + if (!graph.GetVertexIndex(aspp_.labels[0].dst_label, oid, vid)) { + LOG(ERROR) << "vertex not found " + << static_cast(aspp_.labels[0].dst_label) << " " + << oid.to_string(); + RETURN_UNSUPPORTED_ERROR( + "vertex not found" + + std::to_string(static_cast(aspp_.labels[0].dst_label)) + + " " + std::string(oid.to_string())); + } + + auto v = std::make_pair(aspp_.labels[0].dst_label, vid); + { + return PathExpand::all_shortest_paths_with_given_source_and_dest( + graph, std::move(chunk), aspp_, v); + } }); } @@ -506,38 +529,42 @@ class SSSDSPOpr : public IOperator { : spp_(spp), expr_opr_(expr_opr) {} std::string get_operator_name() const override { return "SSSDSPOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - Value vertex = [&]() { - if (expr_opr_.has_param()) { - auto name = expr_opr_.param().name(); - auto val = params.at(name).GetValue(); - return Value::INT64(val); - } else { - const auto& c = expr_opr_.const_(); - return Value::INT64(c.i64()); - } - }(); - vid_t vid; - if (!graph.GetVertexIndex(spp_.labels[0].dst_label, vertex, vid)) { - LOG(ERROR) << "vertex not found" << spp_.labels[0].dst_label << " " - << vertex.to_string(); - RETURN_UNSUPPORTED_ERROR( - "vertex not found" + - std::to_string(static_cast(spp_.labels[0].dst_label)) + " " + - vertex.to_string()); - } - - auto v = std::make_pair(spp_.labels[0].dst_label, vid); - - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return PathExpand::single_source_single_dest_shortest_path( - graph, std::move(chunk), spp_, v); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + Value vertex = [&]() { + if (expr_opr_.has_param()) { + auto name = expr_opr_.param().name(); + auto val = params.at(name).GetValue(); + return Value::INT64(val); + } else { + const auto& c = expr_opr_.const_(); + return Value::INT64(c.i64()); + } + }(); + vid_t vid; + if (!graph.GetVertexIndex(spp_.labels[0].dst_label, vertex, vid)) { + LOG(ERROR) << "vertex not found" << spp_.labels[0].dst_label << " " + << vertex.to_string(); + RETURN_UNSUPPORTED_ERROR( + "vertex not found" + + std::to_string(static_cast(spp_.labels[0].dst_label)) + + " " + vertex.to_string()); + } + + auto v = std::make_pair(spp_.labels[0].dst_label, vid); + + { + return PathExpand::single_source_single_dest_shortest_path( + graph, std::move(chunk), spp_, v); + } }); } @@ -642,15 +669,17 @@ class PathExpandVOpr : public IOperator { public: explicit PathExpandVOpr(const PathExpandParams& pep) : pep_(pep) {} - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return PathExpand::edge_expand_v(graph, std::move(chunk), pep_); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + { return PathExpand::edge_expand_v(graph, std::move(chunk), pep_); } }); } std::string get_operator_name() const override { return "PathExpandVOpr"; } @@ -743,15 +772,17 @@ class PathExpandOpr : public IOperator { std::string get_operator_name() const override { return "PathExpandOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return PathExpand::edge_expand_p(graph, std::move(chunk), pep_); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + { return PathExpand::edge_expand_p(graph, std::move(chunk), pep_); } }); } @@ -768,18 +799,22 @@ class PathExpandOprWithPred : public IOperator { return "PathExpandOprWithPred"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - auto expr = pred_->bind(&graph, params); - GeneralPred predicate_wrapper(std::move(expr)); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return PathExpand::edge_expand_p_with_pred(graph, std::move(chunk), - pep_, predicate_wrapper); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + auto expr = pred_->bind(&graph, params); + GeneralPred predicate_wrapper(std::move(expr)); + { + return PathExpand::edge_expand_p_with_pred(graph, std::move(chunk), + pep_, predicate_wrapper); + } }); } @@ -798,23 +833,27 @@ class AnyWeightedShortestPathOpr : public IOperator { return "WeightedShortestPathOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - auto expr = weight_->bind(&graph, params); - auto weight_func = [&expr](const LabelTriplet& label, vid_t src, vid_t dst, - const void* data_ptr) { - return expr->Cast() - .eval_edge(label, src, dst, data_ptr) - .GetValue(); - }; - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return PathExpand::any_weighted_shortest_path(graph, std::move(chunk), - pep_, weight_func); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + auto expr = weight_->bind(&graph, params); + auto weight_func = [&expr](const LabelTriplet& label, vid_t src, + vid_t dst, const void* data_ptr) { + return expr->Cast() + .eval_edge(label, src, dst, data_ptr) + .GetValue(); + }; + { + return PathExpand::any_weighted_shortest_path( + graph, std::move(chunk), pep_, weight_func); + } }); } diff --git a/src/execution/execute/ops/retrieve/procedure_call.cc b/src/execution/execute/ops/retrieve/procedure_call.cc index b6c3b8007..de2a1da18 100644 --- a/src/execution/execute/ops/retrieve/procedure_call.cc +++ b/src/execution/execute/ops/retrieve/procedure_call.cc @@ -36,27 +36,49 @@ class ProcedureCallOpr : public IOperator { std::string get_operator_name() const override { return "ProcedureCallOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - (void) ctx; - (void) timer; - if (callFunction_ == nullptr) { - THROW_RUNTIME_ERROR("ProcedureCallOpr: callFunction is nullptr"); - } - if (unboundInput_ == nullptr) { - THROW_RUNTIME_ERROR("ProcedureCallOpr: unbound input is nullptr"); - } - if (callFunction_->execFunc == nullptr) { - THROW_RUNTIME_ERROR("ProcedureCallOpr: execFunc is nullptr"); - } - // bindParams returns a per-Eval bound input; nullptr means no deferred - // params and the unbound template is safe to exec as-is. - auto boundInput = unboundInput_->bindParams(params); - const auto& input = boundInput ? *boundInput : *unboundInput_; - return neug::result( - callFunction_->execFunc(input, graph)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + // Legacy extension ABI: Context conversion is confined to this + // boundary. + + while (true) { + auto next_result = input.Next(); + if (!next_result) { + return error_stream(next_result.error()); + } + auto next = std::move(*next_result); + if (!next) { + break; + } + } + + (void) timer; + if (callFunction_ == nullptr) { + THROW_RUNTIME_ERROR("ProcedureCallOpr: callFunction is nullptr"); + } + if (unboundInput_ == nullptr) { + THROW_RUNTIME_ERROR("ProcedureCallOpr: unbound input is nullptr"); + } + if (callFunction_->execFunc == nullptr) { + THROW_RUNTIME_ERROR("ProcedureCallOpr: execFunc is nullptr"); + } + // bindParams returns a per-Eval bound input; nullptr means no + // deferred params and the unbound template is safe to exec as-is. + auto boundInput = unboundInput_->bindParams(params); + const auto& bound = boundInput ? *boundInput : *unboundInput_; + auto output_result = neug::result( + callFunction_->execFunc(bound, graph)); + if (!output_result) { + return error_stream(output_result.error()); + } + auto output = std::move(*output_result); + return stream_from_context(std::move(output)); + }); } }; diff --git a/src/execution/execute/ops/retrieve/project.cc b/src/execution/execute/ops/retrieve/project.cc index b2e7551e3..0f7635d2a 100644 --- a/src/execution/execute/ops/retrieve/project.cc +++ b/src/execution/execute/ops/retrieve/project.cc @@ -47,38 +47,40 @@ class ProjectOpr : public IOperator { ~ProjectOpr() {} - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - if (is_select_columns_) { - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - ContextChunk ret; - for (auto& p : select_columns_mapping_) { - ret.set(p.second, chunk.get(p.first)); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph, params, + timer](ContextChunk&& chunk) -> result { + if (is_select_columns_) { + { + ContextChunk ret; + for (auto& p : select_columns_mapping_) { + ret.set(p.second, chunk.get(p.first)); + } + return ret; } - return ret; - }); - } - - std::vector exprs; - - for (size_t i = 0; i < expr_builders_.size(); ++i) { - if (!expr_builders_[i]) { - exprs.emplace_back(fallback_expr_builders_[i]->build(graph, params), - nullptr, fallback_expr_builders_[i]->alias()); - continue; - } else { - exprs.emplace_back(expr_builders_[i]->build(graph, params), - fallback_expr_builders_[i]->build(graph, params), - expr_builders_[i]->alias()); - } - } + } + + std::vector exprs; + + for (size_t i = 0; i < expr_builders_.size(); ++i) { + if (!expr_builders_[i]) { + exprs.emplace_back( + fallback_expr_builders_[i]->build(graph, params), nullptr, + fallback_expr_builders_[i]->alias()); + continue; + } else { + exprs.emplace_back( + expr_builders_[i]->build(graph, params), + fallback_expr_builders_[i]->build(graph, params), + expr_builders_[i]->alias()); + } + } - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Project::project(std::move(chunk), exprs, is_append_); + { return Project::project(std::move(chunk), exprs, is_append_); } }); } @@ -179,41 +181,44 @@ class ProjectOrderByOprBeta : public IOperator { return "ProjectOrderByOprBeta"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - const auto& graph = - dynamic_cast(graph_interface); - - auto cmp_func = [&](const DataChunk& chunk) -> GeneralComparer { - GeneralComparer cmp; - for (const auto& pair : order_by_pairs_) { - cmp.add_keys(chunk.get(pair.first), pair.second); - } - return cmp; - }; + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return reduce_stream( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + const auto& graph = + dynamic_cast(graph_interface); + + auto cmp_func = [&](const DataChunk& chunk) -> GeneralComparer { + GeneralComparer cmp; + for (const auto& pair : order_by_pairs_) { + cmp.add_keys(chunk.get(pair.first), pair.second); + } + return cmp; + }; - std::vector exprs; + std::vector exprs; - for (size_t i = 0; i < expr_builders_.size(); ++i) { - if (!expr_builders_[i]) { - exprs.emplace_back( - ProjectOp(fallback_expr_builders_[i]->build(graph, params), nullptr, - fallback_expr_builders_[i]->alias())); - continue; - } - exprs.emplace_back( - ProjectOp(expr_builders_[i]->build(graph, params), - fallback_expr_builders_[i]->build(graph, params), - expr_builders_[i]->alias())); - } - ctx.ensure_single_chunk("ProjectOrderByOprBeta"); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Project::project_order_by_fuse( - graph, params, std::move(chunk), std::move(exprs), cmp_func, - lower_bound_, upper_bound_, order_by_keys_, first_pair_); + for (size_t i = 0; i < expr_builders_.size(); ++i) { + if (!expr_builders_[i]) { + exprs.emplace_back( + ProjectOp(fallback_expr_builders_[i]->build(graph, params), + nullptr, fallback_expr_builders_[i]->alias())); + continue; + } + exprs.emplace_back( + ProjectOp(expr_builders_[i]->build(graph, params), + fallback_expr_builders_[i]->build(graph, params), + expr_builders_[i]->alias())); + } + { + return Project::project_order_by_fuse( + graph, params, std::move(chunk), std::move(exprs), cmp_func, + lower_bound_, upper_bound_, order_by_keys_, first_pair_); + } }); } diff --git a/src/execution/execute/ops/retrieve/scan.cc b/src/execution/execute/ops/retrieve/scan.cc index 7886a519f..eee87cf9a 100644 --- a/src/execution/execute/ops/retrieve/scan.cc +++ b/src/execution/execute/ops/retrieve/scan.cc @@ -67,52 +67,57 @@ class FilterOidsGPredOpr : public IOperator { std::unique_ptr&& pred) : params_(params), oids_(oids), pred_(std::move(pred)) {} - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - ctx = Context(); - ctx.append_chunk(DataChunk()); - const auto& rhs_op = oids_.expression().operators(0); - if ((rhs_op.has_const_() && rhs_op.const_().has_none()) || - (rhs_op.has_param() && params.at(rhs_op.param().name()).IsNull())) { - static const std::vector no_oids; - auto empty_chunk = Scan::filter_oids(std::move(ctx.chunk(0)), graph, - params_, DummyPred(), no_oids); - if (!empty_chunk) { - return tl::make_unexpected(empty_chunk.error()); - } - ctx.chunk(0) = std::move(*empty_chunk); - return ctx; - } - std::vector oid_values = ScanUtils::parse_ids(oids_, params); - if (oids_.cmp() == common::Logical::WITHIN) { - oid_values = deduplicate_ids(std::move(oid_values)); - } + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + return generate_chunk([this, &graph, params, + timer]() -> result { + ContextChunk chunk; - if (pred_ == nullptr) { - if (params_.tables.size() == 1 && oid_values.size() == 1) { - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Scan::find_vertex_with_oid(std::move(chunk), graph, - params_.tables[0], - oid_values[0], params_.alias); - }); - } - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Scan::filter_oids(std::move(chunk), graph, params_, - DummyPred(), oid_values); - }); - } else { - auto pred = pred_->bind(&graph, params); - GeneralPred predicate_wrapper(std::move(pred)); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Scan::filter_oids(std::move(chunk), graph, params_, - predicate_wrapper, oid_values); + const auto& rhs_op = oids_.expression().operators(0); + if ((rhs_op.has_const_() && rhs_op.const_().has_none()) || + (rhs_op.has_param() && + params.at(rhs_op.param().name()).IsNull())) { + static const std::vector no_oids; + auto empty_chunk = Scan::filter_oids( + std::move(chunk), graph, params_, DummyPred(), no_oids); + if (!empty_chunk) { + return tl::make_unexpected(empty_chunk.error()); + } + chunk = std::move(*empty_chunk); + return std::move(chunk); + } + std::vector oid_values = ScanUtils::parse_ids(oids_, params); + if (oids_.cmp() == common::Logical::WITHIN) { + oid_values = deduplicate_ids(std::move(oid_values)); + } + + if (pred_ == nullptr) { + if (params_.tables.size() == 1 && oid_values.size() == 1) { + { + return Scan::find_vertex_with_oid( + std::move(chunk), graph, params_.tables[0], oid_values[0], + params_.alias); + } + } + { + return Scan::filter_oids(std::move(chunk), graph, params_, + DummyPred(), oid_values); + } + } else { + auto pred = pred_->bind(&graph, params); + GeneralPred predicate_wrapper(std::move(pred)); + { + return Scan::filter_oids(std::move(chunk), graph, params_, + predicate_wrapper, oid_values); + } + } }); - } + }); } std::string get_operator_name() const override { @@ -133,17 +138,22 @@ class ScanWithSPredOpr : public IOperator { std::string get_operator_name() const override { return "ScanWithSPredOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - ctx = Context(); - ctx.append_chunk(DataChunk()); - - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Scan::scan_vertex_with_special_vertex_predicate( - std::move(chunk), graph, scan_params_, config_, params); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + return generate_chunk( + [this, &graph, params, timer]() -> result { + ContextChunk chunk; + + { + return Scan::scan_vertex_with_special_vertex_predicate( + std::move(chunk), graph, scan_params_, config_, params); + } + }); }); } @@ -157,27 +167,32 @@ class ScanWithGPredOpr : public IOperator { ScanWithGPredOpr(const ScanParams& scan_params, std::unique_ptr pred) : scan_params_(scan_params), pred_(std::move(pred)) {} - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - ctx = Context(); - ctx.append_chunk(DataChunk()); - if (pred_ == nullptr) { - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Scan::scan_vertex(std::move(chunk), graph, scan_params_, - DummyPred()); - }); - } else { - auto pred = pred_->bind(&graph, params); - GeneralPred pred_wrapper(std::move(pred)); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Scan::scan_vertex(std::move(chunk), graph, scan_params_, - pred_wrapper); - }); - } + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph, params, + timer](Stream&& input) mutable -> Stream { + return generate_chunk( + [this, &graph, params, timer]() -> result { + ContextChunk chunk; + + if (pred_ == nullptr) { + { + return Scan::scan_vertex(std::move(chunk), graph, + scan_params_, DummyPred()); + } + } else { + auto pred = pred_->bind(&graph, params); + GeneralPred pred_wrapper(std::move(pred)); + { + return Scan::scan_vertex(std::move(chunk), graph, + scan_params_, pred_wrapper); + } + } + }); + }); } std::string get_operator_name() const override { return "ScanWithGPredOpr"; } @@ -256,17 +271,24 @@ class DummySourceOpr : public IOperator { public: DummySourceOpr() {} - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - Context out; - ContextChunk chunk; - ValueColumnBuilder builder; - builder.push_back_opt(0); - chunk.set(-1, builder.finish()); - out.append_chunk(std::move(chunk)); - return out; + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return defer_stream( + std::move(input), + [this, &graph_interface, params, + timer](Stream&& input) mutable -> Stream { + return generate_chunk([this, &graph_interface, params, + timer]() -> result { + ContextChunk chunk; + ValueColumnBuilder builder; + builder.push_back_opt(0); + chunk.set(-1, builder.finish()); + + return std::move(chunk); + }); + }); } std::string get_operator_name() const override { return "DummySourceOpr"; } diff --git a/src/execution/execute/ops/retrieve/select.cc b/src/execution/execute/ops/retrieve/select.cc index 200d9e5a8..5f42e364f 100644 --- a/src/execution/execute/ops/retrieve/select.cc +++ b/src/execution/execute/ops/retrieve/select.cc @@ -41,50 +41,56 @@ class SelectIdNeOpr : public IOperator { std::string get_operator_name() const override { return "SelectIdNeOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - auto expr = pred_->bind(&graph_interface, params); - neug::execution::GeneralPred fallback_pred(std::move(expr)); - const auto& name = prop_name_; - int64_t oid = params.count(param_name_) - ? params.at(param_name_).GetValue() - : 0; - - return ctx.apply_chunks([&](ContextChunk&& chunk) - -> neug::result { - auto col = chunk.get(tag_); - if ((!col->is_optional()) && - col->column_type() == ContextColumnType::kVertex) { - auto vertex_col = std::dynamic_pointer_cast(col); - auto labels = vertex_col->get_labels_set(); - if (labels.size() == 1 && - name == graph_interface.schema().get_vertex_primary_key_name( - *labels.begin())) { - auto label = *labels.begin(); - vid_t vid; - if (graph_interface.GetVertexIndex(label, Value::INT64(oid), vid)) { - if (vertex_col->vertex_column_type() == VertexColumnType::kSingle) { - const SLVertexColumn& sl_vertex_col = - *(dynamic_cast(vertex_col.get())); - return Select::select( - std::move(chunk), - [&sl_vertex_col, vid](const DataChunk&, size_t i) { - return sl_vertex_col.get_vertex(i).vid_ != vid; - }); - } else { - return Select::select( - std::move(chunk), - [&vertex_col, vid](const DataChunk&, size_t i) { - return vertex_col->get_vertex(i).vid_ != vid; - }); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + auto expr = pred_->bind(&graph_interface, params); + neug::execution::GeneralPred fallback_pred(std::move(expr)); + const auto& name = prop_name_; + int64_t oid = params.count(param_name_) + ? params.at(param_name_).GetValue() + : 0; + + { + auto col = chunk.get(tag_); + if ((!col->is_optional()) && + col->column_type() == ContextColumnType::kVertex) { + auto vertex_col = std::dynamic_pointer_cast(col); + auto labels = vertex_col->get_labels_set(); + if (labels.size() == 1 && + name == graph_interface.schema().get_vertex_primary_key_name( + *labels.begin())) { + auto label = *labels.begin(); + vid_t vid; + if (graph_interface.GetVertexIndex(label, Value::INT64(oid), + vid)) { + if (vertex_col->vertex_column_type() == + VertexColumnType::kSingle) { + const SLVertexColumn& sl_vertex_col = *( + dynamic_cast(vertex_col.get())); + return Select::select( + std::move(chunk), + [&sl_vertex_col, vid](const DataChunk&, size_t i) { + return sl_vertex_col.get_vertex(i).vid_ != vid; + }); + } else { + return Select::select( + std::move(chunk), + [&vertex_col, vid](const DataChunk&, size_t i) { + return vertex_col->get_vertex(i).vid_ != vid; + }); + } + } + } } + return Select::select(std::move(chunk), fallback_pred); } - } - } - return Select::select(std::move(chunk), fallback_pred); - }); + }); } private: @@ -101,15 +107,16 @@ class SelectOpr : public IOperator { std::string get_operator_name() const override { return "SelectOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - auto expr = pred_->bind(&graph, params); - neug::execution::GeneralPred expr_wrapper(std::move(expr)); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Select::select(std::move(chunk), expr_wrapper); + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph, params, + timer](ContextChunk&& chunk) -> result { + auto expr = pred_->bind(&graph, params); + neug::execution::GeneralPred expr_wrapper(std::move(expr)); + { return Select::select(std::move(chunk), expr_wrapper); } }); } diff --git a/src/execution/execute/ops/retrieve/sink.cc b/src/execution/execute/ops/retrieve/sink.cc index 5f5ef371a..6880dcaad 100644 --- a/src/execution/execute/ops/retrieve/sink.cc +++ b/src/execution/execute/ops/retrieve/sink.cc @@ -30,10 +30,11 @@ class SinkOpr : public IOperator { public: explicit SinkOpr(const std::vector& tag_ids) : tag_ids_(tag_ids) {} - neug::result Eval(IStorageInterface& graph, const ParamsMap& params, - Context&& ctx, OprTimer* timer) override { - ctx.tag_ids = tag_ids_; - return ctx; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + OprTimer* timer) override { + input.set_metadata(StreamMetadata{tag_ids_}); + return std::move(input); } std::string get_operator_name() const override { return "SinkOpr"; } diff --git a/src/execution/execute/ops/retrieve/tc_fuse.cc b/src/execution/execute/ops/retrieve/tc_fuse.cc index fd1cb6be6..f7f850316 100644 --- a/src/execution/execute/ops/retrieve/tc_fuse.cc +++ b/src/execution/execute/ops/retrieve/tc_fuse.cc @@ -73,17 +73,21 @@ class TCOpr : public IOperator { std::string get_operator_name() const override { return "TCOpr"; } - neug::result Eval( - IStorageInterface& graph_interface, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - auto& graph = dynamic_cast(graph_interface); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return EdgeExpand::tc(graph, std::move(chunk), labels_, - input_tag_, alias1_, alias2_, is_lt_, - params.at(param_name_)); - }); + Stream Eval(IStorageInterface& graph_interface, + const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks(std::move(input), + [this, &graph_interface, params, + timer](ContextChunk&& chunk) -> result { + auto& graph = dynamic_cast( + graph_interface); + { + return EdgeExpand::tc( + graph, std::move(chunk), labels_, input_tag_, + alias1_, alias2_, is_lt_, params.at(param_name_)); + } + }); } private: diff --git a/src/execution/execute/ops/retrieve/unfold.cc b/src/execution/execute/ops/retrieve/unfold.cc index d4ae5efc7..a1c4ddc36 100644 --- a/src/execution/execute/ops/retrieve/unfold.cc +++ b/src/execution/execute/ops/retrieve/unfold.cc @@ -52,24 +52,22 @@ class UnfoldOpr : public IOperator { std::string get_operator_name() const override { return "UnfoldOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - if (key_.has_value()) { - auto key_val = key_.value(); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Unfold::unfold(std::move(chunk), key_val, alias_); - }); - } else { - auto expr = expr_->bind(&graph, params); - auto& record_expr = expr->Cast(); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return Unfold::unfold(std::move(chunk), record_expr, alias_); - }); - } + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks( + std::move(input), + [this, &graph, params, + timer](ContextChunk&& chunk) -> result { + if (key_.has_value()) { + auto key_val = key_.value(); + { return Unfold::unfold(std::move(chunk), key_val, alias_); } + } else { + auto expr = expr_->bind(&graph, params); + auto& record_expr = expr->Cast(); + { return Unfold::unfold(std::move(chunk), record_expr, alias_); } + } + }); } private: diff --git a/src/execution/execute/ops/retrieve/union.cc b/src/execution/execute/ops/retrieve/union.cc index 6a155b106..6dfd6b392 100644 --- a/src/execution/execute/ops/retrieve/union.cc +++ b/src/execution/execute/ops/retrieve/union.cc @@ -36,33 +36,46 @@ class UnionOpr : public IOperator { std::string get_operator_name() const override { return "UnionOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - std::vector chunks; - for (auto& plan : sub_plans_) { - neug::execution::Context n_ctx = ctx; - std::unique_ptr sub_timer = - (timer != nullptr) ? std::make_unique() - : nullptr; - auto ret = plan.Execute(graph, std::move(n_ctx), params, sub_timer.get()); - if (NEUG_UNLIKELY(timer != nullptr)) { - timer->add_child(std::move(sub_timer)); - } - if (!ret) { - return ret; - } - ret.value().ensure_single_chunk("UnionOpr::sub_plan"); - chunks.emplace_back(std::move(ret.value().chunk(0))); - } - auto union_result = Union::union_op(std::move(chunks)); - if (!union_result) { - return tl::make_unexpected(union_result.error()); - } - Context out; - out.append_chunk(std::move(union_result.value())); - return out; + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + struct State { + Stream input; + std::optional> seed; + Stream branch; + size_t index = 0; + }; + auto metadata = input.metadata(); + auto state = std::make_shared(); + state->input = std::move(input); + return Stream( + [this, &graph, params, timer, metadata, + state]() mutable -> Stream::NextResult { + if (!state->seed) { + GS_AUTO(seed, collect_batches(std::move(state->input))); + state->seed = std::move(seed); + } + while (true) { + GS_AUTO(next, state->branch.Next()); + if (next) { + // UNION has no anonymous output head, matching the union kernel. + next->head().reset(); + return next; + } + if (state->index == sub_plans_.size()) { + return std::optional{}; + } + auto sub_timer = timer ? std::make_unique() : nullptr; + auto* child = sub_timer.get(); + if (timer) { + timer->add_child(std::move(sub_timer)); + } + auto branch = sub_plans_[state->index++].ExecuteStream( + graph, stream_from_batches(*state->seed, metadata), params, + child); + state->branch = std::move(branch); + } + }); } void build_explain_children(OprTimer* parent_timer, const ParamsMap& params, diff --git a/src/execution/execute/ops/retrieve/vertex.cc b/src/execution/execute/ops/retrieve/vertex.cc index a763fa7b1..13063c92d 100644 --- a/src/execution/execute/ops/retrieve/vertex.cc +++ b/src/execution/execute/ops/retrieve/vertex.cc @@ -39,25 +39,27 @@ class GetVFromEdgesOpr : public IOperator { std::string get_operator_name() const override { return "GetVFromEdgesOpr"; } - neug::result Eval( - IStorageInterface& graph, const ParamsMap& params, - neug::execution::Context&& ctx, - neug::execution::OprTimer* timer) override { - if (pred_ != nullptr) { - auto expr = pred_->bind(&graph, params); - GeneralPred pred(std::move(expr)); - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return GetV::get_vertex_from_edges(graph, std::move(chunk), - v_params_, pred); - }); - } else { - return ctx.apply_chunks( - [&](ContextChunk&& chunk) -> neug::result { - return GetV::get_vertex_from_edges(graph, std::move(chunk), - v_params_, DummyPred()); - }); - } + Stream Eval(IStorageInterface& graph, const ParamsMap& params, + Stream&& input, + neug::execution::OprTimer* timer) override { + return map_chunks(std::move(input), + [this, &graph, params, + timer](ContextChunk&& chunk) -> result { + if (pred_ != nullptr) { + auto expr = pred_->bind(&graph, params); + GeneralPred pred(std::move(expr)); + { + return GetV::get_vertex_from_edges( + graph, std::move(chunk), v_params_, pred); + } + } else { + { + return GetV::get_vertex_from_edges( + graph, std::move(chunk), v_params_, + DummyPred()); + } + } + }); } private: diff --git a/src/execution/execute/pipeline.cc b/src/execution/execute/pipeline.cc index 0fbf31c33..217e5a193 100644 --- a/src/execution/execute/pipeline.cc +++ b/src/execution/execute/pipeline.cc @@ -16,6 +16,7 @@ #include "neug/execution/execute/pipeline.h" #include +#include #include #include #include @@ -28,49 +29,85 @@ namespace neug { namespace execution { class OprTimer; +namespace { +Status operator_error(const Status& error, const std::string& name) { + return Status(error.error_code(), "Execution failed at operator: [" + name + + "], " + error.error_message()); +} + +// Pulling upstream happens inside downstream Eval/Next. Charge that time only +// to its producer, rather than counting it twice in PROFILE. +class StreamTimerScope { + public: + StreamTimerScope(OprTimer& timer, double& charged) + : timer_(timer), charged_(charged), before_(charged) { + clock_.start(); + } + ~StreamTimerScope() { + double elapsed = std::max(0.0, clock_.elapsed() - (charged_ - before_)); + timer_.add_elapsed(elapsed); + charged_ += elapsed; + } + + private: + OprTimer& timer_; + double& charged_; + double before_; + TimerUnit clock_; +}; +} // namespace + neug::result Pipeline::Execute(IStorageInterface& graph, Context&& ctx, const ParamsMap& params, OprTimer* timer) { - neug::Status status = Status::OK(); - TimerUnit tu; - OprTimer* cur_timer = timer; - std::unique_ptr next_timer = nullptr; + auto stream = + ExecuteStream(graph, stream_from_context(std::move(ctx)), params, timer); + return materialize(std::move(stream)); +} + +Stream Pipeline::ExecuteStream(IStorageInterface& graph, + Stream stream, + const ParamsMap& params, + OprTimer* timer) { + auto charged = timer ? std::make_shared(0.0) : nullptr; + auto* current_timer = timer; for (size_t i = 0; i < operators_.size(); ++i) { - if (NEUG_UNLIKELY(timer != nullptr)) { - tu.start(); + const auto name = operators_[i]->get_operator_name(); + if (current_timer) { + current_timer->set_name(name); } - TRY_HANDLE_ALL_WITH_EXCEPTION( - neug::result, - [&]() -> neug::result { - auto ret = - operators_[i]->Eval(graph, params, std::move(ctx), cur_timer); - if (!ret) { - return ret; - } - if (NEUG_UNLIKELY(timer != nullptr)) { - cur_timer->set_name(operators_[i]->get_operator_name()); - cur_timer->add_num_tuples(ret.value().row_num()); - cur_timer->record(tu); - if (i + 1 < operators_.size()) { - next_timer = std::make_unique(); - cur_timer->set_next(std::move(next_timer)); - cur_timer = cur_timer->next(); + auto output = + operators_[i]->Eval(graph, params, std::move(stream), current_timer); + auto metadata = output.metadata(); + auto producer = std::make_shared>(std::move(output)); + auto pull = [producer, name]() -> Stream::NextResult { + auto next = producer->Next(); + if (!next) { + return tl::unexpected(operator_error(next.error(), name)); + } + return next; + }; + if (current_timer) { + stream = Stream( + [pull = std::move(pull), current_timer, + charged]() -> Stream::NextResult { + StreamTimerScope scope(*current_timer, *charged); + auto next = pull(); + if (next && *next) { + current_timer->add_num_tuples((**next).row_num()); } - } - return ret; - }, - [&](const neug::Status& err) { - status = neug::Status(err.error_code(), - "Execution failed at operator: [" + - operators_[i]->get_operator_name() + "], " + - err.error_message()); - }, - [&ctx](neug::result&& res) { ctx = std::move(res.value()); }); - if (!status.ok()) { - RETURN_ERROR(status); + return next; + }, + std::move(metadata)); + } else { + stream = Stream(std::move(pull), std::move(metadata)); + } + if (current_timer && i + 1 < operators_.size()) { + current_timer->set_next(std::make_unique()); + current_timer = current_timer->next(); } } - return ctx; + return std::move(stream); } neug::result> Pipeline::explain_tree( diff --git a/src/storages/loader/loader_utils.cc b/src/storages/loader/loader_utils.cc index 703d1b532..0a97d408c 100644 --- a/src/storages/loader/loader_utils.cc +++ b/src/storages/loader/loader_utils.cc @@ -884,7 +884,9 @@ struct CsvSupplierRuntime { config.double_quote, config.delimiter, config.use_threads, stream_factory_) .count(); - reset_reader(); + if (row_num_ > rows_to_skip_) { + reset_reader(); + } } std::shared_ptr get_next_chunk() { diff --git a/src/utils/io/read/common/row_expression_filter.cc b/src/utils/io/read/common/row_expression_filter.cc index 1d35130eb..cfca53af5 100644 --- a/src/utils/io/read/common/row_expression_filter.cc +++ b/src/utils/io/read/common/row_expression_filter.cc @@ -24,7 +24,6 @@ #include "neug/common/types/value.h" #include "neug/execution/expression/expr.h" #include "neug/generated/proto/plan/expr.pb.h" -#include "neug/storages/loader/loader_utils.h" #include "neug/utils/exception/exception.h" #include "neug/utils/io/read/common/type_converter.h" @@ -156,24 +155,6 @@ bool RowExpressionFilter::eval(size_t row) const { return !evaluator_ || evaluator_(row); } -DataChunk read_all_chunks( - const std::vector>& suppliers) { - std::vector> chunks; - for (const auto& supplier : suppliers) { - if (!supplier) { - THROW_INVALID_ARGUMENT_EXCEPTION("Data chunk supplier is null"); - } - while (true) { - auto chunk = supplier->GetNextChunk(); - if (!chunk) { - break; - } - chunks.push_back(std::move(chunk)); - } - } - return merge_chunks(std::move(chunks)); -} - DataChunk merge_chunks(std::vector> chunks) { size_t total_rows = 0; size_t column_count = 0; diff --git a/src/utils/io/read/csv/csv_reader.cc b/src/utils/io/read/csv/csv_reader.cc index a90d83972..524ba9966 100644 --- a/src/utils/io/read/csv/csv_reader.cc +++ b/src/utils/io/read/csv/csv_reader.cc @@ -30,7 +30,6 @@ #include #include -#include "neug/execution/common/context.h" #include "neug/generated/proto/plan/expr.pb.h" #include "neug/storages/loader/loader_utils.h" #include "neug/utils/exception/exception.h" @@ -42,16 +41,6 @@ namespace neug { namespace reader { -namespace { - -CsvReadConfig read_config_for_supplier(const CsvReadConfig& config) { - CsvReadConfig read_config = config; - read_config.include_columns = config.column_names; - return read_config; -} - -} // namespace - CsvReader::CsvReader(std::shared_ptr sharedState, std::unique_ptr optionsBuilder) : sharedState_(std::move(sharedState)), @@ -59,103 +48,63 @@ CsvReader::CsvReader(std::shared_ptr sharedState, CsvReader::~CsvReader() = default; -void CsvReader::read(std::shared_ptr /*localState*/, - execution::Context& ctx) { - if (!sharedState_) { - THROW_INVALID_ARGUMENT_EXCEPTION("SharedState is null"); +std::shared_ptr CsvReader::getDataChunkSupplier() { + if (!sharedState_ || !optionsBuilder_) { + THROW_INVALID_ARGUMENT_EXCEPTION("CSV reader is not initialized"); } - if (!optionsBuilder_) { - THROW_INVALID_ARGUMENT_EXCEPTION("Options builder is null"); - } - auto config = optionsBuilder_->build(); - if (!optionsBuilder_->projectColumns(config)) { - LOG(WARNING) << "Failed to set column projection, using all columns"; - } - - const auto& fileSchema = sharedState_->schema.file; - ReadOptions readOpts; - const bool use_batch_read = readOpts.batch_read.get(fileSchema.options); - - auto read_config = read_config_for_supplier(config); - if (sharedState_->skipRows) { - // Need all columns to evaluate row-filter expression; - // full_read will project afterwards. - read_config.include_columns = config.column_names; - } else if (!sharedState_->projectColumns.empty()) { - if (use_batch_read) { - // batch_read streams chunks directly to the consumer without - // post-projection, so push column projection down to the supplier. - read_config.include_columns = config.include_columns; - } else { - // full_read handles projection via project_chunk(). - read_config.include_columns = config.column_names; - } - } - - const auto& paths = fileSchema.paths; - if (paths.empty()) { + optionsBuilder_->projectColumns(config); + if (sharedState_->schema.file.paths.empty()) { THROW_INVALID_ARGUMENT_EXCEPTION("No file paths provided"); } - std::vector> suppliers; - suppliers.reserve(paths.size()); - for (const auto& path : paths) { - suppliers.push_back(std::make_shared( - path, read_config, - io::bindInputStream(sharedState_->stream_opener, path))); - } - - if (use_batch_read && !sharedState_->skipRows) { - batch_read(suppliers, ctx); - } else { - full_read(suppliers, ctx, config); - } -} - -void CsvReader::full_read( - const std::vector>& suppliers, - execution::Context& output, const CsvReadConfig& output_config) { - auto merged = read_all_chunks(suppliers); - - if (merged.col_num() == 0) { - // No data rows at all (e.g. a header-only file). Emit an empty result; - // there is nothing to validate, filter or project. - output.clear(); - return; - } - - int expected_cols = sharedState_->columnNum(); - if (expected_cols > 0 && - static_cast(merged.col_num()) != expected_cols && - sharedState_->projectColumns.empty()) { - THROW_IO_EXCEPTION( - "Column number mismatch between schema and CSV data, schema: " + - std::to_string(expected_cols) + - ", data: " + std::to_string(merged.col_num())); - } - - auto filtered = - filter_chunk(merged, sharedState_->skipRows, output_config.column_names, - sharedState_->parameters); - auto projected = project_chunk(filtered, output_config.column_names, - sharedState_->projectColumns.empty() - ? output_config.include_columns - : sharedState_->projectColumns); - - output.clear(); - output.append_chunk(std::move(projected)); -} - -void CsvReader::batch_read( - const std::vector>& suppliers, - execution::Context& output) { - output.clear(); - for (const auto& supplier : suppliers) { - while (auto chunk = supplier->GetNextChunk()) { - output.append_chunk(std::move(*chunk)); + // Each execution owns this supplier. Files are initialized one at a time, + // only when requested, and filtering/projection happen before yielding. + class CsvStreamSupplier final : public IDataChunkSupplier { + public: + CsvStreamSupplier(std::shared_ptr state, + CsvReadConfig config) + : state_(std::move(state)), config_(std::move(config)) {} + + int64_t RowNum() const override { return -1; } + + std::shared_ptr GetNextChunk() override { + while (true) { + if (!current_) { + if (file_ == state_->schema.file.paths.size()) { + return nullptr; + } + const auto& path = state_->schema.file.paths[file_++]; + auto read_config = config_; + read_config.include_columns = + state_->skipRows ? config_.column_names : config_.include_columns; + current_ = std::make_shared( + path, read_config, + io::bindInputStream(state_->stream_opener, path)); + } + auto chunk = current_->GetNextChunk(); + if (!chunk) { + current_.reset(); + continue; + } + if (!state_->skipRows) { + return chunk; + } + auto filtered = filter_chunk(*chunk, state_->skipRows, + config_.column_names, state_->parameters); + auto projected = project_chunk(filtered, config_.column_names, + config_.include_columns); + return std::make_shared(std::move(projected)); + } } - } + + private: + std::shared_ptr state_; + CsvReadConfig config_; + size_t file_ = 0; + std::shared_ptr current_; + }; + return std::make_shared(sharedState_, std::move(config)); } result> CsvReader::inferSchema() { diff --git a/src/utils/io/read/json/json_reader.cc b/src/utils/io/read/json/json_reader.cc index 7008f9ea2..85bfa29d6 100644 --- a/src/utils/io/read/json/json_reader.cc +++ b/src/utils/io/read/json/json_reader.cc @@ -31,7 +31,6 @@ #include "neug/common/columns/columns_utils.h" #include "neug/common/types/value.h" -#include "neug/execution/common/context.h" #include "neug/storages/loader/loader_utils.h" #include "neug/utils/exception/exception.h" #include "neug/utils/io/read/common/options.h" @@ -412,12 +411,6 @@ class JsonChunkSupplier : public IDataChunkSupplier { std::unique_ptr line_reader_; // remote JSONL mode }; -JsonReadConfig read_config_for_supplier(const JsonReadConfig& config) { - JsonReadConfig read_config = config; - read_config.include_columns = config.column_names; - return read_config; -} - } // namespace JsonReader::JsonReader(std::shared_ptr sharedState, @@ -427,99 +420,63 @@ JsonReader::JsonReader(std::shared_ptr sharedState, JsonReader::~JsonReader() = default; -void JsonReader::read(std::shared_ptr /*localState*/, - execution::Context& ctx) { +std::shared_ptr JsonReader::getDataChunkSupplier() { if (!sharedState_ || !optionsBuilder_) { - THROW_INVALID_ARGUMENT_EXCEPTION("JsonReader state or builder is null"); + THROW_INVALID_ARGUMENT_EXCEPTION("JSON reader is not initialized"); } - auto config = optionsBuilder_->build(); - if (!optionsBuilder_->projectColumns(config)) { - LOG(WARNING) << "Failed to set column projection, using all columns"; - } - - const auto& fileSchema = sharedState_->schema.file; - ReadOptions readOpts; - const bool use_batch_read = readOpts.batch_read.get(fileSchema.options); - - auto read_config = read_config_for_supplier(config); - if (sharedState_->skipRows) { - // Need all columns to evaluate row-filter expression; - // full_read will project afterwards. - read_config.include_columns = config.column_names; - } else if (!sharedState_->projectColumns.empty()) { - if (use_batch_read) { - // batch_read streams chunks directly to the consumer without - // post-projection, so push column projection down to the supplier. - read_config.include_columns = config.include_columns; - } else { - // full_read handles projection via project_chunk(). - read_config.include_columns = config.column_names; - } - } - - const auto& paths = fileSchema.paths; - if (paths.empty()) { + optionsBuilder_->projectColumns(config); + if (sharedState_->schema.file.paths.empty()) { THROW_INVALID_ARGUMENT_EXCEPTION("No file paths provided"); } - std::vector> suppliers; - suppliers.reserve(paths.size()); - for (const auto& path : paths) { - suppliers.push_back(std::make_shared( - path, read_config, - io::bindInputStream(sharedState_->stream_opener, path))); - } - - if (use_batch_read && !sharedState_->skipRows) { - batch_read(suppliers, ctx); - } else { - full_read(suppliers, ctx, config); - } -} - -void JsonReader::full_read( - const std::vector>& suppliers, - execution::Context& output, const JsonReadConfig& output_config) { - auto merged = read_all_chunks(suppliers); - - if (merged.col_num() == 0) { - // No data rows at all (e.g. an empty JSON array). Emit an empty result; - // there is nothing to validate, filter or project. - output.clear(); - return; - } - - int expected_cols = sharedState_->columnNum(); - if (expected_cols > 0 && - static_cast(merged.col_num()) != expected_cols && - sharedState_->projectColumns.empty()) { - THROW_IO_EXCEPTION( - "Column number mismatch between schema and JSON data, schema: " + - std::to_string(expected_cols) + - ", data: " + std::to_string(merged.col_num())); - } + // Each execution owns this supplier. Files are initialized one at a time, + // only when requested, and filtering/projection happen before yielding. + class JsonStreamSupplier final : public IDataChunkSupplier { + public: + JsonStreamSupplier(std::shared_ptr state, + JsonReadConfig config) + : state_(std::move(state)), config_(std::move(config)) {} - auto filtered = - filter_chunk(merged, sharedState_->skipRows, output_config.column_names, - sharedState_->parameters); - auto projected = project_chunk(filtered, output_config.column_names, - sharedState_->projectColumns.empty() - ? output_config.include_columns - : sharedState_->projectColumns); - output.clear(); - output.append_chunk(std::move(projected)); -} + int64_t RowNum() const override { return -1; } -void JsonReader::batch_read( - const std::vector>& suppliers, - execution::Context& output) { - output.clear(); - for (const auto& supplier : suppliers) { - while (auto chunk = supplier->GetNextChunk()) { - output.append_chunk(std::move(*chunk)); + std::shared_ptr GetNextChunk() override { + while (true) { + if (!current_) { + if (file_ == state_->schema.file.paths.size()) { + return nullptr; + } + const auto& path = state_->schema.file.paths[file_++]; + auto read_config = config_; + read_config.include_columns = + state_->skipRows ? config_.column_names : config_.include_columns; + current_ = std::make_shared( + path, read_config, + io::bindInputStream(state_->stream_opener, path)); + } + auto chunk = current_->GetNextChunk(); + if (!chunk) { + current_.reset(); + continue; + } + if (!state_->skipRows) { + return chunk; + } + auto filtered = filter_chunk(*chunk, state_->skipRows, + config_.column_names, state_->parameters); + auto projected = project_chunk(filtered, config_.column_names, + config_.include_columns); + return std::make_shared(std::move(projected)); + } } - } + + private: + std::shared_ptr state_; + JsonReadConfig config_; + size_t file_ = 0; + std::shared_ptr current_; + }; + return std::make_shared(sharedState_, std::move(config)); } result> JsonReader::inferSchema() { diff --git a/tests/execution/CMakeLists.txt b/tests/execution/CMakeLists.txt index e2068b6a2..41749e970 100644 --- a/tests/execution/CMakeLists.txt +++ b/tests/execution/CMakeLists.txt @@ -1,5 +1,6 @@ add_neug_test( execution_test + test_stream.cc test_logical_expr.cc test_plan_parser.cc test_query_cache.cc diff --git a/tests/execution/test_stream.cc b/tests/execution/test_stream.cc new file mode 100644 index 000000000..0247dd267 --- /dev/null +++ b/tests/execution/test_stream.cc @@ -0,0 +1,376 @@ +/** 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 + +#include "neug/common/columns/array_columns.h" +#include "neug/common/columns/list_columns.h" +#include "neug/common/columns/value_columns.h" +#include "neug/execution/common/operators/retrieve/sink.h" +#include "neug/execution/common/stream.h" +#include "neug/execution/execute/ops/batch/batch_update_utils.h" +#include "neug/execution/execute/ops/retrieve/sink.h" +#include "neug/execution/execute/pipeline.h" +#include "neug/storages/graph/property_graph.h" + +namespace neug::execution { +namespace { +using ChunkStream = Stream; + +DataChunk chunk(int64_t value, int alias = 0) { + ValueColumnBuilder builder; + builder.push_back_opt(value); + DataChunk out; + out.set(alias, builder.finish()); + return out; +} + +TEST(StreamTest, NestedColumnsMergeAcrossBatchesAndSerialize) { + std::function expect_value; + expect_value = [&](const Value& actual, const Value& expected) { + ASSERT_EQ(actual.type(), expected.type()); + ASSERT_EQ(actual.IsNull(), expected.IsNull()); + if (expected.IsNull()) { + return; + } + if (expected.type().id() == DataTypeId::kList || + expected.type().id() == DataTypeId::kArray) { + const auto& actual_children = expected.type().id() == DataTypeId::kList + ? ListValue::GetChildren(actual) + : ArrayValue::GetChildren(actual); + const auto& expected_children = expected.type().id() == DataTypeId::kList + ? ListValue::GetChildren(expected) + : ArrayValue::GetChildren(expected); + ASSERT_EQ(actual_children.size(), expected_children.size()); + for (size_t i = 0; i < actual_children.size(); ++i) { + expect_value(actual_children[i], expected_children[i]); + } + } else { + EXPECT_EQ(actual, expected); + } + }; + const auto list_type = DataType::List(DataType::INT32); + const auto array_type = DataType::Array(DataType::INT32, 2); + const auto nested_type = DataType::Array(array_type, 2); + const auto pair = + Value::ARRAY(array_type, {Value::INT32(1), Value(DataType::INT32)}); + for (const auto& values : std::vector>{ + {Value::LIST(DataType::INT32, + {Value::INT32(1), Value(DataType::INT32)}), + Value(list_type), Value::LIST(DataType::INT32, {}), + Value::LIST(DataType::INT32, {Value::INT32(4)})}, + {pair, Value(array_type), pair, Value(array_type)}, + {Value::ARRAY(nested_type, {pair, pair}), Value(nested_type), + Value::ARRAY(nested_type, {pair, pair}), Value(nested_type)}}) { + std::vector batches; + for (size_t begin = 0; begin < values.size(); begin += 2) { + auto builder = ColumnsUtils::create_builder(values.front().type()); + builder->push_back_elem(values[begin]); + builder->push_back_elem(values[begin + 1]); + ContextChunk batch; + batch.set(0, builder->finish()); + batches.push_back(std::move(batch)); + } + auto merged = collect_chunk(stream_from_batches(batches)); + ASSERT_TRUE(merged); + ASSERT_EQ(merged->row_num(), values.size()); + for (size_t row = 0; row < values.size(); ++row) { + expect_value(merged->get(0)->get_elem(row), values[row]); + } + EXPECT_EQ(merged->head(), merged->get(0)); + auto result = materialize( + stream_from_batches(std::move(batches), StreamMetadata{{0}})); + ASSERT_TRUE(result); + PropertyGraph graph; + GraphView view(graph); + StorageReadInterface storage(view, 0); + QueryResponse response; + Sink::sink_results(*result, storage, &response); + EXPECT_EQ(response.row_count(), values.size()); + EXPECT_EQ(response.arrays_size(), 1); + } +} + +TEST(StreamTest, IsLazyAndReleasesCursorOnCancellation) { + int pulls = 0; + std::weak_ptr weak; + { + auto lifetime = std::make_shared(0); + weak = lifetime; + ChunkStream stream([&, lifetime]() -> ChunkStream::NextResult { + ++pulls; + return std::optional(std::in_place, chunk(pulls)); + }); + EXPECT_EQ(pulls, 0); + ASSERT_TRUE(stream.Next()); + EXPECT_EQ(pulls, 1); + } + EXPECT_TRUE(weak.expired()); + EXPECT_EQ(pulls, 1); +} + +TEST(StreamTest, ErrorIsTerminalAndIsNotEndOfStream) { + int pulls = 0; + ChunkStream stream([&]() -> ChunkStream::NextResult { + ++pulls; + if (pulls == 1) { + return std::optional(std::in_place, chunk(1)); + } + THROW_IO_EXCEPTION("late read failure"); + }); + ASSERT_TRUE(stream.Next()); + auto failed = stream.Next(); + ASSERT_FALSE(failed); + EXPECT_NE(failed.error().error_message().find("late read failure"), + std::string::npos); + EXPECT_FALSE(stream.Next()); + EXPECT_EQ(pulls, 2); +} + +TEST(StreamTest, PreservesHeadsTagsSparseAliasesAndEmptyBatches) { + auto first = chunk(7, 3); + auto head = first.get(3); + Context original; + original.tag_ids = {-1, 3}; + original.append_chunk(std::move(first), head); + ValueColumnBuilder empty_builder; + DataChunk empty; + empty.set(3, empty_builder.finish()); + original.append_chunk(std::move(empty)); + original.append_chunk(DataChunk(), head); + auto restored = materialize(stream_from_context(std::move(original))); + ASSERT_TRUE(restored); + EXPECT_EQ(restored->tag_ids, (std::vector{-1, 3})); + ASSERT_EQ(restored->chunk_num(), 3); + EXPECT_EQ(restored->chunk(0).get(-1), restored->chunk(0).get(3)); + EXPECT_EQ(restored->chunk(1).row_num(), 0); + EXPECT_EQ(restored->chunk(1).get(3)->elem_type(), DataType::INT64); + EXPECT_EQ(restored->chunk(2).row_num(), 1); +} + +TEST(StreamTest, BatchTransformPreservesColumnIdentityAndDoesNotReadAhead) { + int pulls = 0; + auto data = chunk(42, 3); + auto column = data.get(3); + ChunkStream source( + [&]() -> ChunkStream::NextResult { + if (++pulls > 1) { + THROW_IO_EXCEPTION("must not read ahead"); + } + return std::optional(std::in_place, std::move(data), + column); + }, + StreamMetadata{{3, -1}}); + auto mapped = map_chunks(std::move(source), + [](ContextChunk&& batch) -> result { + return std::move(batch); + }); + EXPECT_EQ(pulls, 0); + auto first = mapped.Next(); + ASSERT_TRUE(first); + ASSERT_TRUE(*first); + EXPECT_EQ((**first).get(3), column); + EXPECT_EQ((**first).head(), column); + EXPECT_EQ(mapped.metadata().output_columns, (std::vector{3, -1})); + EXPECT_EQ(pulls, 1); +} + +TEST(StreamTest, StorageBridgePreservesMappingAndLateError) { + int pulls = 0; + ChunkStream stream([&]() -> ChunkStream::NextResult { + if (++pulls == 2) { + return tl::unexpected( + Status(StatusCode::ERR_IO_ERROR, "bad second batch")); + } + auto input = chunk(5, 2); + input.set(0, chunk(9).get(0)); + return std::optional(std::in_place, std::move(input)); + }); + ops::StreamChunkSupplier supplier(std::move(stream), {{2, "a"}, {0, "b"}}); + EXPECT_EQ(pulls, 0); + EXPECT_EQ(supplier.RowNum(), -1); + auto first = supplier.GetNextChunk(); + ASSERT_NE(first, nullptr); + EXPECT_EQ(first->get(0)->get_elem(0).GetValue(), 5); + EXPECT_EQ(first->get(1)->get_elem(0).GetValue(), 9); + EXPECT_EQ(supplier.GetNextChunk(), nullptr); + EXPECT_EQ(supplier.status().error_code(), StatusCode::ERR_IO_ERROR); +} + +TEST(StreamTest, CopyResultPreservesCardinalityWithoutInputColumns) { + auto output = materialize(ops::batch_insert_result(1000000)); + ASSERT_TRUE(output); + EXPECT_EQ(output->row_num(), 1000000); + EXPECT_EQ(output->col_num(), 0); + EXPECT_TRUE(output->tag_ids.empty()); +} + +struct Counts { + int produced = 0; + int consumed = 0; +}; +class CountingSource final : public IOperator { + public: + explicit CountingSource(Counts& counts) : counts_(counts) {} + std::string get_operator_name() const override { return "CountingSource"; } + ChunkStream Eval(IStorageInterface&, const ParamsMap&, ChunkStream&&, + OprTimer*) override { + return ChunkStream([this]() -> ChunkStream::NextResult { + EXPECT_EQ(counts_.produced, counts_.consumed); + if (counts_.produced == 3) { + return std::optional{}; + } + return std::optional(std::in_place, + chunk(++counts_.produced)); + }); + } + + private: + Counts& counts_; +}; +class CountingProject final : public IOperator { + public: + explicit CountingProject(Counts& counts) : counts_(counts) {} + std::string get_operator_name() const override { return "CountingProject"; } + Stream Eval(IStorageInterface&, const ParamsMap&, + Stream&& input, OprTimer*) override { + return map_chunks(std::move(input), + [this](ContextChunk&& chunk) -> result { + ++counts_.consumed; + EXPECT_EQ(chunk.row_num(), 1); + return std::move(chunk); + }); + } + + private: + Counts& counts_; +}; + +TEST(StreamTest, DeferredInitializationRunsOnceAndReportsErrorsOnNext) { + int initialized = 0; + auto stream = defer_stream( + Stream(), + [&](Stream &&) -> Stream { + ++initialized; + return error_stream(Status::InternalError("init failed")); + }); + EXPECT_EQ(initialized, 0); + auto first = stream.Next(); + ASSERT_FALSE(first); + EXPECT_NE(first.error().ToString().find("init failed"), std::string::npos); + EXPECT_EQ(initialized, 1); + auto second = stream.Next(); + ASSERT_FALSE(second); + EXPECT_EQ(second.error().ToString(), first.error().ToString()); + EXPECT_EQ(initialized, 1); + + auto throwing = + defer_stream(Stream(), + [](Stream &&) -> Stream { + THROW_IO_EXCEPTION("opening source failed"); + }); + auto error = throwing.Next(); + ASSERT_FALSE(error); + EXPECT_NE(error.error().ToString().find("opening source failed"), + std::string::npos); +} + +TEST(StreamTest, SinkMetadataPreservesOutputOrderAndEmptyResults) { + PropertyGraph graph; + GraphView view(graph); + StorageReadInterface storage(view, 0); + physical::PhysicalPlan plan; + auto* sink = plan.add_plan()->mutable_opr()->mutable_sink(); + for (int alias : {3, 0, 3}) { + sink->add_tags()->mutable_tag()->set_value(alias); + } + ops::SinkOprBuilder builder; + auto built = builder.Build(Schema(), ContextMeta(), plan, 0); + ASSERT_TRUE(built); + std::vector> operators; + operators.push_back(std::move(built->first)); + Pipeline pipeline(std::move(operators)); + for (bool empty : {false, true}) { + Context input; + if (!empty) { + auto data = chunk(7, 3); + data.set(0, chunk(9).get(0)); + input.append_chunk(std::move(data)); + } + auto result = pipeline.Execute(storage, std::move(input), {}, nullptr); + ASSERT_TRUE(result); + EXPECT_EQ(result->tag_ids, (std::vector{3, 0, 3})); + EXPECT_EQ(result->row_num(), empty ? 0 : 1); + } +} + +TEST(StreamTest, PipelineReportsInitializationFailureOnlyWhenPulled) { + class FailingSource final : public IOperator { + public: + explicit FailingSource(int& calls) : calls_(calls) {} + std::string get_operator_name() const override { return "FailingSource"; } + Stream Eval(IStorageInterface&, const ParamsMap&, + Stream&& input, + OprTimer*) override { + return defer_stream( + std::move(input), + [this](Stream &&) -> Stream { + ++calls_; + THROW_IO_EXCEPTION("source initialization failed"); + }); + } + + private: + int& calls_; + }; + int calls = 0; + std::vector> operators; + operators.push_back(std::make_unique(calls)); + Pipeline pipeline(std::move(operators)); + PropertyGraph graph; + GraphView view(graph); + StorageReadInterface storage(view, 0); + auto output = + pipeline.ExecuteStream(storage, Stream(), {}, nullptr); + EXPECT_EQ(calls, 0); + auto error = output.Next(); + ASSERT_FALSE(error); + EXPECT_NE(error.error().ToString().find("FailingSource"), std::string::npos); + EXPECT_EQ(calls, 1); + EXPECT_FALSE(output.Next()); + EXPECT_EQ(calls, 1); +} + +TEST(StreamTest, PipelinePullsThroughChunkwiseOperatorsAndProfilesRows) { + Counts counts; + std::vector> operators; + operators.push_back(std::make_unique(counts)); + operators.push_back(std::make_unique(counts)); + Pipeline pipeline(std::move(operators)); + PropertyGraph graph; + GraphView view(graph); + StorageReadInterface storage(view, 0); + OprTimer timer; + auto result = pipeline.Execute(storage, Context(), {}, &timer); + ASSERT_TRUE(result) << result.error().ToString(); + EXPECT_EQ(result->row_num(), 3); + EXPECT_EQ(counts.produced, 3); + EXPECT_EQ(counts.consumed, 3); + auto profile = OprTimer::ToProfileResult(&timer); + ASSERT_EQ(profile.operators_size(), 2); + EXPECT_EQ(profile.operators(0).output_rows(), 3); + EXPECT_EQ(profile.operators(1).output_rows(), 3); +} +} // namespace +} // namespace neug::execution diff --git a/tests/utils/json_test.cc b/tests/utils/json_test.cc index 35fc9a6ff..23770971a 100644 --- a/tests/utils/json_test.cc +++ b/tests/utils/json_test.cc @@ -24,6 +24,7 @@ #include "neug/compiler/common/case_insensitive_map.h" #include "neug/execution/common/context.h" #include "neug/generated/proto/plan/basic_type.pb.h" +#include "neug/storages/loader/loader_utils.h" #include "neug/utils/io/read/common/options.h" #include "neug/utils/io/read/common/schema.h" #include "neug/utils/io/reader.h" @@ -166,13 +167,17 @@ TEST_F(JsonTest, TestJsonArray) { "\"name\": \"Bob\", \"age\": 30}]"); auto sharedState = createSharedState( "test_json_array.json", {"id", "name", "age"}, - {createUInt32Type(), createStringType(), createDoubleType()}, - {{"batch_read", "false"}}); + {createUInt32Type(), createStringType(), createDoubleType()}, {}); auto reader = createJsonReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); @@ -192,14 +197,19 @@ TEST_F(JsonTest, TestJsonArrayColumn) { createJsonFile("test_json_array_column.json", "[{\"id\": 1, \"readings\": [1, 2, 3]}, " "{\"id\": 2, \"readings\": [4, 5, 6]}]"); - auto sharedState = createSharedState( - "test_json_array_column.json", {"id", "readings"}, - {createUInt32Type(), createInt64ArrayType(3)}, {{"batch_read", "false"}}); + auto sharedState = + createSharedState("test_json_array_column.json", {"id", "readings"}, + {createUInt32Type(), createInt64ArrayType(3)}, {}); auto reader = createJsonReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 2); EXPECT_EQ(ctx.row_num(), 2); @@ -219,13 +229,18 @@ TEST_F(JsonTest, TestJsonNullArrayColumnIsOptional) { "[{\"id\":1,\"readings\":[1,2]}," "{\"id\":2,\"readings\":null}," "{\"id\":3,\"readings\":[3,4]}]"); - auto sharedState = createSharedState( - "test_json_null_array_column.json", {"id", "readings"}, - {createUInt32Type(), createInt32ArrayType(2)}, {{"batch_read", "false"}}); + auto sharedState = + createSharedState("test_json_null_array_column.json", {"id", "readings"}, + {createUInt32Type(), createInt32ArrayType(2)}, {}); auto reader = createJsonReader(sharedState); execution::Context ctx; - reader->read(std::make_shared(), ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } ASSERT_EQ(ctx.row_num(), 3); auto readings_col = ctx.chunk(0).columns()[1]; @@ -250,13 +265,18 @@ TEST_F(JsonTest, TestJsonArrayColumnLengthMismatch) { "[{\"id\": 1, \"readings\": [1, 2]}]"); auto sharedState = createSharedState( "test_json_array_length_mismatch.json", {"id", "readings"}, - {createUInt32Type(), createInt64ArrayType(3)}, {{"batch_read", "false"}}); + {createUInt32Type(), createInt64ArrayType(3)}, {}); auto reader = createJsonReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; try { - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } FAIL() << "Expected an ARRAY length mismatch"; } catch (const std::exception& error) { EXPECT_NE( @@ -269,15 +289,20 @@ TEST_F(JsonTest, TestJsonArrayColumnLengthMismatch) { TEST_F(JsonTest, TestJsonArrayColumnRejectsNonArray) { createJsonFile("test_json_non_array.json", "[{\"id\": 1, \"readings\": 42}]"); - auto sharedState = createSharedState( - "test_json_non_array.json", {"id", "readings"}, - {createUInt32Type(), createInt64ArrayType(3)}, {{"batch_read", "false"}}); + auto sharedState = + createSharedState("test_json_non_array.json", {"id", "readings"}, + {createUInt32Type(), createInt64ArrayType(3)}, {}); auto reader = createJsonReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; try { - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } FAIL() << "Expected a non-array conversion error"; } catch (const std::exception& error) { EXPECT_NE(std::string(error.what()) @@ -294,11 +319,15 @@ TEST_F(JsonTest, TestJsonRecursiveListColumn) { "{\"id\":4,\"nested\":[[null,[\"h\"]],null]}]"); auto sharedState = createSharedState("test_json_recursive_list.json", {"id", "nested"}, - {createUInt32Type(), createNestedStringListType()}, - {{"batch_read", "false"}}); + {createUInt32Type(), createNestedStringListType()}, {}); auto reader = createJsonReader(sharedState); execution::Context ctx; - reader->read(std::make_shared(), ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } ASSERT_EQ(ctx.row_num(), 4); auto nested = ctx.chunk(0).columns()[1]; @@ -337,14 +366,18 @@ TEST_F(JsonTest, TestJsonArrayStreaming) { "\"name\": \"Bob\", \"age\": 30}]"); auto sharedState = createSharedState( "test_json_array_stream.json", {"id", "name", "age"}, - {createUInt32Type(), createStringType(), createDoubleType()}, - {{"batch_read", "false"}}); + {createUInt32Type(), createStringType(), createDoubleType()}, {}); sharedState->stream_opener = localStreamOpener(); auto reader = createJsonReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); @@ -366,13 +399,17 @@ TEST_F(JsonTest, TestEmptyJsonArrayReturnsEmpty) { createJsonFile("test_json_array_empty.json", "[]"); auto sharedState = createSharedState( "test_json_array_empty.json", {"id", "name", "age"}, - {createUInt32Type(), createStringType(), createDoubleType()}, - {{"batch_read", "false"}}); + {createUInt32Type(), createStringType(), createDoubleType()}, {}); auto reader = createJsonReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - EXPECT_NO_THROW(reader->read(localState, ctx)); + EXPECT_NO_THROW([&] { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + }()); EXPECT_EQ(ctx.col_num(), 0); EXPECT_EQ(ctx.row_num(), 0); } @@ -382,14 +419,18 @@ TEST_F(JsonTest, TestEmptyJsonArrayStreamingReturnsEmpty) { createJsonFile("test_json_array_empty_stream.json", "[]"); auto sharedState = createSharedState( "test_json_array_empty_stream.json", {"id", "name", "age"}, - {createUInt32Type(), createStringType(), createDoubleType()}, - {{"batch_read", "false"}}); + {createUInt32Type(), createStringType(), createDoubleType()}, {}); sharedState->stream_opener = localStreamOpener(); auto reader = createJsonReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - EXPECT_NO_THROW(reader->read(localState, ctx)); + EXPECT_NO_THROW([&] { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + }()); EXPECT_EQ(ctx.col_num(), 0); EXPECT_EQ(ctx.row_num(), 0); } @@ -403,14 +444,19 @@ TEST_F(JsonTest, TestJsonLinesCRLF) { "{\"id\": 1, \"age\": 25}\r\n" "\r\n" "{\"id\": 2, \"age\": 30}\r"); - auto sharedState = createSharedState("test_jsonl_crlf.json", {"id", "age"}, - {createUInt32Type(), createDoubleType()}, - {{"batch_read", "false"}}); + auto sharedState = + createSharedState("test_jsonl_crlf.json", {"id", "age"}, + {createUInt32Type(), createDoubleType()}, {}); auto reader = createJsonLinesReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - ASSERT_NO_THROW(reader->read(localState, ctx)); + ASSERT_NO_THROW([&] { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + }()); ASSERT_EQ(ctx.row_num(), 2); auto col0 = ctx.chunk(0).columns()[0]; EXPECT_EQ(col0->get_elem(0).GetValue(), 1u); @@ -427,15 +473,20 @@ TEST_F(JsonTest, TestJsonLinesCRLFStreaming) { "{\"id\": 1, \"age\": 25}\r\n" "\r\n" "{\"id\": 2, \"age\": 30}\r"); - auto sharedState = createSharedState( - "test_jsonl_crlf_stream.json", {"id", "age"}, - {createUInt32Type(), createDoubleType()}, {{"batch_read", "false"}}); + auto sharedState = + createSharedState("test_jsonl_crlf_stream.json", {"id", "age"}, + {createUInt32Type(), createDoubleType()}, {}); sharedState->stream_opener = localStreamOpener(); auto reader = createJsonLinesReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - ASSERT_NO_THROW(reader->read(localState, ctx)); + ASSERT_NO_THROW([&] { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + }()); ASSERT_EQ(ctx.row_num(), 2); auto col0 = ctx.chunk(0).columns()[0]; EXPECT_EQ(col0->get_elem(0).GetValue(), 1u); @@ -446,15 +497,20 @@ TEST_F(JsonTest, TestJsonLinesCRLFStreaming) { TEST_F(JsonTest, TestJsonLinesLFUnchanged) { createJsonFile("test_jsonl_lf.json", "{\"id\": 1, \"age\": 25}\n{\"id\": 2, \"age\": 30}\n"); - auto sharedState = createSharedState("test_jsonl_lf.json", {"id", "age"}, - {createUInt32Type(), createDoubleType()}, - {{"batch_read", "false"}}); + auto sharedState = + createSharedState("test_jsonl_lf.json", {"id", "age"}, + {createUInt32Type(), createDoubleType()}, {}); sharedState->stream_opener = localStreamOpener(); auto reader = createJsonLinesReader(sharedState); - auto localState = std::make_shared(); + execution::Context ctx; - ASSERT_NO_THROW(reader->read(localState, ctx)); + ASSERT_NO_THROW([&] { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + }()); ASSERT_EQ(ctx.row_num(), 2); auto col0 = ctx.chunk(0).columns()[0]; EXPECT_EQ(col0->get_elem(1).GetValue(), 2u); diff --git a/tests/utils/test_reader.cc b/tests/utils/test_reader.cc index bf58aa820..e06f6c169 100644 --- a/tests/utils/test_reader.cc +++ b/tests/utils/test_reader.cc @@ -31,6 +31,25 @@ namespace neug { namespace test { +TEST_F(ReaderTest, TestCsvSupplierOpensEachFileOnDemand) { + auto state = createSharedState("lazy_first.csv", {"id"}, {createInt64Type()}, + {{"skip_rows", "1"}, {"batch_size", "1"}}); + state->schema.file.paths.push_back(std::string(ARROW_READER_TEST_DIR) + + "/lazy_second.csv"); + auto supplier = createCsvReader(state)->getDataChunkSupplier(); + EXPECT_EQ(supplier->RowNum(), -1); + // Neither file exists when the stream is constructed. + createCsvFile("lazy_first.csv", "id\n1\n"); + auto first = supplier->GetNextChunk(); + ASSERT_NE(first, nullptr); + EXPECT_EQ(first->get(0)->get_elem(0).GetValue(), 1); + createCsvFile("lazy_second.csv", "id\n2\n"); + auto second = supplier->GetNextChunk(); + ASSERT_NE(second, nullptr); + EXPECT_EQ(second->get(0)->get_elem(0).GetValue(), 2); + EXPECT_EQ(supplier->GetNextChunk(), nullptr); +} + namespace { class VectorChunkSupplier final : public IDataChunkSupplier { @@ -57,6 +76,20 @@ class VectorChunkSupplier final : public IDataChunkSupplier { int64_t row_count_ = 0; }; +DataChunk readAllChunks( + const std::vector>& suppliers) { + std::vector> chunks; + for (const auto& supplier : suppliers) { + if (!supplier) { + THROW_INVALID_ARGUMENT_EXCEPTION("Data chunk supplier is null"); + } + while (auto chunk = supplier->GetNextChunk()) { + chunks.push_back(std::move(chunk)); + } + } + return reader::merge_chunks(std::move(chunks)); +} + std::shared_ptr<::common::Expression> comparison(const std::string& column, ::common::Logical logical, int32_t value) { @@ -153,7 +186,7 @@ TEST_F(ReaderTest, TestReadAllChunksMergesNestedValuesAndArrayNulls) { auto second = std::make_shared( std::vector{nestedChunk(3)}); - auto merged = reader::read_all_chunks({first, second}); + auto merged = readAllChunks({first, second}); ASSERT_EQ(merged.col_num(), 3); ASSERT_EQ(merged.row_num(), 4); EXPECT_EQ(merged.get(0)->get_elem(0).GetValue(), 1); @@ -189,10 +222,8 @@ TEST_F(ReaderTest, TestReadAllChunksRejectsIncompatibleChunks) { auto supplier = std::make_shared( std::vector{one_column, two_columns}); - EXPECT_THROW(reader::read_all_chunks({supplier}), - exception::InvalidArgumentException); - EXPECT_THROW(reader::read_all_chunks({nullptr}), - exception::InvalidArgumentException); + EXPECT_THROW(readAllChunks({supplier}), exception::InvalidArgumentException); + EXPECT_THROW(readAllChunks({nullptr}), exception::InvalidArgumentException); } TEST_F(ReaderTest, TestCommonRowFilterUsesSqlNullSemantics) { @@ -473,15 +504,18 @@ TEST_F(ReaderTest, TestBasicCsvRead) { std::vector> columnTypes = { createInt32Type(), createStringType(), createDoubleType()}; - auto sharedState = - createSharedState("test1.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "false"}}); + auto sharedState = createSharedState("test1.csv", columnNames, columnTypes, + {{"skip_rows", "1"}}); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Verify data: should have 3 columns EXPECT_EQ(ctx.col_num(), 3); @@ -497,16 +531,19 @@ TEST_F(ReaderTest, TestCsvWithTabDelimiter) { std::vector> columnTypes = { createInt32Type(), createStringType(), createDoubleType()}; - auto sharedState = createSharedState( - "test2.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"delim", "\t"}, {"batch_read", "false"}}); + auto sharedState = createSharedState("test2.csv", columnNames, columnTypes, + {{"skip_rows", "1"}, {"delim", "\t"}}); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); @@ -521,17 +558,19 @@ TEST_F(ReaderTest, TestCsvWithCustomQuoting) { std::vector> columnTypes = { createInt32Type(), createStringType(), createDoubleType()}; - auto sharedState = createSharedState("test3.csv", columnNames, columnTypes, - {{"quote", "'"}, - {"delim", ","}, - {"skip_rows", "1"}, - {"batch_read", "false"}}); + auto sharedState = + createSharedState("test3.csv", columnNames, columnTypes, + {{"quote", "'"}, {"delim", ","}, {"skip_rows", "1"}}); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); @@ -545,14 +584,18 @@ TEST_F(ReaderTest, TestCsvWithNoHeader) { std::vector> columnTypes = { createInt32Type(), createStringType(), createDoubleType()}; - auto sharedState = createSharedState("test4.csv", columnNames, columnTypes, - {{"batch_read", "false"}}); + auto sharedState = + createSharedState("test4.csv", columnNames, columnTypes, {}); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); @@ -572,15 +615,19 @@ TEST_F(ReaderTest, TestBatchRead) { std::vector> columnTypes = { createInt32Type(), createStringType(), createDoubleType()}; - auto sharedState = createSharedState( - "test5.csv", columnNames, columnTypes, - {{"batch_read", "true"}, {"batch_size", "1024"}, {"skip_rows", "1"}}); + auto sharedState = + createSharedState("test5.csv", columnNames, columnTypes, + {{"batch_size", "1024"}, {"skip_rows", "1"}}); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Batch mode: data is materialized into Context chunks EXPECT_GT(ctx.chunk_num(), 0); @@ -602,16 +649,19 @@ TEST_F(ReaderTest, TestColumnPruning) { // Project only "id" and "score" columns (exclude "name") std::vector projectColumns = {"id", "score"}; - auto sharedState = createSharedState( - "test6.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "false"}}, projectColumns); + auto sharedState = createSharedState("test6.csv", columnNames, columnTypes, + {{"skip_rows", "1"}}, projectColumns); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Should only have 2 columns (id and score) EXPECT_EQ(ctx.col_num(), 2); @@ -632,16 +682,19 @@ TEST_F(ReaderTest, TestFilterPushdown) { // Filter: score > 90.0 auto filterExpr = createFilterExpression("score", ValueConverter::fromDouble(90.0)); - auto sharedState = createSharedState( - "test7.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "false"}}, {}, filterExpr); + auto sharedState = createSharedState("test7.csv", columnNames, columnTypes, + {{"skip_rows", "1"}}, {}, filterExpr); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Should filter out rows with score <= 90.0 // Expected: Alice (95.5) and Charlie (92.5) - 2 rows @@ -666,15 +719,18 @@ TEST_F(ReaderTest, TestColumnPruningAndFilterPushdown) { createFilterExpression("score", ValueConverter::fromDouble(90.0)); auto sharedState = createSharedState("test8.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "false"}}, - projectColumns, filterExpr); + {{"skip_rows", "1"}}, projectColumns, filterExpr); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Should have 2 columns (id, score) and filtered rows (score > 90.0) EXPECT_EQ(ctx.col_num(), 2); @@ -691,19 +747,22 @@ TEST_F(ReaderTest, TestMultipleFiles) { std::vector> columnTypes = { createInt32Type(), createStringType(), createDoubleType()}; - auto sharedState = - createSharedState("test9a.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "false"}}); + auto sharedState = createSharedState("test9a.csv", columnNames, columnTypes, + {{"skip_rows", "1"}}); // Add second file sharedState->schema.file.paths.push_back(std::string(ARROW_READER_TEST_DIR) + "/test9b.csv"); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Should read all rows from both files (4 rows total) EXPECT_EQ(ctx.col_num(), 3); @@ -721,15 +780,18 @@ TEST_F(ReaderTest, TestForceColumnTypeConversion) { std::vector> columnTypes = { createInt32Type(), createStringType(), createInt64Type()}; - auto sharedState = - createSharedState("test10.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "false"}}); + auto sharedState = createSharedState("test10.csv", columnNames, columnTypes, + {{"skip_rows", "1"}}); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 3); @@ -764,16 +826,19 @@ TEST_F(ReaderTest, TestMultiColumnAndFilterPushdown) { createFilterExpression("score", ValueConverter::fromDouble(90.0)); auto andExpr = createAndExpression(leftExpr, rightExpr); - auto sharedState = createSharedState( - "test11.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "false"}}, {}, andExpr); + auto sharedState = createSharedState("test11.csv", columnNames, columnTypes, + {{"skip_rows", "1"}}, {}, andExpr); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Should have 3 columns EXPECT_EQ(ctx.col_num(), 3); @@ -782,7 +847,7 @@ TEST_F(ReaderTest, TestMultiColumnAndFilterPushdown) { EXPECT_EQ(ctx.row_num(), 2); } -// Test 12: batch_read=true with filter (skipRows) should fallback to full_read +// Test 12: Supplier applies filter (skipRows) to batches TEST_F(ReaderTest, TestBatchReadWithFilter) { createCsvFile("test12.csv", "id|name|score\n1|Alice|95.5\n2|Bob|87.0\n3|Charlie|92.5\n4|" @@ -792,19 +857,22 @@ TEST_F(ReaderTest, TestBatchReadWithFilter) { std::vector> columnTypes = { createInt32Type(), createStringType(), createDoubleType()}; - // Filter: score > 90.0 with batch_read=true + // Filter: score > 90.0 with supplier batches auto filterExpr = createFilterExpression("score", ValueConverter::fromDouble(90.0)); - auto sharedState = createSharedState( - "test12.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "true"}}, {}, filterExpr); + auto sharedState = createSharedState("test12.csv", columnNames, columnTypes, + {{"skip_rows", "1"}}, {}, filterExpr); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Should filter out rows with score <= 90.0 // Expected: Alice (95.5) and Charlie (92.5) - 2 rows @@ -812,7 +880,7 @@ TEST_F(ReaderTest, TestBatchReadWithFilter) { EXPECT_EQ(ctx.row_num(), 2); } -// Test 13: batch_read=true with filter AND column projection +// Test 13: supplier batches with filter AND column projection TEST_F(ReaderTest, TestBatchReadWithFilterAndProjection) { createCsvFile("test13.csv", "id|name|score\n1|Alice|95.5\n2|Bob|87.0\n3|Charlie|92.5\n4|" @@ -823,20 +891,24 @@ TEST_F(ReaderTest, TestBatchReadWithFilterAndProjection) { createInt32Type(), createStringType(), createDoubleType()}; // Project only "id" and "score" columns, filter: score > 90.0, - // batch_read=true + // supplier batches std::vector projectColumns = {"id", "score"}; auto filterExpr = createFilterExpression("score", ValueConverter::fromDouble(90.0)); - auto sharedState = createSharedState( - "test13.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "true"}}, projectColumns, filterExpr); + auto sharedState = + createSharedState("test13.csv", columnNames, columnTypes, + {{"skip_rows", "1"}}, projectColumns, filterExpr); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Should have 2 columns (id, score) and filtered rows (score > 90.0) EXPECT_EQ(ctx.col_num(), 2); @@ -855,15 +927,18 @@ TEST_F(ReaderTest, TestBasicJsonRead) { std::vector> columnTypes = { createInt64Type(), createStringType(), createDoubleType()}; - auto sharedState = - createJsonSharedState("test_json_basic.json", columnNames, columnTypes, - {{"batch_read", "false"}}); + auto sharedState = createJsonSharedState("test_json_basic.json", columnNames, + columnTypes, {}); auto reader = createJsonReader(sharedState, false); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); @@ -877,19 +952,23 @@ TEST_F(ReaderTest, TestJsonNonExistentColumnThrows) { std::vector> columnTypes = { createInt64Type(), createStringType(), createDoubleType()}; - auto sharedState = - createJsonSharedState("test_json_nonexist.json", columnNames, columnTypes, - {{"batch_read", "false"}}); + auto sharedState = createJsonSharedState("test_json_nonexist.json", + columnNames, columnTypes, {}); auto reader = createJsonReader(sharedState, false); - auto localState = std::make_shared(); execution::Context ctx; - EXPECT_THROW(reader->read(localState, ctx), - exception::SchemaMismatchException); + EXPECT_THROW( + [&] { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + }(), + exception::SchemaMismatchException); } -// Test: JSON batch_read=true with filter should fallback to full_read +// Test: JSON supplier applies filters to batches TEST_F(ReaderTest, TestJsonBatchReadWithFilter) { createJsonFile("test_json_filter.json", "{\"id\":1,\"name\":\"Alice\",\"score\":95.5}\n" @@ -901,20 +980,23 @@ TEST_F(ReaderTest, TestJsonBatchReadWithFilter) { std::vector> columnTypes = { createInt64Type(), createStringType(), createDoubleType()}; - // Filter: score > 90.0 with batch_read=true + // Filter: score > 90.0 with supplier batches auto filterExpr = createFilterExpression("score", ValueConverter::fromDouble(90.0)); - auto sharedState = - createJsonSharedState("test_json_filter.json", columnNames, columnTypes, - {{"batch_read", "true"}}); + auto sharedState = createJsonSharedState("test_json_filter.json", columnNames, + columnTypes, {}); sharedState->skipRows = filterExpr; auto reader = createJsonReader(sharedState, false); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Should filter out rows with score <= 90.0 // Expected: Alice (95.5) and Charlie (92.5) - 2 rows @@ -922,7 +1004,7 @@ TEST_F(ReaderTest, TestJsonBatchReadWithFilter) { EXPECT_EQ(ctx.row_num(), 2); } -// Test: JSON batch_read=true with filter AND column projection +// Test: JSON supplier batches with filter AND column projection TEST_F(ReaderTest, TestJsonBatchReadWithFilterAndProjection) { createJsonFile("test_json_filter_proj.json", "{\"id\":1,\"name\":\"Alice\",\"score\":95.5}\n" @@ -934,21 +1016,24 @@ TEST_F(ReaderTest, TestJsonBatchReadWithFilterAndProjection) { std::vector> columnTypes = { createInt64Type(), createStringType(), createDoubleType()}; - // Project only "id" and "score", filter: score > 90.0, batch_read=true + // Project only "id" and "score", filter: score > 90.0, supplier batches auto filterExpr = createFilterExpression("score", ValueConverter::fromDouble(90.0)); - auto sharedState = - createJsonSharedState("test_json_filter_proj.json", columnNames, - columnTypes, {{"batch_read", "true"}}); + auto sharedState = createJsonSharedState("test_json_filter_proj.json", + columnNames, columnTypes, {}); sharedState->skipRows = filterExpr; sharedState->projectColumns = {"id", "score"}; auto reader = createJsonReader(sharedState, false); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } // Should have 2 columns (id, score) and filtered rows (score > 90.0) EXPECT_EQ(ctx.col_num(), 2); @@ -969,16 +1054,19 @@ TEST_F(ReaderTest, TestCsvStreamingRead) { std::vector> columnTypes = { createInt32Type(), createStringType(), createDoubleType()}; - auto sharedState = - createSharedState("stream1.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "false"}}); + auto sharedState = createSharedState("stream1.csv", columnNames, columnTypes, + {{"skip_rows", "1"}}); sharedState->stream_opener = localStreamOpener(); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 3); @@ -998,16 +1086,20 @@ TEST_F(ReaderTest, TestCsvStreamingBatchRead) { std::vector> columnTypes = { createInt32Type(), createStringType(), createDoubleType()}; - auto sharedState = createSharedState( - "stream_batch.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "true"}, {"batch_size", "32"}}); + auto sharedState = + createSharedState("stream_batch.csv", columnNames, columnTypes, + {{"skip_rows", "1"}, {"batch_size", "32"}}); sharedState->stream_opener = localStreamOpener(); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(count_batch_row_num(ctx), 100); @@ -1025,17 +1117,18 @@ TEST_F(ReaderTest, TestCsvStreamingReadWithQuoting) { auto sharedState = createSharedState("stream_quote.csv", columnNames, columnTypes, - {{"quote", "'"}, - {"delim", ","}, - {"skip_rows", "1"}, - {"batch_read", "false"}}); + {{"quote", "'"}, {"delim", ","}, {"skip_rows", "1"}}); sharedState->stream_opener = localStreamOpener(); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); @@ -1053,16 +1146,19 @@ TEST_F(ReaderTest, TestJsonStreamingRead) { std::vector> columnTypes = { createInt64Type(), createStringType(), createDoubleType()}; - auto sharedState = - createJsonSharedState("test_json_stream.json", columnNames, columnTypes, - {{"batch_read", "false"}}); + auto sharedState = createJsonSharedState("test_json_stream.json", columnNames, + columnTypes, {}); sharedState->stream_opener = localStreamOpener(); auto reader = createJsonReader(sharedState, false); - auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + } EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 3); @@ -1078,15 +1174,18 @@ TEST_F(ReaderTest, TestCsvHeaderOnlyFileReturnsEmpty) { std::vector> columnTypes = { createInt32Type(), createStringType(), createDoubleType()}; - auto sharedState = - createSharedState("header_only.csv", columnNames, columnTypes, - {{"skip_rows", "1"}, {"batch_read", "false"}}); + auto sharedState = createSharedState("header_only.csv", columnNames, + columnTypes, {{"skip_rows", "1"}}); auto reader = createCsvReader(sharedState); - auto localState = std::make_shared(); execution::Context ctx; - EXPECT_NO_THROW(reader->read(localState, ctx)); + EXPECT_NO_THROW([&] { + auto supplier = reader->getDataChunkSupplier(); + while (auto chunk = supplier->GetNextChunk()) { + ctx.append_chunk(std::move(*chunk)); + } + }()); EXPECT_EQ(ctx.col_num(), 0); EXPECT_EQ(ctx.row_num(), 0); } diff --git a/tools/python_bind/tests/test_load_array.py b/tools/python_bind/tests/test_load_array.py index 9d7231e37..69f1007f6 100644 --- a/tools/python_bind/tests/test_load_array.py +++ b/tools/python_bind/tests/test_load_array.py @@ -756,7 +756,15 @@ def test_parquet_list_and_array_preserve_null_elements(self): "UNWIND values AS value RETURN value ORDER BY value" ) ) - assert rows == [[1], [None], [3], [4], [5]] + assert rows == [[1], [3], [4], [5], [None]] + + rows = list( + self.conn.execute( + f'LOAD FROM "{list_path}" ' + "UNWIND values AS value RETURN value ORDER BY value DESC" + ) + ) + assert rows == [[None], [5], [4], [3], [1]] rows = list( self.conn.execute(f'LOAD FROM "{array_path}" RETURN id, vec ORDER BY id') diff --git a/tools/python_bind/tests/test_parquet_backend_contract.py b/tools/python_bind/tests/test_parquet_backend_contract.py index a5f14c141..d28c33f30 100644 --- a/tools/python_bind/tests/test_parquet_backend_contract.py +++ b/tools/python_bind/tests/test_parquet_backend_contract.py @@ -111,17 +111,23 @@ def type_file(tmp_path): return path -def test_reader_preserves_types_nulls_and_nested_values(connection, type_file): +@pytest.mark.parametrize("batch_rows", [1, 2]) +@pytest.mark.parametrize("ordered", [False, True]) +def test_reader_preserves_types_nulls_and_nested_values( + connection, type_file, batch_rows, ordered +): path = type_file rows = list( connection.execute( - f'LOAD FROM "{path.as_posix()}" ' + f'LOAD FROM "{path.as_posix()}" (PARQUET_BATCH_ROWS={batch_rows}) ' "RETURN id, enabled, signed_value, unsigned_value, score, label, " "event_date, timestamp_s, timestamp_ms, items, fixed3, matrix2x2 " - "ORDER BY id" + + ("ORDER BY id" if ordered else "") ) ) + if not ordered: + rows.sort(key=lambda row: row[0]) assert rows == [ [ 1, @@ -222,25 +228,25 @@ def test_reader_preserves_types_nulls_and_nested_values(connection, type_file): ), ], ) -@pytest.mark.parametrize("batch_read", [False, True]) +@pytest.mark.parametrize("batch_rows", [1, 2]) def test_reader_preserves_complete_predicates( - connection, type_file, predicate, expected, batch_read + connection, type_file, predicate, expected, batch_rows ): rows = list( connection.execute( f'LOAD FROM "{type_file.as_posix()}" ' - f"(batch_read={str(batch_read).lower()}, row_batch_size=1) " + f"(PARQUET_BATCH_ROWS={batch_rows}) " f"WHERE {predicate} RETURN id ORDER BY id" ) ) assert rows == [[value] for value in expected] -@pytest.mark.parametrize("batch_read", [False, True]) -def test_reader_fallback_binds_current_parameters(connection, type_file, batch_read): +@pytest.mark.parametrize("batch_rows", [1, 2]) +def test_reader_fallback_binds_current_parameters(connection, type_file, batch_rows): query = ( f'LOAD FROM "{type_file.as_posix()}" ' - f"(batch_read={str(batch_read).lower()}, row_batch_size=1) " + f"(PARQUET_BATCH_ROWS={batch_rows}) " "WHERE CASE WHEN score IS NULL THEN $missing ELSE score END > $minimum " "RETURN id ORDER BY id" ) diff --git a/tools/python_bind/tests/test_stream_execution.py b/tools/python_bind/tests/test_stream_execution.py new file mode 100644 index 000000000..ac00730a1 --- /dev/null +++ b/tools/python_bind/tests/test_stream_execution.py @@ -0,0 +1,157 @@ +"""Operator streams retain one COPY pipeline and per-execution reader state.""" + +import pytest + + +def operator_names(result): + return [op["operator_name"] for op in result.get_profile_metrics()["operators"]] + + +@pytest.mark.parametrize("file_format", ["csv", "jsonl", "json"]) +def test_copy_stream_uses_regular_operators(empty_db, tmp_path, file_format): + _, conn = empty_db + path = tmp_path / f"nodes.{file_format}" + content = { + "csv": "id|name\n1|one\n2|two\n3|three\n", + "jsonl": '{"id":1,"name":"one"}\n{"id":2,"name":"two"}\n' + '{"id":3,"name":"three"}\n', + "json": '[{"id":1,"name":"one"},{"id":2,"name":"two"},' + '{"id":3,"name":"three"}]', + } + path.write_text(content[file_format], encoding="utf-8") + conn.execute("CREATE NODE TABLE person(id INT64, name STRING, PRIMARY KEY(id))") + result = conn.execute(f'PROFILE COPY person FROM "{path}" (batch_size=1)') + assert len(result) == 3 + names = operator_names(result) + assert "DataSourceOpr" in names + assert "BatchInsertVertexOpr" in names + assert not any("FusedCSV" in name for name in names) + assert list(conn.execute("MATCH (p:person) RETURN p.id, p.name ORDER BY p.id")) == [ + [1, "one"], + [2, "two"], + [3, "three"], + ] + + +def test_copy_stream_subquery_projection(empty_db, tmp_path): + _, conn = empty_db + path = tmp_path / "projection.csv" + path.write_text("id|score|name\n1|10|one\n2|20|two\n3|30|three\n") + conn.execute( + "CREATE NODE TABLE selected(id INT64, name STRING, value INT64, " + "PRIMARY KEY(id))" + ) + result = conn.execute( + f'PROFILE COPY selected FROM (LOAD FROM "{path}" ' + "(header=true, batch_size=1) WHERE score > 10 " + "RETURN id, name, score + 1 AS value)" + ) + names = operator_names(result) + assert "DataSourceOpr" in names + assert "ProjectOpr" in names + assert "BatchInsertVertexOpr" in names + assert list( + conn.execute("MATCH (n:selected) RETURN n.id, n.name, n.value ORDER BY n.id") + ) == [[2, "two", 21], [3, "three", 31]] + + +def test_copy_stream_late_parse_error_rolls_back(empty_db, tmp_path): + _, conn = empty_db + conn.execute("CREATE NODE TABLE person(id INT64, name STRING, PRIMARY KEY(id))") + conn.execute("CREATE (:person {id: 0, name: 'original'})") + path = tmp_path / "bad.csv" + path.write_text("id|name\n1|first\nnot-an-integer|bad\n") + with pytest.raises(RuntimeError): + conn.execute(f'COPY person FROM "{path}" (batch_size=1)') + assert list(conn.execute("MATCH (p:person) RETURN p.id, p.name")) == [ + [0, "original"] + ] + path.write_text("id|name\n1|first\n2|second\n") + conn.execute(f'COPY person FROM "{path}" (batch_size=1)') + assert list(conn.execute("MATCH (p:person) RETURN p.id ORDER BY p.id")) == [ + [0], + [1], + [2], + ] + + +def test_load_stream_reexpands_glob_each_execution(empty_db, tmp_path): + _, conn = empty_db + (tmp_path / "part1.csv").write_text("id\n1\n") + query = ( + f'LOAD FROM "{tmp_path}/part*.csv" ' + "(header=true, batch_size=1) RETURN id ORDER BY id" + ) + assert list(conn.execute(query)) == [[1]] + (tmp_path / "part2.csv").write_text("id\n2\n") + assert list(conn.execute(query)) == [[1], [2]] + assert list(conn.execute(query)) == [[1], [2]] + + +def test_stream_preserves_literal_heads_and_global_aggregation(empty_db, tmp_path): + _, conn = empty_db + assert list(conn.execute("RETURN 42")) == [[42]] + assert list(conn.execute("UNWIND [1, 2, 3] AS x RETURN x + 1")) == [[2], [3], [4]] + path = tmp_path / "values.csv" + path.write_text("id\n3\n1\n2\n") + source = f'LOAD FROM "{path}" (header=true, batch_size=1)' + assert list(conn.execute(f"{source} RETURN count(*), sum(id)")) == [[3, 6]] + assert list(conn.execute(f"{source} RETURN id ORDER BY id LIMIT 2")) == [[1], [2]] + + +def test_failed_stream_copy_does_not_persist_partial_batch(tmp_path): + from neug import Database + + db_path = str(tmp_path / "durable") + csv_path = tmp_path / "late_error.csv" + csv_path.write_text("id|name\n1|valid\ninvalid|bad\n") + db = Database(db_path=db_path, mode="w") + conn = db.connect() + try: + conn.execute("CREATE NODE TABLE person(id INT64, name STRING, PRIMARY KEY(id))") + conn.execute("CREATE (:person {id: 0, name: 'original'})") + with pytest.raises(RuntimeError): + conn.execute(f'COPY person FROM "{csv_path}" (batch_size=1)') + finally: + conn.close() + db.close() + + db = Database(db_path=db_path, mode="w") + conn = db.connect() + try: + assert list(conn.execute("MATCH (p:person) RETURN p.id, p.name")) == [ + [0, "original"] + ] + finally: + conn.close() + db.close() + + +def test_unwind_limit_stops_upstream_batches(empty_db, tmp_path): + _, conn = empty_db + path = tmp_path / "incremental.csv" + path.write_text("id\n1\n2\n3\n") + result = conn.execute( + f'PROFILE LOAD FROM "{path}" (header=true, batch_size=1) ' + "UNWIND [id, id + 10] AS x RETURN x LIMIT 2" + ) + assert list(result) == [[1], [11]] + source = next( + op + for op in result.get_profile_metrics()["operators"] + if op["operator_name"] == "DataSourceOpr" + ) + assert source["output_rows"] == 1 + + +def test_cross_batch_skip_and_topk(empty_db, tmp_path): + _, conn = empty_db + path = tmp_path / "cross_batch.csv" + path.write_text("id\n4\n1\n4\n3\n2\n") + source = f'LOAD FROM "{path}" (header=true, batch_size=1)' + assert list(conn.execute(f"{source} RETURN id SKIP 2 LIMIT 2")) == [[4], [3]] + assert list(conn.execute(f"{source} RETURN id + 1 AS x ORDER BY x LIMIT 2")) == [ + [2], + [3], + ] + assert list(conn.execute(f"{source} RETURN id LIMIT 0")) == []